Merge pull request cockroachdb#156501 from yuzefovich/backport25.3-155824-156110

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

Author: yuzefovich

pkg/sql/apply_join.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/distsql"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
@@ -35,6 +36,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 +253,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 +273,7 @@ func runPlanInsidePlan(
 	resultWriter rowResultWriter,
 	deferredRoutineSender eval.DeferredRoutineSender,
 	stmtForDistSQLDiagram string,
-) error {
+) (topLevelQueryStats, error) {
 	defer plan.close(ctx)
 	execCfg := params.ExecCfg()
 	recv := MakeDistSQLReceiver(
@@ -286,6 +292,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 +312,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,
@@ -309,9 +322,9 @@ func runPlanInsidePlan(
 			recv,
 			&subqueryResultMemAcc,
 			false, /* skipDistSQLDiagramGeneration */
-			params.p.mustUseLeafTxn(),
+			params.p.innerPlansMustUseLeafTxn(),
 		) {
-			return resultWriter.Err()
+			return recv.stats, resultWriter.Err()
 		}
 	}
 
@@ -324,7 +337,9 @@ func runPlanInsidePlan(
 	planCtx := execCfg.DistSQLPlanner.NewPlanningCtx(ctx, evalCtx, &plannerCopy, plannerCopy.txn, distributeType)
 	planCtx.distSQLProhibitedErr = distSQLProhibitedErr
 	planCtx.stmtType = recv.stmtType
-	planCtx.mustUseLeafTxn = params.p.mustUseLeafTxn()
+	if params.p.innerPlansMustUseLeafTxn() {
+		planCtx.flowConcurrency = distsql.ConcurrencyWithOuterPlan
+	}
 	planCtx.stmtForDistSQLDiagram = stmtForDistSQLDiagram
 
 	// Wrap PlanAndRun in a function call so that we clean up immediately.
@@ -338,10 +353,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 +373,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

@@ -3961,6 +3961,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

@@ -3507,6 +3507,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/server.go

@@ -533,6 +533,28 @@ func newFlow(
 	return rowflow.NewRowBasedFlow(base)
 }
 
+// ConcurrencyKind indicates which concurrency type is present within the local
+// DistSQL flow. Note that inter-node concurrency (i.e. whether we have a
+// distributed plan) is not reflected here.
+type ConcurrencyKind uint32
+
+const (
+	// ConcurrencyHasParallelProcessors, if set, indicates that we have multiple
+	// processors running for the same plan stage.
+	ConcurrencyHasParallelProcessors ConcurrencyKind = (1 << iota)
+	// ConcurrencyStreamer, if set, indicates we have concurrency due to usage
+	// of the Streamer API.
+	ConcurrencyStreamer
+	// ConcurrencyParallelChecks, if set, indicates that we're running
+	// post-query CHECKs in parallel with each other (i.e. the concurrency is
+	// with _other_ local flows).
+	ConcurrencyParallelChecks
+	// ConcurrencyWithOuterPlan, if set, indicates that - if we're running an
+	// "inner" plan (like an apply-join iteration or a routine) - we might have
+	// concurrency with the "outer" plan.
+	ConcurrencyWithOuterPlan
+)
+
 // LocalState carries information that is required to set up a flow with wrapped
 // planNodes.
 type LocalState struct {
@@ -547,13 +569,9 @@ type LocalState struct {
 	// remote flows.
 	IsLocal bool
 
-	// HasConcurrency indicates whether the local flow uses multiple goroutines.
-	HasConcurrency bool
-
-	// MustUseLeaf indicates whether the local flow must use the LeafTxn even if
-	// there is no concurrency in the flow on its own because there would be
-	// concurrency with other flows which prohibits the usage of the RootTxn.
-	MustUseLeaf bool
+	// concurrency tracks the types of concurrency present when accessing the
+	// Txn.
+	concurrency ConcurrencyKind
 
 	// Txn is filled in on the gateway only. It is the RootTxn that the query is running in.
 	// This will be used directly by the flow if the flow has no concurrency and IsLocal is set.
@@ -570,10 +588,22 @@ type LocalState struct {
 	LocalVectorSources map[int32]any
 }
 
+// AddConcurrency marks the given concurrency kinds as present in the local
+// flow.
+func (l *LocalState) AddConcurrency(kind ConcurrencyKind) {
+	l.concurrency |= kind
+}
+
+// GetConcurrency returns the bit-mask representing all concurrency kinds
+// present in the local flow.
+func (l LocalState) GetConcurrency() ConcurrencyKind {
+	return l.concurrency
+}
+
 // MustUseLeafTxn returns true if a LeafTxn must be used. It is valid to call
-// this method only after IsLocal and HasConcurrency have been set correctly.
+// this method only after IsLocal and all concurrency kinds have been set.
 func (l LocalState) MustUseLeafTxn() bool {
-	return !l.IsLocal || l.HasConcurrency || l.MustUseLeaf
+	return !l.IsLocal || l.concurrency != 0
 }
 
 // SetupLocalSyncFlow sets up a synchronous flow on the current (planning) node,

pkg/sql/distsql_physical_planner.go

@@ -989,7 +989,8 @@ type PlanningCtx struct {
 	// isLocal is set to true if we're planning this query on a single node.
 	isLocal bool
 	// distSQLProhibitedErr, if set, indicates why the plan couldn't be
-	// distributed.
+	// distributed. If any part of the plan isn't distributable, then this is
+	// guaranteed to be non-nil.
 	distSQLProhibitedErr error
 	planner              *planner
 
@@ -1031,11 +1032,11 @@ type PlanningCtx struct {
 	// query).
 	subOrPostQuery bool
 
-	// mustUseLeafTxn, if set, indicates that this PlanningCtx is used to handle
+	// flowConcurrency will be non-zero when this PlanningCtx is used to handle
 	// one of the plans that will run in parallel with other plans. As such, the
 	// DistSQL planner will need to use the LeafTxn (even if it's not needed
 	// based on other "regular" factors).
-	mustUseLeafTxn bool
+	flowConcurrency distsql.ConcurrencyKind
 
 	// onFlowCleanup contains non-nil functions that will be called after the
 	// local flow finished running and is being cleaned up. It allows us to
@@ -4371,7 +4372,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 +4381,9 @@ func (dsp *DistSQLPlanner) wrapPlan(
 		useFastPath,
 		firstNotWrapped,
 	)
+	if err != nil {
+		return nil, err
+	}
 
 	localProcIdx := p.AddLocalProcessor(wrapper)
 	var input []execinfrapb.InputSyncSpec