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.

Author: DrewKimball

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.

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)
+}

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),

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),

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
 }
 

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)

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,

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

pkg/sql/distsql_running.go

@@ -470,6 +470,11 @@ func (dsp *DistSQLPlanner) setupFlows(
 			setupReq.JobTag = jobTag.ValueStr()
 		}
 	}
+	if 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 {

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,