kvserver: move LoadBasedRebalancingMode to kvserverbase

wenyihu6 · wenyihu6 · commit bbd2bb4c7ffd · 2025-08-04T10:29:55.000-04:00
Previously, the cluster setting LoadBasedRebalancingMode was in the kvserver
package. This would create an import cycle when future commits introduce
allocator sync. Allocator sync (in the mmaprototypehelpers package) needs to
check if mma is enabled by accessing this setting in kvserver, while kvserver
depends on mmaprototypehelpers to initialize allocator sync on every node. To
resolve this, this commit moves the setting out of kvserver to kvserverbase,
breaking the dependency cycle. This also improves clarity, as the setting fits
better in kvserverbase alongside other cluster settings, rather than in the
already large kvserver package.
diff --git a/pkg/ccl/kvccl/kvfollowerreadsccl/followerreads_test.go b/pkg/ccl/kvccl/kvfollowerreadsccl/followerreads_test.go
@@ -844,7 +844,7 @@ func TestFollowerReadsWithStaleDescriptor(t *testing.T) {
 
 	// Further down, we'll set up the test to pin the lease to store 1. Turn off
 	// load based rebalancing to make sure it doesn't move.
-	kvserver.LoadBasedRebalancingMode.Override(ctx, &settings.SV, kvserver.LBRebalancingOff)
+	kvserverbase.LoadBasedRebalancingMode.Override(ctx, &settings.SV, kvserverbase.LBRebalancingOff)
 
 	n1 := sqlutils.MakeSQLRunner(tc.Conns[0])
 	n1.Exec(t, `CREATE DATABASE t`)
diff --git a/pkg/kv/kvserver/BUILD.bazel b/pkg/kv/kvserver/BUILD.bazel
@@ -223,7 +223,6 @@ go_library(
         "//pkg/util/encoding",
         "//pkg/util/envutil",
         "//pkg/util/errorutil",
-        "//pkg/util/errorutil/unimplemented",
         "//pkg/util/growstack",
         "//pkg/util/grpcutil",
         "//pkg/util/grunning",
diff --git a/pkg/kv/kvserver/allocator/mmaprototypehelpers/BUILD.bazel b/pkg/kv/kvserver/allocator/mmaprototypehelpers/BUILD.bazel
@@ -10,10 +10,10 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/kv/kvpb",
-        "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/allocator",
         "//pkg/kv/kvserver/allocator/mmaprototype",
         "//pkg/kv/kvserver/allocator/storepool",
+        "//pkg/kv/kvserver/kvserverbase",
         "//pkg/roachpb",
         "//pkg/settings/cluster",
         "//pkg/util/log",
diff --git a/pkg/kv/kvserver/allocator/mmaprototypehelpers/kvserver_mma_integration.go b/pkg/kv/kvserver/allocator/mmaprototypehelpers/kvserver_mma_integration.go
@@ -10,10 +10,10 @@ import (
 	"fmt"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
-	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -129,7 +129,7 @@ func (as *AllocatorSync) NonMMAPreTransferLease(
 	log.Infof(ctx, "registering external lease transfer change: usage=%v changes=%v",
 		usage, replicaChanges)
 	var changeIDs []mmaprototype.ChangeID
-	if kvserver.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserver.LBRebalancingMultiMetric {
+	if kvserverbase.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserverbase.LBRebalancingMultiMetric {
 		changeIDs = as.mmAllocator.RegisterExternalChanges(replicaChanges[:])
 		if changeIDs == nil {
 			log.Info(ctx, "mma did not track lease transfer, skipping")
@@ -220,7 +220,7 @@ func (as *AllocatorSync) NonMMAPreChangeReplicas(
 	log.Infof(ctx, "registering external replica change: chgs=%v usage=%v changes=%v",
 		changes, usage, replicaChanges)
 	var changeIDs []mmaprototype.ChangeID
-	if kvserver.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserver.LBRebalancingMultiMetric {
+	if kvserverbase.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserverbase.LBRebalancingMultiMetric {
 		changeIDs = as.mmAllocator.RegisterExternalChanges(replicaChanges)
 		if changeIDs == nil {
 			log.Info(ctx, "cluster does not have a range for the external replica change, skipping")
@@ -347,7 +347,7 @@ func (as *AllocatorSync) PostApply(ctx context.Context, syncChangeID SyncChangeI
 		}
 		delete(as.mu.trackedChanges, syncChangeID)
 	}()
-	if kvserver.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserver.LBRebalancingMultiMetric {
+	if kvserverbase.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserverbase.LBRebalancingMultiMetric {
 		if changeIDs := tracked.changeIDs; changeIDs != nil {
 			log.Infof(ctx, "PostApply: tracked=%v change_ids=%v success: %v", tracked, changeIDs, success)
 			as.updateMetrics(success, tracked.typ, tracked.author)
diff --git a/pkg/kv/kvserver/asim/mmaintegration/BUILD.bazel b/pkg/kv/kvserver/asim/mmaintegration/BUILD.bazel
@@ -9,13 +9,13 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/mmaintegration",
     visibility = ["//visibility:public"],
     deps = [
-        "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/allocator",
         "//pkg/kv/kvserver/allocator/mmaprototype",
         "//pkg/kv/kvserver/allocator/mmaprototypehelpers",
         "//pkg/kv/kvserver/asim/config",
         "//pkg/kv/kvserver/asim/op",
         "//pkg/kv/kvserver/asim/state",
+        "//pkg/kv/kvserver/kvserverbase",
         "//pkg/roachpb",
         "//pkg/util/log",
         "@com_github_cockroachdb_logtags//:logtags",
diff --git a/pkg/kv/kvserver/asim/mmaintegration/mma_store_rebalancer.go b/pkg/kv/kvserver/asim/mmaintegration/mma_store_rebalancer.go
@@ -10,13 +10,13 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototypehelpers"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/op"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/state"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/logtags"
@@ -97,7 +97,7 @@ func (msr *MMAStoreRebalancer) Tick(ctx context.Context, tick time.Time, s state
 		return
 	}
 
-	if kvserver.LoadBasedRebalancingMode.Get(&msr.settings.ST.SV) != kvserver.LBRebalancingMultiMetric {
+	if kvserverbase.LoadBasedRebalancingMode.Get(&msr.settings.ST.SV) != kvserverbase.LBRebalancingMultiMetric {
 		// When the store rebalancer isn't set to use the multi-metric mode, the
 		// legacy store rebalancer is used.
 		return
diff --git a/pkg/kv/kvserver/asim/state/BUILD.bazel b/pkg/kv/kvserver/asim/state/BUILD.bazel
@@ -30,6 +30,7 @@ go_library(
         "//pkg/kv/kvserver/allocator/storepool",
         "//pkg/kv/kvserver/asim/config",
         "//pkg/kv/kvserver/asim/workload",
+        "//pkg/kv/kvserver/kvserverbase",
         "//pkg/kv/kvserver/liveness/livenesspb",
         "//pkg/kv/kvserver/load",
         "//pkg/kv/kvserver/split",
diff --git a/pkg/kv/kvserver/asim/state/impl.go b/pkg/kv/kvserver/asim/state/impl.go
@@ -27,6 +27,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/workload"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/raft"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
@@ -1397,7 +1398,7 @@ func (s *state) RegisterConfigChangeListener(listener ConfigChangeListener) {
 func (s *state) SetClusterSetting(Key string, Value interface{}) {
 	switch Key {
 	case "LBRebalancingMode":
-		kvserver.LoadBasedRebalancingMode.Override(context.Background(), &s.settings.ST.SV, kvserver.LBRebalancingMode(Value.(int64)))
+		kvserverbase.LoadBasedRebalancingMode.Override(context.Background(), &s.settings.ST.SV, kvserverbase.LBRebalancingMode(Value.(int64)))
 	default:
 		panic("other cluster settings not supported")
 	}
diff --git a/pkg/kv/kvserver/asim/storerebalancer/BUILD.bazel b/pkg/kv/kvserver/asim/storerebalancer/BUILD.bazel
@@ -20,6 +20,7 @@ go_library(
         "//pkg/kv/kvserver/asim/op",
         "//pkg/kv/kvserver/asim/state",
         "//pkg/kv/kvserver/kvflowcontrol/rac2",
+        "//pkg/kv/kvserver/kvserverbase",
         "//pkg/raft",
         "//pkg/roachpb",
         "//pkg/util/hlc",
diff --git a/pkg/kv/kvserver/asim/storerebalancer/store_rebalancer.go b/pkg/kv/kvserver/asim/storerebalancer/store_rebalancer.go
@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/op"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/state"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/raft"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 )
@@ -190,7 +191,7 @@ func (src *storeRebalancerControl) phasePrologue(
 			s, src.storeID,
 			kvserver.LBRebalancingObjective(src.settings.LBRebalancingObjective).ToDimension(),
 		),
-		kvserver.LoadBasedRebalancingMode.Get(&src.settings.ST.SV),
+		kvserverbase.LoadBasedRebalancingMode.Get(&src.settings.ST.SV),
 	)
 
 	if !src.sr.ShouldRebalanceStore(ctx, rctx) {
diff --git a/pkg/kv/kvserver/asim/tests/BUILD.bazel b/pkg/kv/kvserver/asim/tests/BUILD.bazel
@@ -39,7 +39,6 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":tests"],
     deps = [
-        "//pkg/kv/kvserver",
         "//pkg/kv/kvserver/allocator/allocatorimpl",
         "//pkg/kv/kvserver/asim/assertion",
         "//pkg/kv/kvserver/asim/config",
@@ -49,6 +48,7 @@ go_test(
         "//pkg/kv/kvserver/asim/metrics",
         "//pkg/kv/kvserver/asim/scheduled",
         "//pkg/kv/kvserver/asim/state",
+        "//pkg/kv/kvserver/kvserverbase",
         "//pkg/kv/kvserver/liveness/livenesspb",
         "//pkg/spanconfig/spanconfigtestutils",
         "//pkg/testutils/datapathutils",
diff --git a/pkg/kv/kvserver/asim/tests/datadriven_simulation_test.go b/pkg/kv/kvserver/asim/tests/datadriven_simulation_test.go
@@ -17,7 +17,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/allocatorimpl"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/assertion"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
@@ -26,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/history"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/scheduled"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/state"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/spanconfig/spanconfigtestutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
@@ -485,7 +485,7 @@ func TestDataDriven(t *testing.T) {
 							event.SetSimulationSettingsEvent{
 								IsClusterSetting: true,
 								Key:              "LBRebalancingMode",
-								Value:            int64(kvserver.LBRebalancingLeasesAndReplicas),
+								Value:            int64(kvserverbase.LBRebalancingLeasesAndReplicas),
 							})
 					},
 					"mma-only": func(eg *gen.StaticEvents) {
@@ -495,7 +495,7 @@ func TestDataDriven(t *testing.T) {
 							event.SetSimulationSettingsEvent{
 								IsClusterSetting: true,
 								Key:              "LBRebalancingMode",
-								Value:            int64(kvserver.LBRebalancingMultiMetric),
+								Value:            int64(kvserverbase.LBRebalancingMultiMetric),
 							})
 					},
 					// Both the replicate/lease queues and the MMA are enabled.
@@ -506,7 +506,7 @@ func TestDataDriven(t *testing.T) {
 							event.SetSimulationSettingsEvent{
 								IsClusterSetting: true,
 								Key:              "LBRebalancingMode",
-								Value:            int64(kvserver.LBRebalancingMultiMetric),
+								Value:            int64(kvserverbase.LBRebalancingMultiMetric),
 							})
 					},
 				}
diff --git a/pkg/kv/kvserver/kvserverbase/BUILD.bazel b/pkg/kv/kvserver/kvserverbase/BUILD.bazel
@@ -21,8 +21,10 @@ go_library(
         "//pkg/roachpb",
         "//pkg/settings",
         "//pkg/settings/cluster",
+        "//pkg/util/buildutil",
         "//pkg/util/debugutil",
         "//pkg/util/errorutil",
+        "//pkg/util/errorutil/unimplemented",
         "//pkg/util/hlc",
         "//pkg/util/log",
         "//pkg/util/quotapool",
diff --git a/pkg/kv/kvserver/kvserverbase/base.go b/pkg/kv/kvserver/kvserverbase/base.go
@@ -17,6 +17,8 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
+	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
 	"github.com/cockroachdb/cockroach/pkg/util/quotapool"
 	"github.com/cockroachdb/redact"
 )
@@ -112,6 +114,50 @@ var MVCCGCQueueEnabled = settings.RegisterBoolSetting(
 	true,
 )
 
+// LoadBasedRebalancingMode controls whether range rebalancing takes
+// additional variables such as write load and disk usage into account.
+// If disabled, rebalancing is done purely based on replica count.
+var LoadBasedRebalancingMode = settings.RegisterEnumSetting(
+	settings.SystemOnly,
+	"kv.allocator.load_based_rebalancing",
+	"whether to rebalance based on the distribution of load across stores",
+	"leases and replicas",
+	map[LBRebalancingMode]string{
+		LBRebalancingOff:               "off",
+		LBRebalancingLeasesOnly:        "leases",
+		LBRebalancingLeasesAndReplicas: "leases and replicas",
+		LBRebalancingMultiMetric:       "multi-metric",
+	},
+	settings.WithPublic,
+	settings.WithValidateEnum(func(enumStr string) error {
+		if buildutil.CrdbTestBuild || enumStr != "multi-metric" {
+			return nil
+		}
+		return unimplemented.NewWithIssue(
+			103320, "multi-metric rebalancing not supported for production use")
+	}),
+)
+
+// LBRebalancingMode controls if and when we do store-level rebalancing
+// based on load.
+type LBRebalancingMode int64
+
+const (
+	// LBRebalancingOff means that we do not do store-level rebalancing
+	// based on load statistics.
+	LBRebalancingOff LBRebalancingMode = iota
+	// LBRebalancingLeasesOnly means that we rebalance leases based on
+	// store-level load imbalances.
+	LBRebalancingLeasesOnly
+	// LBRebalancingLeasesAndReplicas means that we rebalance both leases and
+	// replicas based on store-level load imbalances.
+	LBRebalancingLeasesAndReplicas
+	// LBRebalancingMultiMetric means that the store rebalancer yields to the
+	// multi-metric store rebalancer, balancing both leases and replicas based on
+	// store-level load imbalances.
+	LBRebalancingMultiMetric
+)
+
 // RangeFeedRefreshInterval is injected from kvserver to avoid import cycles
 // when accessed from kvcoord.
 var RangeFeedRefreshInterval *settings.DurationSetting
diff --git a/pkg/kv/kvserver/mma_store_rebalancer.go b/pkg/kv/kvserver/mma_store_rebalancer.go
@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -65,7 +66,7 @@ func newMMAStoreRebalancer(
 func (m *mmaStoreRebalancer) run(ctx context.Context, stopper *stop.Stopper) {
 	timer := time.NewTicker(jitteredInterval(allocator.LoadBasedRebalanceInterval.Get(&m.st.SV)))
 	defer timer.Stop()
-	log.Infof(ctx, "starting multi-metric store rebalancer with mode=%v", LoadBasedRebalancingMode.Get(&m.st.SV))
+	log.Infof(ctx, "starting multi-metric store rebalancer with mode=%v", kvserverbase.LoadBasedRebalancingMode.Get(&m.st.SV))
 
 	for {
 		select {
@@ -77,7 +78,7 @@ func (m *mmaStoreRebalancer) run(ctx context.Context, stopper *stop.Stopper) {
 			// Wait out the first tick before doing anything since the store is still
 			// starting up and we might as well wait for some stats to accumulate.
 			timer.Reset(jitteredInterval(allocator.LoadBasedRebalanceInterval.Get(&m.st.SV)))
-			if LoadBasedRebalancingMode.Get(&m.st.SV) != LBRebalancingMultiMetric {
+			if kvserverbase.LoadBasedRebalancingMode.Get(&m.st.SV) != kvserverbase.LBRebalancingMultiMetric {
 				continue
 			}
 
diff --git a/pkg/kv/kvserver/replicate_queue_test.go b/pkg/kv/kvserver/replicate_queue_test.go
@@ -30,6 +30,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/allocatorimpl"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/plan"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
@@ -83,7 +84,7 @@ func TestReplicateQueueRebalance(t *testing.T) {
 	for _, server := range tc.Servers {
 		st := server.ClusterSettings()
 		st.Manual.Store(true)
-		kvserver.LoadBasedRebalancingMode.Override(ctx, &st.SV, kvserver.LBRebalancingOff)
+		kvserverbase.LoadBasedRebalancingMode.Override(ctx, &st.SV, kvserverbase.LBRebalancingOff)
 	}
 
 	const newRanges = 10
diff --git a/pkg/kv/kvserver/store_rebalancer.go b/pkg/kv/kvserver/store_rebalancer.go
diff --git a/pkg/kv/kvserver/store_rebalancer_test.go b/pkg/kv/kvserver/store_rebalancer_test.go
diff --git a/pkg/sql/multitenant_admin_function_test.go b/pkg/sql/multitenant_admin_function_test.go
