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

Author: DrewKimball

pkg/sql/distsql_running.go

@@ -470,7 +470,7 @@ func (dsp *DistSQLPlanner) setupFlows(
 			setupReq.JobTag = jobTag.ValueStr()
 		}
 	}
-	if localState.Txn != nil {
+	if evalCtx.SessionData().PropagateAdmissionHeaderToLeafTransactions && localState.Txn != nil {
 		// Propagate the admission control header so that leaf transactions
 		// correctly inherit it.
 		setupReq.LeafTxnAdmissionHeader = localState.Txn.AdmissionHeader()

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

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

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

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

pkg/sql/opt/exec/execbuilder/testdata/lookup_join

@@ -2627,6 +2627,26 @@ SELECT message FROM [SHOW TRACE FOR SESSION] WHERE MESSAGE LIKE 'initializing le
 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
 

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

pkg/sql/vars.go

@@ -4001,6 +4001,23 @@ var varGen = map[string]sessionVar{
 			return strconv.FormatInt(vecindex.SearchBeamSizeSetting.Get(sv), 10)
 		},
 	},
+
+	// 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) {