sql: fix top-level query stats when "inner" plans are present

Previously, whenever we ran an "inner" plan (via `runPlanInsidePlan`
helper), we ignored the top-level query stats for the "inner" plan. As
a result, reads and writes performed by the inner plan weren't reflected
in the "outer" top-level query stats. This is now fixed by adjusting the
routines, apply joins, and recursive CTEs to propagate their metrics as
ProducerMetadata objects.

Note that for routines the access to the DistSQL infra is rather
difficult, so we plumbed the access via the planner straight into the
DistSQLReceiver, and even though it's ugly, it should work in practice.
The only alternative I see is adding the necessary reference into the
`eval.Context`, but then it gets tricky to actually set that and ensure
the right copy of the eval context is observed by all routines (plus
we'd need to make the copy or restore the original state somehow), so
I chose to not pursue it.

Release note (bug fix): Previously, CockroachDB didn't include reads and
writes performed by routines (user-defined functions and stored
procedures) as well as apply joins into `bytes read`, `rows read`, and
`rows written` statement execution statistics, and this is now fixed.
The bug has been present since before 23.2 version.

Author: yuzefovich

pkg/sql/apply_join.go

@@ -35,6 +35,10 @@ type applyJoinNode struct {
 	// The data source with no outer columns.
 	singleInputPlanNode
 
+	// forwarder allows propagating the ProducerMetadata towards the
+	// DistSQLReceiver.
+	forwarder metadataForwarder
+
 	// pred represents the join predicate.
 	pred *joinPredicate
 
@@ -248,12 +252,13 @@ func (a *applyJoinNode) runNextRightSideIteration(params runParams, leftRow tree
 	}
 	plan := p.(*planComponents)
 	rowResultWriter := NewRowResultWriter(&a.run.rightRows)
-	if err := runPlanInsidePlan(
-		ctx, params, plan, rowResultWriter,
-		nil /* deferredRoutineSender */, "", /* stmtForDistSQLDiagram */
-	); err != nil {
+	queryStats, err := runPlanInsidePlan(
+		ctx, params, plan, rowResultWriter, nil /* deferredRoutineSender */, "", /* stmtForDistSQLDiagram */
+	)
+	if err != nil {
 		return err
 	}
+	forwardInnerQueryStats(a.forwarder, queryStats)
 	a.run.rightRowsIterator = newRowContainerIterator(ctx, a.run.rightRows)
 	return nil
 }
@@ -267,7 +272,7 @@ func runPlanInsidePlan(
 	resultWriter rowResultWriter,
 	deferredRoutineSender eval.DeferredRoutineSender,
 	stmtForDistSQLDiagram string,
-) error {
+) (topLevelQueryStats, error) {
 	defer plan.close(ctx)
 	execCfg := params.ExecCfg()
 	recv := MakeDistSQLReceiver(
@@ -286,6 +291,11 @@ func runPlanInsidePlan(
 	// before we can produce any "outer" rows to be returned to the client, so
 	// we make sure to unset pausablePortal field on the planner.
 	plannerCopy.pausablePortal = nil
+	// Avoid any possible metadata confusion by unsetting the
+	// routineMetadataForwarder (if there is a routine in the inner plan that
+	// needs it, then the plannerCopy will be updated during the inner plan
+	// setup).
+	plannerCopy.routineMetadataForwarder = nil
 
 	// planner object embeds the extended eval context, so we will modify that
 	// (which won't affect the outer planner's extended eval context), and we'll
@@ -301,6 +311,8 @@ func runPlanInsidePlan(
 		// return from this method (after the main query is executed).
 		subqueryResultMemAcc := params.p.Mon().MakeBoundAccount()
 		defer subqueryResultMemAcc.Close(ctx)
+		// Note that planAndRunSubquery updates recv.stats with top-level
+		// subquery stats.
 		if !execCfg.DistSQLPlanner.PlanAndRunSubqueries(
 			ctx,
 			&plannerCopy,
@@ -311,7 +323,7 @@ func runPlanInsidePlan(
 			false, /* skipDistSQLDiagramGeneration */
 			params.p.mustUseLeafTxn(),
 		) {
-			return resultWriter.Err()
+			return recv.stats, resultWriter.Err()
 		}
 	}
 
@@ -338,10 +350,10 @@ func runPlanInsidePlan(
 
 	// Check if there was an error interacting with the resultWriter.
 	if recv.commErr != nil {
-		return recv.commErr
+		return recv.stats, recv.commErr
 	}
 	if resultWriter.Err() != nil {
-		return resultWriter.Err()
+		return recv.stats, resultWriter.Err()
 	}
 
 	plannerCopy.autoCommit = false
@@ -358,10 +370,10 @@ func runPlanInsidePlan(
 	// need to update the plan for cleanup purposes before proceeding.
 	*plan = plannerCopy.curPlan.planComponents
 	if recv.commErr != nil {
-		return recv.commErr
+		return recv.stats, recv.commErr
 	}
 
-	return resultWriter.Err()
+	return recv.stats, resultWriter.Err()
 }
 
 func (a *applyJoinNode) Values() tree.Datums {

pkg/sql/conn_executor.go

@@ -3957,6 +3957,7 @@ func (ex *connExecutor) initPlanner(ctx context.Context, p *planner) {
 	p.noticeSender = nil
 	p.preparedStatements = ex.getPrepStmtsAccessor()
 	p.sqlCursors = ex.getCursorAccessor()
+	p.routineMetadataForwarder = nil
 	p.storedProcTxnState = ex.getStoredProcTxnStateAccessor()
 	p.createdSequences = ex.getCreatedSequencesAccessor()
 

pkg/sql/conn_executor_exec.go

@@ -3496,6 +3496,8 @@ type topLevelQueryStats struct {
 	// client receiving the PGWire protocol messages (as well as construcing
 	// those messages).
 	clientTime time.Duration
+	// NB: when adding another field here, consider whether
+	// forwardInnerQueryStats method needs an adjustment.
 }
 
 func (s *topLevelQueryStats) add(other *topLevelQueryStats) {

pkg/sql/distsql_physical_planner.go

@@ -4371,7 +4371,7 @@ func (dsp *DistSQLPlanner) wrapPlan(
 	// expecting is in fact RowsAffected. RowsAffected statements return a single
 	// row with the number of rows affected by the statement, and are the only
 	// types of statement where it's valid to invoke a plan's fast path.
-	wrapper := newPlanNodeToRowSource(
+	wrapper, err := newPlanNodeToRowSource(
 		n,
 		runParams{
 			extendedEvalCtx: &evalCtx,
@@ -4380,6 +4380,9 @@ func (dsp *DistSQLPlanner) wrapPlan(
 		useFastPath,
 		firstNotWrapped,
 	)
+	if err != nil {
+		return nil, err
+	}
 
 	localProcIdx := p.AddLocalProcessor(wrapper)
 	var input []execinfrapb.InputSyncSpec

pkg/sql/distsql_running.go

@@ -1005,6 +1005,15 @@ func (dsp *DistSQLPlanner) Run(
 		return
 	}
 
+	if len(flows) == 1 && evalCtx.Txn != nil && evalCtx.Txn.Type() == kv.RootTxn {
+		// If we have a fully local plan and a RootTxn, we don't expect any
+		// concurrency, so it's safe to use the DistSQLReceiver to push the
+		// metadata into directly from routines.
+		if planCtx.planner != nil {
+			planCtx.planner.routineMetadataForwarder = recv
+		}
+	}
+
 	if finishedSetupFn != nil {
 		finishedSetupFn(flow)
 	}
@@ -1131,6 +1140,8 @@ type DistSQLReceiver struct {
 	}
 }
 
+var _ metadataForwarder = &DistSQLReceiver{}
+
 // rowResultWriter is a subset of CommandResult to be used with the
 // DistSQLReceiver. It's implemented by RowResultWriter.
 type rowResultWriter interface {
@@ -1503,6 +1514,35 @@ func (r *DistSQLReceiver) checkConcurrentError() {
 	}
 }
 
+type metadataForwarder interface {
+	forwardMetadata(metadata *execinfrapb.ProducerMetadata)
+}
+
+// forwardInnerQueryStats propagates the query stats of "inner" plans as
+// metadata via the forwarder.
+func forwardInnerQueryStats(f metadataForwarder, stats topLevelQueryStats) {
+	if !buildutil.CrdbTestBuild && f == nil {
+		// Safety measure in production builds in case the forwarder is nil for
+		// some reason.
+		return
+	}
+	meta := execinfrapb.GetProducerMeta()
+	meta.Metrics = execinfrapb.GetMetricsMeta()
+	meta.Metrics.BytesRead = stats.bytesRead
+	meta.Metrics.RowsRead = stats.rowsRead
+	meta.Metrics.RowsWritten = stats.rowsWritten
+	// stats.networkEgressEstimate and stats.clientTime are ignored since they
+	// only matter at the "true" top-level query (and actually should be zero
+	// here anyway).
+	f.forwardMetadata(meta)
+}
+
+func (r *DistSQLReceiver) forwardMetadata(metadata *execinfrapb.ProducerMetadata) {
+	// Note that we don't use pushMeta method directly in order to go through
+	// the testing callback path.
+	r.Push(nil /* row */, metadata)
+}
+
 // pushMeta takes in non-empty metadata object and pushes it to the result
 // writer. Possibly updated status is returned.
 func (r *DistSQLReceiver) pushMeta(meta *execinfrapb.ProducerMetadata) execinfra.ConsumerStatus {

pkg/sql/distsql_running_test.go

@@ -1214,3 +1214,127 @@ SELECT id, details FROM jobs AS j INNER JOIN cte1 ON id = job_id WHERE id = 1;
 	require.Equal(t, 1, id)
 	require.Equal(t, 1, details)
 }
+
+// TestTopLevelQueryStats verifies that top-level query stats are collected
+// correctly, including when the query executes "plans inside plans".
+func TestTopLevelQueryStats(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	// testQuery will be updated throughout the test to the current target.
+	var testQuery atomic.Value
+	// The callback will send number of rows read and rows written (for each
+	// ProducerMetadata.Metrics object) on these channels, respectively.
+	rowsReadCh, rowsWrittenCh := make(chan int64), make(chan int64)
+	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+		Knobs: base.TestingKnobs{
+			SQLExecutor: &ExecutorTestingKnobs{
+				DistSQLReceiverPushCallbackFactory: func(_ context.Context, query string) func(rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) (rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) {
+					if target := testQuery.Load(); target == nil || target.(string) != query {
+						return nil
+					}
+					return func(row rowenc.EncDatumRow, batch coldata.Batch, meta *execinfrapb.ProducerMetadata) (rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) {
+						if meta != nil && meta.Metrics != nil {
+							rowsReadCh <- meta.Metrics.RowsRead
+							rowsWrittenCh <- meta.Metrics.RowsWritten
+						}
+						return row, batch, meta
+					}
+				},
+			},
+		},
+	})
+	defer s.Stopper().Stop(context.Background())
+
+	if _, err := sqlDB.Exec(`
+CREATE TABLE t (k INT PRIMARY KEY);
+INSERT INTO t SELECT generate_series(1, 10);
+CREATE FUNCTION no_reads() RETURNS INT AS 'SELECT 1' LANGUAGE SQL;
+CREATE FUNCTION reads() RETURNS INT AS 'SELECT count(*) FROM t' LANGUAGE SQL;
+CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x' LANGUAGE SQL;
+`); err != nil {
+		t.Fatal(err)
+	}
+
+	for _, tc := range []struct {
+		name           string
+		query          string
+		expRowsRead    int64
+		expRowsWritten int64
+	}{
+		{
+			name:           "simple read",
+			query:          "SELECT k FROM t",
+			expRowsRead:    10,
+			expRowsWritten: 0,
+		},
+		{
+			name:           "simple write",
+			query:          "INSERT INTO t SELECT generate_series(11, 42)",
+			expRowsRead:    0,
+			expRowsWritten: 32,
+		},
+		{
+			name: "read with apply join",
+			query: `SELECT (
+    WITH foo AS MATERIALIZED (SELECT k FROM t AS x WHERE x.k = y.k)
+    SELECT * FROM foo
+  ) FROM t AS y`,
+			expRowsRead:    84, // scanning the table twice
+			expRowsWritten: 0,
+		},
+		{
+			name:           "routine, no reads",
+			query:          "SELECT no_reads()",
+			expRowsRead:    0,
+			expRowsWritten: 0,
+		},
+		{
+			name:           "routine, reads",
+			query:          "SELECT reads()",
+			expRowsRead:    42,
+			expRowsWritten: 0,
+		},
+		{
+			name:           "routine, write",
+			query:          "SELECT write(43)",
+			expRowsRead:    0,
+			expRowsWritten: 1,
+		},
+		{
+			name:           "routine, multiple reads and writes",
+			query:          "SELECT reads(), write(44), reads(), write(45), write(46), reads()",
+			expRowsRead:    133, // first read is 43 rows, second is 44, third is 46
+			expRowsWritten: 3,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			testQuery.Store(tc.query)
+			errCh := make(chan error)
+			// Spin up the worker goroutine which will actually execute the
+			// query.
+			go func() {
+				defer close(errCh)
+				_, err := sqlDB.Exec(tc.query)
+				errCh <- err
+			}()
+			// In the main goroutine, loop until the query is completed while
+			// accumulating the top-level query stats.
+			var rowsRead, rowsWritten int64
+		LOOP:
+			for {
+				select {
+				case read := <-rowsReadCh:
+					rowsRead += read
+				case written := <-rowsWrittenCh:
+					rowsWritten += written
+				case err := <-errCh:
+					require.NoError(t, err)
+					break LOOP
+				}
+			}
+			require.Equal(t, tc.expRowsRead, rowsRead)
+			require.Equal(t, tc.expRowsWritten, rowsWritten)
+		})
+	}
+}

pkg/sql/plan_node_to_row_source.go

@@ -23,10 +23,6 @@ import (
 	"github.com/cockroachdb/errors"
 )
 
-type metadataForwarder interface {
-	forwardMetadata(metadata *execinfrapb.ProducerMetadata)
-}
-
 type planNodeToRowSource struct {
 	execinfra.ProcessorBase
 
@@ -50,6 +46,7 @@ type planNodeToRowSource struct {
 var _ execinfra.LocalProcessor = &planNodeToRowSource{}
 var _ execreleasable.Releasable = &planNodeToRowSource{}
 var _ execopnode.OpNode = &planNodeToRowSource{}
+var _ metadataForwarder = &planNodeToRowSource{}
 
 var planNodeToRowSourcePool = sync.Pool{
 	New: func() interface{} {
@@ -59,7 +56,7 @@ var planNodeToRowSourcePool = sync.Pool{
 
 func newPlanNodeToRowSource(
 	source planNode, params runParams, fastPath bool, firstNotWrapped planNode,
-) *planNodeToRowSource {
+) (*planNodeToRowSource, error) {
 	p := planNodeToRowSourcePool.Get().(*planNodeToRowSource)
 	*p = planNodeToRowSource{
 		ProcessorBase:   p.ProcessorBase,
@@ -85,7 +82,39 @@ func newPlanNodeToRowSource(
 	} else {
 		p.row = make(rowenc.EncDatumRow, len(p.outputTypes))
 	}
-	return p
+	// Find any planNodes that need a way to propagate metadata before we get to
+	// the firstNotWrapped - this planNodeToRowSource adapter will be the
+	// forwarder for all of them.
+	var setForwarder func(parent planNode) error
+	setForwarder = func(parent planNode) error {
+		switch t := parent.(type) {
+		case *applyJoinNode:
+			t.forwarder = p
+		case *recursiveCTENode:
+			t.forwarder = p
+		}
+		for i, n := 0, parent.InputCount(); i < n; i++ {
+			child, err := parent.Input(i)
+			if err != nil {
+				return err
+			}
+			if child == p.firstNotWrapped {
+				// Once we get to the firstNotWrapped, we stop the recursion
+				// since all remaining planNodes aren't our responsibility.
+				return nil
+			}
+			if err = setForwarder(child); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+	err := setForwarder(p.node)
+	if err != nil {
+		p.Release()
+		return nil, err
+	}
+	return p, nil
 }
 
 // MustBeStreaming implements the execinfra.Processor interface.