Merge #144427

craig[bot] · DrewKimball · craig[bot] · commit 7392ad1321f2 · 2025-04-22T01:20:49.000Z
144427: sql,kv: set the admission control header for leaf transactions r=michae2 a=DrewKimball #### kv: set the admission control header for leaf transactions This commit propagates the admission control header from the root transaction to any leaf transactions it creates. This ensures that users of leaf transactions (e.g. streamer, DistSQL, FK checks) correctly respect the user's requested quality of service. Fixes #144421 Release note (bug fix): Fixed a bug that could cause queries that perform work in parallel to ignore the requested quality-of-service level. Affected operations include lookup joins, DistSQL execution, and foreign-key checks. #### sql: add session seting for leaf txn admission control behavior This commit adds a `propagate_admission_header_to_leaf_transactions` session setting to control whether leaf txns inherit the admission header from root txns. The setting is on by default. Informs #144421 Release note: None Co-authored-by: Drew Kimball <drewk@cockroachlabs.com>
diff --git a/pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go b/pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go
@@ -1065,7 +1065,7 @@ func TestTxnMultipleCoord(t *testing.T) {
 	// New create a second, leaf coordinator.
 	leafInputState, err := txn.GetLeafTxnInputState(ctx)
 	require.NoError(t, err)
-	txn2 := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState)
+	txn2 := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState, nil /* header */)
 
 	// Start the second transaction.
 	key2 := roachpb.Key("b")
@@ -2802,7 +2802,7 @@ func TestLeafTxnClientRejectError(t *testing.T) {
 	require.NoError(t, err)
 
 	// New create a second, leaf coordinator.
-	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState, nil /* header */)
 
 	if _, err := leafTxn.Get(ctx, errKey); !testutils.IsError(err, "TransactionAbortedError") {
 		t.Fatalf("expected injected err, got: %v", err)
@@ -2832,7 +2832,7 @@ func TestUpdateRootWithLeafFinalStateInAbortedTxn(t *testing.T) {
 	txn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
 	leafInputState, err := txn.GetLeafTxnInputState(ctx)
 	require.NoError(t, err)
-	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0, leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0, leafInputState, nil /* header */)
 
 	finalState, err := leafTxn.GetLeafTxnFinalState(ctx)
 	if err != nil {
@@ -3038,7 +3038,7 @@ func TestTxnTypeCompatibleWithBatchRequest(t *testing.T) {
 	rootTxn := kv.NewTxn(ctx, s.DB, 0 /* gatewayNodeID */)
 	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
 	require.NoError(t, err)
-	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState, nil /* header */)
 
 	// A LeafTxn is not compatible with locking requests.
 	// 1. Locking Get requests.
diff --git a/pkg/kv/kvclient/kvcoord/txn_test.go b/pkg/kv/kvclient/kvcoord/txn_test.go
@@ -1624,7 +1624,7 @@ func TestTxnBasicBufferedWrites(t *testing.T) {
 			}
 			for _, txnForRead := range []*kv.Txn{
 				txn,
-				kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, tis),
+				kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, tis, nil /* header */),
 			} {
 				if err = readFn(txnForRead); err != nil {
 					return err
@@ -2287,3 +2287,42 @@ func TestTxnBufferedWriteReadYourOwnWrites(t *testing.T) {
 	})
 	require.NoError(t, err)
 }
