Merge #156110

156110: sql: harden DistSQL interactions with routines r=yuzefovich a=yuzefovich

**sql: unconditionally check DistSQL supportability**

This commit hardens c17591e. In
particular, in that change we disabled the usage of the Streamer
whenever some part of the plan isn't distributable. However, population
of `distSQLProhibitedErr` was left on a best-effort basis. In
particular, if we have `distsql=off` (as well as a couple of other
conditions), we'd skip the supportability check. This meant that we
could still end up using the Streamer in some illegal cases - one
example that I came up with is if we have a routine (which currently
prohibits distsql), it might access the RootTxn concurrently with the
LeafTxn access by the Streamer, which is no bueno. This commit makes it
so that we do DistSQL supportability check unconditionally. Note that
this shouldn't really have a performance impact - after all, we do
expect this check to be done pretty much all the time (unless someone
runs with `distsql=off`).

Additionally, this change happens to fix a nil pointer error around the
recent fix to top-level query stats in presence of routines. Namely,
that patch only set the metadata forwarder for local plans if we end up
using the RootTxn, but in a query where we happened to use the streamer
with routines, (before this patch) we'd end up with LeafTxn and unset
metadata forwarder. Now we'll have streamer disabled, so we'll have the
RootTxn and the forwarder will be set.

Fixes: #155955.

**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

Co-authored-by: Yahor Yuzefovich <yahor@cockroachlabs.com>

Author: craig[bot]

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

@@ -545,6 +545,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 {
@@ -559,13 +581,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.
@@ -582,10 +600,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

@@ -977,7 +977,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
 
@@ -1019,11 +1020,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

@@ -754,7 +754,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
@@ -777,15 +777,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 {
@@ -875,23 +879,18 @@ func (dsp *DistSQLPlanner) Run(
 			// that we might have a plan where some expression (e.g. a cast to
 			// an Oid type) uses the planner's txn (which is the RootTxn), so
 			// it'd be illegal to use LeafTxns for a part of such plan.
-			// TODO(yuzefovich): this check is both excessive and insufficient.
-			// For example:
-			// - it disables the usage of the Streamer when a subquery has an
-			// Oid type, but that would have no impact on usage of the Streamer
-			// in the main query;
-			// - it might allow the usage of the Streamer even when the internal
-			// executor is used by a part of the plan, and the IE would use the
-			// RootTxn. Arguably, this would be a bug in not prohibiting the
-			// DistSQL altogether.
+			// TODO(yuzefovich): this check could be excessive. For example, it
+			// disables the usage of the Streamer when a subquery has an Oid
+			// type (due to a serialization issue), but that would have no
+			// impact on usage of the Streamer in the main query.
 			if !containsLocking && !mustUseRootTxn && planCtx.distSQLProhibitedErr == nil {
 				if evalCtx.SessionData().StreamerEnabled {
 					for _, proc := range plan.Processors {
 						if jr := proc.Spec.Core.JoinReader; jr != nil {
 							// 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
 						}
 					}
@@ -1014,11 +1013,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
 		}
 	}
@@ -1994,7 +2019,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
 		}
@@ -2061,7 +2086,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(
@@ -2074,7 +2099,7 @@ func (dsp *DistSQLPlanner) PlanAndRunSubqueries(
 			recv,
 			subqueryResultMemAcc,
 			skipDistSQLDiagramGeneration,
-			mustUseLeafTxn,
+			innerPlansMustUseLeafTxn,
 		); err != nil {
 			recv.SetError(err)
 			return false
@@ -2098,7 +2123,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)
@@ -2110,7 +2135,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)
 	}
@@ -2729,7 +2756,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/distsql_running_test.go

@@ -1221,12 +1221,13 @@ func TestTopLevelQueryStats(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
+	ctx := context.Background()
 	// 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{
+	srv, 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) {
@@ -1244,11 +1245,13 @@ func TestTopLevelQueryStats(t *testing.T) {
 			},
 		},
 	})
-	defer s.Stopper().Stop(context.Background())
+	defer srv.Stopper().Stop(ctx)
+	conn, err := sqlDB.Conn(ctx)
+	require.NoError(t, err)
 
 	if _, err := sqlDB.Exec(`
-CREATE TABLE t (k INT PRIMARY KEY);
-INSERT INTO t SELECT generate_series(1, 10);
+CREATE TABLE t (k INT PRIMARY KEY, i INT, v INT, INDEX(i));
+INSERT INTO t SELECT i, 1, 1 FROM generate_series(1, 10) AS g(i);
 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;
@@ -1259,6 +1262,7 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
 	for _, tc := range []struct {
 		name           string
 		query          string
+		setup, cleanup string // optional
 		expRowsRead    int64
 		expRowsWritten int64
 	}{
@@ -1268,6 +1272,16 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
 			expRowsRead:    10,
 			expRowsWritten: 0,
 		},
+		{
+			name:    "routine and index join (used to be powered by streamer)",
+			query:   "SELECT v FROM t@t_i_idx WHERE reads() > 0",
+			setup:   "SET distsql=off",
+			cleanup: "RESET distsql",
+			// 10 rows for secondary index, 10 for index join into primary, and
+			// then for each row do ten-row-scan in the routine.
+			expRowsRead:    120,
+			expRowsWritten: 0,
+		},
 		{
 			name:           "simple write",
 			query:          "INSERT INTO t SELECT generate_series(11, 42)",
@@ -1309,13 +1323,23 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
+			if tc.setup != "" {
+				_, err := conn.ExecContext(ctx, tc.setup)
+				require.NoError(t, err)
+			}
+			if tc.cleanup != "" {
+				defer func() {
+					_, err := conn.ExecContext(ctx, tc.cleanup)
+					require.NoError(t, err)
+				}()
+			}
 			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)
+				_, err := conn.ExecContext(ctx, tc.query)
 				errCh <- err
 			}()
 			// In the main goroutine, loop until the query is completed while