Merge #152181

152181: kvadmission,server: use the range's tenant ID for admission control r=arulajmani a=sumeerbhola

This behavior applies only to non-admin requests, and if the cluster setting is not disabled.

Fixes #135853

Epic: none

Release note (ops change): The new cluster setting kvadmission.use_range_tenant_id_for_non_admin.enabled can be used to disable the behavior where admission control uses the range's tenantID for non-admin requests. The default behavior is enabled.

Co-authored-by: sumeerbhola <[email protected]>

Author: craig[bot]

pkg/BUILD.bazel

@@ -261,6 +261,7 @@ ALL_TESTS = [
     "//pkg/kv/kvserver/gc:gc_test",
     "//pkg/kv/kvserver/idalloc:idalloc_test",
     "//pkg/kv/kvserver/intentresolver:intentresolver_test",
+    "//pkg/kv/kvserver/kvadmission:kvadmission_test",
     "//pkg/kv/kvserver/kvflowcontrol/node_rac2:node_rac2_test",
     "//pkg/kv/kvserver/kvflowcontrol/rac2:rac2_test",
     "//pkg/kv/kvserver/kvflowcontrol/replica_rac2:replica_rac2_test",
@@ -1553,6 +1554,7 @@ GO_TARGETS = [
     "//pkg/kv/kvserver/intentresolver:intentresolver",
     "//pkg/kv/kvserver/intentresolver:intentresolver_test",
     "//pkg/kv/kvserver/kvadmission:kvadmission",
+    "//pkg/kv/kvserver/kvadmission:kvadmission_test",
     "//pkg/kv/kvserver/kvflowcontrol/kvflowcontrolpb:kvflowcontrolpb",
     "//pkg/kv/kvserver/kvflowcontrol/kvflowinspectpb:kvflowinspectpb",
     "//pkg/kv/kvserver/kvflowcontrol/node_rac2:node_rac2",

pkg/kv/kvserver/kvadmission/BUILD.bazel

@@ -1,4 +1,4 @@
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "kvadmission",
@@ -25,3 +25,17 @@ go_library(
         "@com_github_cockroachdb_pebble//:pebble",
     ],
 )
+
+go_test(
+    name = "kvadmission_test",
+    srcs = ["kvadmission_test.go"],
+    embed = [":kvadmission"],
+    deps = [
+        "//pkg/kv/kvpb",
+        "//pkg/roachpb",
+        "//pkg/settings/cluster",
+        "//pkg/util/admission",
+        "//pkg/util/admission/admissionpb",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/kv/kvserver/kvadmission/kvadmission.go

@@ -142,14 +142,32 @@ var ConnectedStoreExpiration = settings.RegisterDurationSetting(
 	5*time.Minute,
 )
 
+// useRangeTenantIDForNonAdminEnabled determines whether the range's tenant ID
+// is used by admission control when called by the system tenant. When false,
+// the requester's tenant ID (i.e., the system tenant ID) is used.
+var useRangeTenantIDForNonAdminEnabled = settings.RegisterBoolSetting(
+	settings.SystemOnly,
+	"kvadmission.use_range_tenant_id_for_non_admin.enabled",
+	"when true, and the caller is the system tenant, the tenantID used by admission control "+
+		"for non-admin requests is overridden to the range's tenantID",
+	true,
+)
+
 // Controller provides admission control for the KV layer.
 type Controller interface {
 	// AdmitKVWork must be called before performing KV work.
 	// BatchRequest.AdmissionHeader and BatchRequest.Replica.StoreID must be
-	// populated for admission to work correctly. If err is non-nil, the
-	// returned handle can be ignored. If err is nil, AdmittedKVWorkDone must be
-	// called after the KV work is done executing.
-	AdmitKVWork(context.Context, roachpb.TenantID, *kvpb.BatchRequest) (Handle, error)
+	// populated for admission to work correctly. The requestTenantID represents
+	// the authenticated caller and must be populated. The rangeTenantID
+	// represents the tenant of the range on which the work is being performed
+	// -- in rare cases it may be unpopulated.
+	//
+	// If err is non-nil, the returned handle can be ignored. If err is nil,
+	// AdmittedKVWorkDone must be called after the KV work is done executing.
+	AdmitKVWork(
+		_ context.Context, requestTenantID roachpb.TenantID, rangeTenantID roachpb.TenantID,
+		_ *kvpb.BatchRequest,
+	) (Handle, error)
 	// AdmittedKVWorkDone is called after the admitted KV work is done
 	// executing.
 	AdmittedKVWorkDone(Handle, *StoreWriteBytes)
@@ -269,50 +287,16 @@ func MakeController(
 // TODO(irfansharif): There's a fair bit happening here and there's no test
 // coverage. Fix that.
 func (n *controllerImpl) AdmitKVWork(
-	ctx context.Context, tenantID roachpb.TenantID, ba *kvpb.BatchRequest,
-) (handle Handle, retErr error) {
-	ah := Handle{tenantID: tenantID}
+	ctx context.Context,
+	requestTenantID roachpb.TenantID,
+	rangeTenantID roachpb.TenantID,
+	ba *kvpb.BatchRequest,
+) (_ Handle, retErr error) {
 	if n.kvAdmissionQ == nil {
-		return ah, nil
-	}
-
-	bypassAdmission := ba.IsAdmin()
-	source := ba.AdmissionHeader.Source
-	if !roachpb.IsSystemTenantID(tenantID.ToUint64()) {
-		// Request is from a SQL node.
-		bypassAdmission = false
-		source = kvpb.AdmissionHeader_FROM_SQL
-	}
-	if source == kvpb.AdmissionHeader_OTHER {
-		bypassAdmission = true
-	}
-	// TODO(abaptist): Revisit and deprecate this setting in v23.1.
-	if admission.KVBulkOnlyAdmissionControlEnabled.Get(&n.settings.SV) {
-		if admissionpb.WorkPriority(ba.AdmissionHeader.Priority) >= admissionpb.NormalPri {
-			bypassAdmission = true
-		}
-	}
-	// LeaseInfo requests are used as makeshift replica health probes by
-	// DistSender circuit breakers, make sure they bypass AC.
-	//
-	// TODO(erikgrinaker): the various bypass conditions here should be moved to
-	// one or more request flags.
-	if ba.IsSingleLeaseInfoRequest() {
-		bypassAdmission = true
-	}
-	createTime := ba.AdmissionHeader.CreateTime
-	if !bypassAdmission && createTime == 0 {
-		// TODO(sumeer): revisit this for multi-tenant. Specifically, the SQL use
-		// of zero CreateTime needs to be revisited. It should use high priority.
-		createTime = timeutil.Now().UnixNano()
-	}
-	admissionInfo := admission.WorkInfo{
-		TenantID:        tenantID,
-		Priority:        admissionpb.WorkPriority(ba.AdmissionHeader.Priority),
-		CreateTime:      createTime,
-		BypassAdmission: bypassAdmission,
+		return Handle{}, nil
 	}
-
+	admissionInfo := workInfoForBatch(n.settings, requestTenantID, rangeTenantID, ba)
+	ah := Handle{tenantID: admissionInfo.TenantID}
 	admissionEnabled := true
 	// Don't subject HeartbeatTxnRequest to the storeAdmissionQ. Even though
 	// it would bypass admission, it would consume a slot. When writes are
@@ -323,13 +307,14 @@ func (n *controllerImpl) AdmitKVWork(
 	if ba.IsWrite() && !ba.IsSingleHeartbeatTxnRequest() {
 		var admitted bool
 		attemptFlowControl := kvflowcontrol.Enabled.Get(&n.settings.SV)
-		if attemptFlowControl && !bypassAdmission {
+		if attemptFlowControl && !admissionInfo.BypassAdmission {
 			kvflowHandle, found := n.kvflowHandles.LookupReplicationAdmissionHandle(ba.RangeID)
 			if !found {
 				return Handle{}, nil
 			}
 			var err error
-			admitted, err = kvflowHandle.Admit(ctx, admissionInfo.Priority, timeutil.FromUnixNanos(createTime))
+			admitted, err = kvflowHandle.Admit(
+				ctx, admissionInfo.Priority, timeutil.FromUnixNanos(admissionInfo.CreateTime))
 			if err != nil {
 				return Handle{}, err
 			} else if admitted {
@@ -654,3 +639,57 @@ func (wb *StoreWriteBytes) Release() {
 	}
 	storeWriteBytesPool.Put(wb)
 }
+
+func workInfoForBatch(
+	st *cluster.Settings,
+	requestTenantID roachpb.TenantID,
+	rangeTenantID roachpb.TenantID,
+	ba *kvpb.BatchRequest,
+) admission.WorkInfo {
+	bypassAdmission := ba.IsAdmin()
+	source := ba.AdmissionHeader.Source
+	tenantID := requestTenantID
+	if requestTenantID.IsSystem() {
+		if useRangeTenantIDForNonAdminEnabled.Get(&st.SV) && !bypassAdmission &&
+			rangeTenantID.IsSet() {
+			tenantID = rangeTenantID
+		}
+		// Else, either it is an admin request (common), or rangeTenantID is not
+		// set (rare), or the cluster setting is disabled, so continue using the
+		// SystemTenantID.
+	} else {
+		// Request is from a SQL node.
+		bypassAdmission = false
+		source = kvpb.AdmissionHeader_FROM_SQL
+	}
+	if source == kvpb.AdmissionHeader_OTHER {
+		bypassAdmission = true
+	}
+	// TODO(abaptist): Revisit and deprecate this setting in v23.1.
+	if admission.KVBulkOnlyAdmissionControlEnabled.Get(&st.SV) {
+		if admissionpb.WorkPriority(ba.AdmissionHeader.Priority) >= admissionpb.NormalPri {
+			bypassAdmission = true
+		}
+	}
+	// LeaseInfo requests are used as makeshift replica health probes by
+	// DistSender circuit breakers, make sure they bypass AC.
+	//
+	// TODO(erikgrinaker): the various bypass conditions here should be moved to
+	// one or more request flags.
+	if ba.IsSingleLeaseInfoRequest() {
+		bypassAdmission = true
+	}
+	createTime := ba.AdmissionHeader.CreateTime
+	if !bypassAdmission && createTime == 0 {
+		// TODO(sumeer): revisit this for multi-tenant. Specifically, the SQL use
+		// of zero CreateTime needs to be revisited. It should use high priority.
+		createTime = timeutil.Now().UnixNano()
+	}
+	admissionInfo := admission.WorkInfo{
+		TenantID:        tenantID,
+		Priority:        admissionpb.WorkPriority(ba.AdmissionHeader.Priority),
+		CreateTime:      createTime,
+		BypassAdmission: bypassAdmission,
+	}
+	return admissionInfo
+}

pkg/kv/kvserver/kvadmission/kvadmission_test.go

@@ -0,0 +1,130 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+// Package kvadmission is the integration layer between KV and admission
+// control.
+
+package kvadmission
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/util/admission"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWorkInfoForBatch(t *testing.T) {
+	testCases := []struct {
+		name                    string
+		isAdmin                 bool
+		isLeaseInfoRequest      bool
+		ah                      kvpb.AdmissionHeader
+		requestTenant           int
+		rangeTenant             int
+		useRangeTenantIDEnabled bool
+		bulkOnlyEnabled         bool
+
+		expected admission.WorkInfo
+	}{
+		{
+			name: "other-from-system-tenant-bypasses",
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.BulkNormalPri), CreateTime: 2, Source: kvpb.AdmissionHeader_OTHER},
+			requestTenant: 1, rangeTenant: 4,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(1),
+				Priority: admissionpb.BulkNormalPri, CreateTime: 2, BypassAdmission: true,
+			}},
+		{
+			name: "other-from-non-system-tenant-does-not-bypass",
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.BulkNormalPri), CreateTime: 2, Source: kvpb.AdmissionHeader_OTHER},
+			requestTenant: 2, rangeTenant: 4,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(2),
+				Priority: admissionpb.BulkNormalPri, CreateTime: 2, BypassAdmission: false,
+			}},
+		{
+			name:               "is-lease-info-request-bypasses",
+			isLeaseInfoRequest: true,
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.BulkNormalPri), CreateTime: 2, Source: kvpb.AdmissionHeader_FROM_SQL},
+			requestTenant: 2, rangeTenant: 4,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(2),
+				Priority: admissionpb.BulkNormalPri, CreateTime: 2, BypassAdmission: true,
+			}},
+		{
+			name:    "is-admin-from-system-tenant-bypasses",
+			isAdmin: true,
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.LowPri), CreateTime: 5, Source: kvpb.AdmissionHeader_ROOT_KV},
+			requestTenant: 1, rangeTenant: 4,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(1),
+				Priority: admissionpb.LowPri, CreateTime: 5, BypassAdmission: true,
+			}},
+		{
+			name:    "is-admin-from-non-system-tenant-does-not-bypass",
+			isAdmin: true,
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.LowPri), CreateTime: 5, Source: kvpb.AdmissionHeader_ROOT_KV},
+			requestTenant: 2, rangeTenant: 4,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(2),
+				Priority: admissionpb.LowPri, CreateTime: 5, BypassAdmission: false,
+			}},
+		{
+			name: "bulk-only-enabled-from-non-system-tenant-bypasses",
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.NormalPri), CreateTime: 5, Source: kvpb.AdmissionHeader_FROM_SQL},
+			requestTenant: 2, rangeTenant: 4, bulkOnlyEnabled: true,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(2),
+				Priority: admissionpb.NormalPri, CreateTime: 5, BypassAdmission: true,
+			}},
+		{
+			name: "system-tenant-uses-range-tenant",
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.NormalPri), CreateTime: 5, Source: kvpb.AdmissionHeader_FROM_SQL},
+			requestTenant: 1, rangeTenant: 4, useRangeTenantIDEnabled: true,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(4),
+				Priority: admissionpb.NormalPri, CreateTime: 5, BypassAdmission: false,
+			}},
+		{
+			name: "non-system-tenant-does-not-use-range-tenant",
+			ah: kvpb.AdmissionHeader{
+				Priority: int32(admissionpb.NormalPri), CreateTime: 5, Source: kvpb.AdmissionHeader_FROM_SQL},
+			requestTenant: 2, rangeTenant: 4, useRangeTenantIDEnabled: true,
+			expected: admission.WorkInfo{
+				TenantID: roachpb.MustMakeTenantID(2),
+				Priority: admissionpb.NormalPri, CreateTime: 5, BypassAdmission: false,
+			}},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			st := cluster.MakeTestingClusterSettings()
+			ctx := context.Background()
+			ba := kvpb.BatchRequest{}
+			ba.AdmissionHeader = tc.ah
+			if tc.isLeaseInfoRequest {
+				ba.Add(&kvpb.LeaseInfoRequest{})
+			} else if tc.isAdmin {
+				ba.Add(&kvpb.AdminSplitRequest{})
+			}
+			useRangeTenantIDForNonAdminEnabled.Override(ctx, &st.SV, tc.useRangeTenantIDEnabled)
+			admission.KVBulkOnlyAdmissionControlEnabled.Override(ctx, &st.SV, tc.bulkOnlyEnabled)
+			workInfo := workInfoForBatch(st, roachpb.MustMakeTenantID(uint64(tc.requestTenant)),
+				roachpb.MustMakeTenantID(uint64(tc.rangeTenant)), &ba)
+			require.Equal(t, tc.expected, workInfo)
+		})
+	}
+}

pkg/server/node.go

@@ -1621,8 +1621,19 @@ func (n *Node) batchInternal(
 		defer log.Event(ctx, "node sending response")
 	}
 
+	var rangeTenantIDForAC roachpb.TenantID
+	// The computation of rangeTenantIDForAC fails open in case of error.
+	if len(args.Requests) > 0 {
+		k, err := keys.Addr(args.Requests[0].GetInner().Header().Key)
+		if err == nil {
+			_, tid, err := keys.DecodeTenantPrefix(roachpb.Key(k))
+			if err == nil {
+				rangeTenantIDForAC = tid
+			}
+		}
+	}
 	tStart := timeutil.Now()
-	handle, err := n.storeCfg.KVAdmissionController.AdmitKVWork(ctx, tenID, args)
+	handle, err := n.storeCfg.KVAdmissionController.AdmitKVWork(ctx, tenID, rangeTenantIDForAC, args)
 	if err != nil {
 		return nil, err
 	}