+
+// TestLeafTransactionAdmissionHeader tests that the admission control header is
+// correctly set for a new leaf txn.
+func TestLeafTransactionAdmissionHeader(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	s := createTestDB(t)
+	defer s.Stop()
+	priorityOptions := []admissionpb.WorkPriority{
+		admissionpb.LowPri,
+		admissionpb.BulkLowPri,
+		admissionpb.UserLowPri,
+		admissionpb.BulkNormalPri,
+		admissionpb.NormalPri,
+		admissionpb.LockingNormalPri,
+		admissionpb.UserHighPri,
+		admissionpb.LockingUserHighPri,
+		admissionpb.HighPri,
+	}
+	rnd, _ := randutil.NewTestRand()
+	priority := priorityOptions[rnd.Intn(len(priorityOptions))]
+
+	ctx := context.Background()
+	rootTxn := kv.NewTxnWithAdmissionControl(
+		ctx, s.DB, 0 /* gatewayNodeID */, kvpb.AdmissionHeader_FROM_SQL, priority)
+	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
+	require.NoError(t, err)
+
+	rootAdmissionHeader := rootTxn.AdmissionHeader()
+	leafTxn := kv.NewLeafTxn(ctx, s.DB, 0 /* gatewayNodeID */, leafInputState, &rootAdmissionHeader)
+	leafHeader := leafTxn.AdmissionHeader()
+	expectedLeafHeader := kvpb.AdmissionHeader{
+		Priority:   int32(priority),
+		CreateTime: rootAdmissionHeader.CreateTime,
+		Source:     kvpb.AdmissionHeader_FROM_SQL,
+	}
+	require.Equal(t, expectedLeafHeader, leafHeader)
+}
diff --git a/pkg/kv/kvclient/kvstreamer/streamer_accounting_test.go b/pkg/kv/kvclient/kvstreamer/streamer_accounting_test.go
@@ -80,7 +80,7 @@ func TestStreamerMemoryAccounting(t *testing.T) {
 		if err != nil {
 			panic(err)
 		}
-		leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState)
+		leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState, nil /* header */)
 		metrics := MakeMetrics()
 		s := NewStreamer(
 			s.DistSenderI().(*kvcoord.DistSender),
diff --git a/pkg/kv/kvclient/kvstreamer/streamer_test.go b/pkg/kv/kvclient/kvstreamer/streamer_test.go
@@ -50,7 +50,7 @@ func getStreamer(
 	if err != nil {
 		panic(err)
 	}
-	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState, nil /* header */)
 	metrics := kvstreamer.MakeMetrics()
 	return kvstreamer.NewStreamer(
 		s.DistSenderI().(*kvcoord.DistSender),
diff --git a/pkg/kv/txn.go b/pkg/kv/txn.go
@@ -231,7 +231,11 @@ func NewTxnFromProto(
 
 // NewLeafTxn instantiates a new leaf transaction.
 func NewLeafTxn(
-	ctx context.Context, db *DB, gatewayNodeID roachpb.NodeID, tis *roachpb.LeafTxnInputState,
+	ctx context.Context,
+	db *DB,
+	gatewayNodeID roachpb.NodeID,
+	tis *roachpb.LeafTxnInputState,
+	header *kvpb.AdmissionHeader,
 ) *Txn {
 	if db == nil {
 		panic(errors.WithContextTags(
@@ -246,6 +250,19 @@ func NewLeafTxn(
 	txn.mu.ID = tis.Txn.ID
 	txn.mu.userPriority = roachpb.NormalUserPriority
 	txn.mu.sender = db.factory.LeafTransactionalSender(tis)
+	if header != nil {
+		if admissionpb.WorkPriority(header.Priority) != admissionpb.NormalPri {
+			log.VEventf(ctx, 2,
+				"initializing leaf txn admission control header with priority: %v",
+				admissionpb.WorkPriority(header.Priority),
+			)
+		}
+		txn.admissionHeader = kvpb.AdmissionHeader{
+			CreateTime: header.CreateTime,
+			Priority:   header.Priority,
+			Source:     header.Source,
+		}
+	}
 	return txn
 }
 
diff --git a/pkg/kv/txn_external_test.go b/pkg/kv/txn_external_test.go
@@ -1100,7 +1100,7 @@ func TestUpdateRootWithLeafFinalStateReadsBelowRefreshTimestamp(t *testing.T) {
 		// Fork off a leaf transaction before the root is refreshed.
 		leafInputState, err := txn.GetLeafTxnInputState(ctx)
 		require.NoError(t, err)
-		leafTxn := kv.NewLeafTxn(ctx, db, 0, leafInputState)
+		leafTxn := kv.NewLeafTxn(ctx, db, 0, leafInputState, nil /* header */)
 
 		writeTS, err := performConflictingWrite(ctx, keyB)
 		require.NoError(t, err)
diff --git a/pkg/sql/colflow/colbatch_scan_test.go b/pkg/sql/colflow/colbatch_scan_test.go
@@ -65,7 +65,7 @@ func TestColBatchScanMeta(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState, nil /* header */)
 	flowCtx := execinfra.FlowCtx{
 		EvalCtx: &evalCtx,
 		Mon:     evalCtx.TestingMon,
diff --git a/pkg/sql/distsql/server.go b/pkg/sql/distsql/server.go
@@ -260,7 +260,8 @@ func (ds *ServerImpl) setupFlow(
 		}
 		// The flow will run in a LeafTxn because we do not want each distributed
 		// Txn to heartbeat the transaction.
-		return kv.NewLeafTxn(ctx, ds.DB.KV(), roachpb.NodeID(req.Flow.Gateway), tis), nil
+		nodeID := roachpb.NodeID(req.Flow.Gateway)
+		return kv.NewLeafTxn(ctx, ds.DB.KV(), nodeID, tis, &req.LeafTxnAdmissionHeader), nil
 	}
 
 	var evalCtx *eval.Context
diff --git a/pkg/sql/distsql_running.go b/pkg/sql/distsql_running.go
@@ -470,6 +470,11 @@ func (dsp *DistSQLPlanner) setupFlows(
 			setupReq.JobTag = jobTag.ValueStr()
 		}
 	}
+	if evalCtx.SessionData().PropagateAdmissionHeaderToLeafTransactions && localState.Txn != nil {
+		// Propagate the admission control header so that leaf transactions
+		// correctly inherit it.
+		setupReq.LeafTxnAdmissionHeader = localState.Txn.AdmissionHeader()
+	}
 
 	var isVectorized bool
 	if vectorizeMode := evalCtx.SessionData().VectorizeMode; vectorizeMode != sessiondatapb.VectorizeOff {
diff --git a/pkg/sql/exec_util.go b/pkg/sql/exec_util.go
@@ -4178,6 +4178,10 @@ func (m *sessionDataMutator) SetVectorSearchBeamSize(val int32) {
 	m.data.VectorSearchBeamSize = val
 }
 
+func (m *sessionDataMutator) SetPropagateAdmissionHeaderToLeafTransactions(val bool) {
+	m.data.PropagateAdmissionHeaderToLeafTransactions = val
+}
+
 // Utility functions related to scrubbing sensitive information on SQL Stats.
 
 // quantizeCounts ensures that the Count field in the
diff --git a/pkg/sql/execinfrapb/api.proto b/pkg/sql/execinfrapb/api.proto
@@ -15,6 +15,7 @@ option go_package = "github.com/cockroachdb/cockroach/pkg/sql/execinfrapb";
 import "gogoproto/gogo.proto";
 import "google/protobuf/timestamp.proto";
 
+import "kv/kvpb/api.proto";
 import "roachpb/data.proto";
 import "sql/execinfrapb/data.proto";
 import "sql/execinfrapb/processors.proto";
@@ -35,6 +36,10 @@ message SetupFlowRequest {
   // flows expect to run in a txn, but some, like backfills, don't.
   optional roachpb.LeafTxnInputState leaf_txn_input_state = 7;
 
+  // LeafTxnAdmissionHeader is used to initialize the admission control
+  // header for the flow's txn, if LeafTxnInputState is set.
+  optional roachpb.AdmissionHeader leaf_txn_admission_header = 12 [(gogoproto.nullable) = false];
+
   // Version of distsqlrun protocol; a server accepts a certain range of
   // versions, up to its own version. See server.go for more details.
   optional uint32 version = 5 [(gogoproto.nullable) = false,
diff --git a/pkg/sql/internal_test.go b/pkg/sql/internal_test.go
@@ -487,7 +487,7 @@ func TestInternalExecutorInLeafTxnDoesNotPanic(t *testing.T) {
 
 	ltis, err := rootTxn.GetLeafTxnInputState(ctx)
 	require.NoError(t, err)
-	leafTxn := kv.NewLeafTxn(ctx, kvDB, roachpb.NodeID(1), ltis)
+	leafTxn := kv.NewLeafTxn(ctx, kvDB, roachpb.NodeID(1), ltis, nil /* header */)
 
 	ie := s.InternalExecutor().(*sql.InternalExecutor)
 	_, err = ie.ExecEx(
diff --git a/pkg/sql/logictest/testdata/logic_test/information_schema b/pkg/sql/logictest/testdata/logic_test/information_schema
@@ -4055,6 +4055,7 @@ plan_cache_mode                                            auto
 plpgsql_use_strict_into                                    off
 prefer_lookup_joins_for_fks                                off
 prepared_statements_cache_size                             0 B
+propagate_admission_header_to_leaf_transactions            on
 propagate_input_ordering                                   off
 recursion_depth_limit                                      1000
 register_latch_wait_contention_events                      off
diff --git a/pkg/sql/logictest/testdata/logic_test/pg_catalog b/pkg/sql/logictest/testdata/logic_test/pg_catalog
@@ -3064,6 +3064,7 @@ plan_cache_mode                                            auto                N
 plpgsql_use_strict_into                                    off                 NULL      NULL        NULL        string
 prefer_lookup_joins_for_fks                                off                 NULL      NULL        NULL        string
 prepared_statements_cache_size                             0 B                 NULL      NULL        NULL        string
+propagate_admission_header_to_leaf_transactions            on                  NULL      NULL        NULL        string
 propagate_input_ordering                                   off                 NULL      NULL        NULL        string
 recursion_depth_limit                                      1000                NULL      NULL        NULL        string
 register_latch_wait_contention_events                      off                 NULL      NULL        NULL        string
@@ -3284,6 +3285,7 @@ plan_cache_mode                                            auto                N
 plpgsql_use_strict_into                                    off                 NULL  user     NULL      off                 off
 prefer_lookup_joins_for_fks                                off                 NULL  user     NULL      off                 off
 prepared_statements_cache_size                             0 B                 NULL  user     NULL      0 B                 0 B
+propagate_admission_header_to_leaf_transactions            on                  NULL  user     NULL      on                  on
 propagate_input_ordering                                   off                 NULL  user     NULL      off                 off
 recursion_depth_limit                                      1000                NULL  user     NULL      1000                1000
 register_latch_wait_contention_events                      off                 NULL  user     NULL      off                 off
@@ -3503,6 +3505,7 @@ plan_cache_mode                                            NULL    NULL     NULL
 plpgsql_use_strict_into                                    NULL    NULL     NULL     NULL        NULL
 prefer_lookup_joins_for_fks                                NULL    NULL     NULL     NULL        NULL
 prepared_statements_cache_size                             NULL    NULL     NULL     NULL        NULL
+propagate_admission_header_to_leaf_transactions            NULL    NULL     NULL     NULL        NULL
 propagate_input_ordering                                   NULL    NULL     NULL     NULL        NULL
 recursion_depth_limit                                      NULL    NULL     NULL     NULL        NULL
 register_latch_wait_contention_events                      NULL    NULL     NULL     NULL        NULL
diff --git a/pkg/sql/logictest/testdata/logic_test/show_source b/pkg/sql/logictest/testdata/logic_test/show_source
@@ -180,6 +180,7 @@ plan_cache_mode                                            auto
 plpgsql_use_strict_into                                    off
 prefer_lookup_joins_for_fks                                off
 prepared_statements_cache_size                             0 B
+propagate_admission_header_to_leaf_transactions            on
 propagate_input_ordering                                   off
 recursion_depth_limit                                      1000
 register_latch_wait_contention_events                      off
diff --git a/pkg/sql/opt/exec/execbuilder/testdata/lookup_join b/pkg/sql/opt/exec/execbuilder/testdata/lookup_join
@@ -2605,3 +2605,49 @@ vectorized: true
 
 statement ok
 RESET optimizer_plan_lookup_joins_with_reverse_scans
+
+# Make sure the admission control priority is correctly set for leaf txns.
+subtest regression_144421
+
+statement ok
+SET default_transaction_quality_of_service=background
+
+statement ok
+SET TRACING = "on", cluster;
+
+statement ok
+SELECT small.a, large.c FROM small INNER LOOKUP JOIN large ON small.a = large.b
+
+statement ok
+SET TRACING = "off";
+
+query T
+SELECT message FROM [SHOW TRACE FOR SESSION] WHERE MESSAGE LIKE 'initializing leaf txn%';
+----
+initializing leaf txn admission control header with priority: user-low-pri
+initializing leaf txn admission control header with priority: user-low-pri
+
+# Turning off admission header propagation should result in the old behavior.
+statement ok
+SET propagate_admission_header_to_leaf_transactions = false;
+
+statement ok
+SET TRACING = "on", cluster;
+
+statement ok
+SELECT small.a, large.c FROM small INNER LOOKUP JOIN large ON small.a = large.b
+
+statement ok
+SET TRACING = "off";
+
+query T
+SELECT message FROM [SHOW TRACE FOR SESSION] WHERE MESSAGE LIKE 'initializing leaf txn%';
+----
+
+statement ok
+RESET propagate_admission_header_to_leaf_transactions;
+
+statement ok
+RESET default_transaction_quality_of_service
+
+subtest end
diff --git a/pkg/sql/rowexec/inverted_joiner_test.go b/pkg/sql/rowexec/inverted_joiner_test.go
@@ -776,7 +776,7 @@ func TestInvertedJoinerDrain(t *testing.T) {
 	rootTxn := kv.NewTxn(ctx, s.DB(), s.NodeID())
 	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
 	require.NoError(t, err)
-	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState, nil /* header */)
 	flowCtx := execinfra.FlowCtx{
 		EvalCtx: &evalCtx,
 		Mon:     evalCtx.TestingMon,
diff --git a/pkg/sql/rowexec/joinreader_test.go b/pkg/sql/rowexec/joinreader_test.go
@@ -1366,7 +1366,7 @@ func TestJoinReaderDrain(t *testing.T) {
 	rootTxn := kv.NewTxn(ctx, s.DB(), s.NodeID())
 	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
 	require.NoError(t, err)
-	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState, nil /* header */)
 
 	flowCtx := execinfra.FlowCtx{
 		EvalCtx: &evalCtx,
diff --git a/pkg/sql/rowexec/tablereader_test.go b/pkg/sql/rowexec/tablereader_test.go
@@ -330,7 +330,7 @@ func TestTableReaderDrain(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState, nil /* header */)
 	flowCtx := execinfra.FlowCtx{
 		EvalCtx: &evalCtx,
 		Mon:     evalCtx.TestingMon,
diff --git a/pkg/sql/rowexec/zigzagjoiner_test.go b/pkg/sql/rowexec/zigzagjoiner_test.go
@@ -760,7 +760,7 @@ func TestZigzagJoinerDrain(t *testing.T) {
 	rootTxn := kv.NewTxn(ctx, s.DB(), s.NodeID())
 	leafInputState, err := rootTxn.GetLeafTxnInputState(ctx)
 	require.NoError(t, err)
-	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState)
+	leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.NodeID(), leafInputState, nil /* header */)
 	flowCtx := execinfra.FlowCtx{
 		EvalCtx: &evalCtx,
 		Mon:     evalCtx.TestingMon,
diff --git a/pkg/sql/sessiondatapb/local_only_session_data.proto b/pkg/sql/sessiondatapb/local_only_session_data.proto
@@ -664,6 +664,9 @@ message LocalOnlySessionData {
   // are considered at each level of the search tree. Higher values increase
   // search accuracy but require more processing resources.
   int32 vector_search_beam_size = 168;
+  // PropagateAdmissionHeaderToLeafTransactions, when true, causes leaf
+  // transactions to inherit the admission header from the root transaction.
+  bool propagate_admission_header_to_leaf_transactions = 169;
   ///////////////////////////////////////////////////////////////////////////
   // WARNING: consider whether a session parameter you're adding needs to  //
   // be propagated to the remote nodes. If so, that parameter should live  //
diff --git a/pkg/sql/vars.go b/pkg/sql/vars.go
@@ -4000,6 +4000,23 @@ var varGen = map[string]sessionVar{
 			return "32"
 		},
 	},
+
+	// CockroachDB extension.
+	`propagate_admission_header_to_leaf_transactions`: {
+		GetStringVal: makePostgresBoolGetStringValFn(`propagate_admission_header_to_leaf_transactions`),
+		Set: func(_ context.Context, m sessionDataMutator, s string) error {
+			b, err := paramparse.ParseBoolVar("propagate_admission_header_to_leaf_transactions", s)
+			if err != nil {
+				return err
+			}
+			m.SetPropagateAdmissionHeaderToLeafTransactions(b)
+			return nil
+		},
+		Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {
+			return formatBoolAsPostgresSetting(evalCtx.SessionData().PropagateAdmissionHeaderToLeafTransactions), nil
+		},
+		GlobalDefault: globalTrue,
+	},
 }
 
 func ReplicationModeFromString(s string) (sessiondatapb.ReplicationMode, error) {

Original file line number	Diff line number	Diff line change
`@@ -80,7 +80,7 @@ func TestStreamerMemoryAccounting(t *testing.T) {`
`80`	`80`	`if err != nil {`
`81`	`81`	`panic(err)`
`82`	`82`	`}`
`83`		`- leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState)`
	`83`	`+ leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState, nil /* header */)`
`84`	`84`	`metrics := MakeMetrics()`
`85`	`85`	`s := NewStreamer(`
`86`	`86`	`s.DistSenderI().(*kvcoord.DistSender),`
Original file line number	Diff line number	Diff line change
`@@ -50,7 +50,7 @@ func getStreamer(`
`50`	`50`	`if err != nil {`
`51`	`51`	`panic(err)`
`52`	`52`	`}`
`53`		`- leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState)`
	`53`	`+ leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState, nil /* header */)`
`54`	`54`	`metrics := kvstreamer.MakeMetrics()`
`55`	`55`	`return kvstreamer.NewStreamer(`
`56`	`56`	`s.DistSenderI().(*kvcoord.DistSender),`
Original file line number	Diff line number	Diff line change
`@@ -65,7 +65,7 @@ func TestColBatchScanMeta(t *testing.T) {`
`65`	`65`	`if err != nil {`
`66`	`66`	`t.Fatal(err)`
`67`	`67`	`}`
`68`		`- leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState)`
	`68`	`+ leafTxn := kv.NewLeafTxn(ctx, s.DB(), s.DistSQLPlanningNodeID(), leafInputState, nil /* header */)`
`69`	`69`	`flowCtx := execinfra.FlowCtx{`
`70`	`70`	`EvalCtx: &evalCtx,`
`71`	`71`	`Mon: evalCtx.TestingMon,`
Original file line number	Diff line number	Diff line change
`@@ -260,7 +260,8 @@ func (ds *ServerImpl) setupFlow(`
`260`	`260`	`}`
`261`	`261`	`// The flow will run in a LeafTxn because we do not want each distributed`
`262`	`262`	`// Txn to heartbeat the transaction.`
`263`		`- return kv.NewLeafTxn(ctx, ds.DB.KV(), roachpb.NodeID(req.Flow.Gateway), tis), nil`
	`263`	`+ nodeID := roachpb.NodeID(req.Flow.Gateway)`
	`264`	`+ return kv.NewLeafTxn(ctx, ds.DB.KV(), nodeID, tis, &req.LeafTxnAdmissionHeader), nil`
`264`	`265`	`}`
`265`	`266`
`266`	`267`	`var evalCtx *eval.Context`
Original file line number	Diff line number	Diff line change
`@@ -470,6 +470,11 @@ func (dsp *DistSQLPlanner) setupFlows(`
`470`	`470`	`setupReq.JobTag = jobTag.ValueStr()`
`471`	`471`	`}`
`472`	`472`	`}`
	`473`	`+ if evalCtx.SessionData().PropagateAdmissionHeaderToLeafTransactions && localState.Txn != nil {`
	`474`	`+ // Propagate the admission control header so that leaf transactions`
	`475`	`+ // correctly inherit it.`
	`476`	`+ setupReq.LeafTxnAdmissionHeader = localState.Txn.AdmissionHeader()`
	`477`	`+ }`
`473`	`478`
`474`	`479`	`var isVectorized bool`
`475`	`480`	`if vectorizeMode := evalCtx.SessionData().VectorizeMode; vectorizeMode != sessiondatapb.VectorizeOff {`
Original file line number	Diff line number	Diff line change
`@@ -4178,6 +4178,10 @@ func (m *sessionDataMutator) SetVectorSearchBeamSize(val int32) {`
`4178`	`4178`	`m.data.VectorSearchBeamSize = val`
`4179`	`4179`	`}`
`4180`	`4180`
	`4181`	`+func (m *sessionDataMutator) SetPropagateAdmissionHeaderToLeafTransactions(val bool) {`
	`4182`	`+ m.data.PropagateAdmissionHeaderToLeafTransactions = val`
	`4183`	`+}`
	`4184`	`+`
`4181`	`4185`	`// Utility functions related to scrubbing sensitive information on SQL Stats.`
`4182`	`4186`
`4183`	`4187`	`// quantizeCounts ensures that the Count field in the`