sql: harden recent fix to top-level query stats with routines

This commit relaxes the requirements for when it's safe to set the
routine metadata forwarder. Previously, we only set it when we end up
using the RootTxn in a local plan, but now we examine all possible kinds
of concurrency within the local plan to see whether it's actually safe
to set the metadata forwarder. This seems like a nice cleanup as well.

Note that this commit independently fixes the issue mentioned in the
previous commit.

Release note: None

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"
@@ -321,7 +322,7 @@ func runPlanInsidePlan(
 			recv,
 			&subqueryResultMemAcc,
 			false, /* skipDistSQLDiagramGeneration */
-			params.p.mustUseLeafTxn(),
+			params.p.innerPlansMustUseLeafTxn(),
 		) {
 			return recv.stats, resultWriter.Err()
 		}
@@ -336,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.

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

@@ -1032,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

pkg/sql/distsql_running.go

@@ -745,7 +745,7 @@ func (dsp *DistSQLPlanner) Run(
 	// the line.
 	localState.EvalContext = evalCtx
 	localState.IsLocal = planCtx.isLocal
-	localState.MustUseLeaf = planCtx.mustUseLeafTxn
+	localState.AddConcurrency(planCtx.flowConcurrency)
 	localState.Txn = txn
 	localState.LocalProcs = plan.LocalProcessors
 	localState.LocalVectorSources = plan.LocalVectorSources
@@ -768,15 +768,19 @@ func (dsp *DistSQLPlanner) Run(
 		// cannot create a LeafTxn, so we cannot parallelize scans.
 		planCtx.parallelizeScansIfLocal = false
 		for _, flow := range flows {
-			localState.HasConcurrency = localState.HasConcurrency || execinfra.HasParallelProcessors(flow)
+			if execinfra.HasParallelProcessors(flow) {
+				localState.AddConcurrency(distsql.ConcurrencyHasParallelProcessors)
+			}
 		}
 	} else {
 		if planCtx.isLocal && noMutations && planCtx.parallelizeScansIfLocal {
 			// Even though we have a single flow on the gateway node, we might
 			// have decided to parallelize the scans. If that's the case, we
 			// will need to use the Leaf txn.
 			for _, flow := range flows {
-				localState.HasConcurrency = localState.HasConcurrency || execinfra.HasParallelProcessors(flow)
+				if execinfra.HasParallelProcessors(flow) {
+					localState.AddConcurrency(distsql.ConcurrencyHasParallelProcessors)
+				}
 			}
 		}
 		if noMutations {
@@ -877,7 +881,7 @@ func (dsp *DistSQLPlanner) Run(
 							// Both index and lookup joins, with and without
 							// ordering, are executed via the Streamer API that has
 							// concurrency.
-							localState.HasConcurrency = true
+							localState.AddConcurrency(distsql.ConcurrencyStreamer)
 							break
 						}
 					}
@@ -1000,11 +1004,37 @@ 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 {
+	if len(flows) == 1 && planCtx.planner != nil {
+		// We have a fully local plan, so check whether it'll be safe to use the
+		// DistSQLReceiver to push the metadata into directly from routines
+		// (which is the case when we don't have any concurrency between
+		// routines themselves as well as a routine and the "head" processor -
+		// the one pushing into the DistSQLReceiver).
+		var safe bool
+		if evalCtx.Txn != nil && evalCtx.Txn.Type() == kv.RootTxn {
+			// We have a RootTxn, so we don't expect any concurrency whatsoever.
+			safe = true
+		} else {
+			// We have a LeafTxn, so we need to examine what kind of concurrency
+			// is present in the flow.
+			var safeConcurrency distsql.ConcurrencyKind
+			// We don't care whether we use the Streamer API - it has
+			// concurrency only at the KV client level and below.
+			safeConcurrency |= distsql.ConcurrencyStreamer
+			// If we have "outer plan" concurrency, the "inner" and the
+			// "outer" plans have their own DistSQLReceivers.
+			//
+			// Note that the same is the case with parallel CHECKs concurrency,
+			// but then planCtx.planner is shared between goroutines, so we'll
+			// avoid mutating it. (We can't have routines in post-query CHECKs
+			// since only FK and UNIQUE checks are run in parallel.)
+			safeConcurrency |= distsql.ConcurrencyWithOuterPlan
+			unsafeConcurrency := ^safeConcurrency
+			if localState.GetConcurrency()&unsafeConcurrency == 0 {
+				safe = true
+			}
+		}
+		if safe {
 			planCtx.planner.routineMetadataForwarder = recv
 		}
 	}
@@ -1859,7 +1889,7 @@ func (dsp *DistSQLPlanner) PlanAndRunAll(
 			// Skip the diagram generation since on this "main" query path we
 			// can get it via the statement bundle.
 			true,  /* skipDistSQLDiagramGeneration */
-			false, /* mustUseLeafTxn */
+			false, /* innerPlansMustUseLeafTxn */
 		) {
 			return recv.commErr
 		}
@@ -1926,7 +1956,7 @@ func (dsp *DistSQLPlanner) PlanAndRunSubqueries(
 	recv *DistSQLReceiver,
 	subqueryResultMemAcc *mon.BoundAccount,
 	skipDistSQLDiagramGeneration bool,
-	mustUseLeafTxn bool,
+	innerPlansMustUseLeafTxn bool,
 ) bool {
 	for planIdx, subqueryPlan := range subqueryPlans {
 		if err := dsp.planAndRunSubquery(
@@ -1939,7 +1969,7 @@ func (dsp *DistSQLPlanner) PlanAndRunSubqueries(
 			recv,
 			subqueryResultMemAcc,
 			skipDistSQLDiagramGeneration,
-			mustUseLeafTxn,
+			innerPlansMustUseLeafTxn,
 		); err != nil {
 			recv.SetError(err)
 			return false
@@ -1963,7 +1993,7 @@ func (dsp *DistSQLPlanner) planAndRunSubquery(
 	recv *DistSQLReceiver,
 	subqueryResultMemAcc *mon.BoundAccount,
 	skipDistSQLDiagramGeneration bool,
-	mustUseLeafTxn bool,
+	innerPlansMustUseLeafTxn bool,
 ) error {
 	subqueryDistribution, distSQLProhibitedErr := planner.getPlanDistribution(ctx, subqueryPlan.plan)
 	distribute := DistributionType(LocalDistribution)
@@ -1975,7 +2005,9 @@ func (dsp *DistSQLPlanner) planAndRunSubquery(
 	subqueryPlanCtx.stmtType = tree.Rows
 	subqueryPlanCtx.skipDistSQLDiagramGeneration = skipDistSQLDiagramGeneration
 	subqueryPlanCtx.subOrPostQuery = true
-	subqueryPlanCtx.mustUseLeafTxn = mustUseLeafTxn
+	if innerPlansMustUseLeafTxn {
+		subqueryPlanCtx.flowConcurrency = distsql.ConcurrencyWithOuterPlan
+	}
 	if planner.instrumentation.ShouldSaveFlows() {
 		subqueryPlanCtx.saveFlows = getDefaultSaveFlowsFunc(ctx, planner, planComponentTypeSubquery)
 	}
@@ -2590,7 +2622,9 @@ func (dsp *DistSQLPlanner) planAndRunPostquery(
 	}
 	postqueryPlanCtx.associateNodeWithComponents = associateNodeWithComponents
 	postqueryPlanCtx.collectExecStats = planner.instrumentation.ShouldCollectExecStats()
-	postqueryPlanCtx.mustUseLeafTxn = parallelCheck
+	if parallelCheck {
+		postqueryPlanCtx.flowConcurrency = distsql.ConcurrencyParallelChecks
+	}
 
 	postqueryPhysPlan, physPlanCleanup, err := dsp.createPhysPlan(ctx, postqueryPlanCtx, postqueryPlan)
 	defer physPlanCleanup()

pkg/sql/planner.go

@@ -1071,8 +1071,9 @@ func (p *planner) ClearTableStatsCache() {
 	}
 }
 
-// mustUseLeafTxn returns true if inner plans must use a leaf transaction.
-func (p *planner) mustUseLeafTxn() bool {
+// innerPlansMustUseLeafTxn returns true if inner plans must use a leaf
+// transaction.
+func (p *planner) innerPlansMustUseLeafTxn() bool {
 	return atomic.LoadInt32(&p.atomic.innerPlansMustUseLeafTxn) >= 1
 }
 

pkg/sql/schema_changer.go

@@ -470,7 +470,7 @@ func (sc *SchemaChanger) backfillQueryIntoTable(
 					ctx, localPlanner, localPlanner.ExtendedEvalContextCopy,
 					localPlanner.curPlan.subqueryPlans, recv, &subqueryResultMemAcc,
 					false, /* skipDistSQLDiagramGeneration */
-					false, /* mustUseLeafTxn */
+					false, /* innerPlansMustUseLeafTxn */
 				) {
 					if planAndRunErr = rw.Err(); planAndRunErr != nil {
 						return