kvserver: add LBRebalancingMultiMetricAndCount

This commit removes the cluster setting
DisableCountBasedRebalancingIfMMAEnabled, renames LBRebalancingMultiMetric to
LBRebalancingMultiMetricOnly, and introduces a new mode:
LBRebalancingMultiMetricAndCount. To keep the datadriven test output unchanged,
all previous references to mma in asim now use LBRebalancingMultiMetricAndCount,
which matches the prior behavior. I’ll follow up with a new asim mode in future
commits.

Author: wenyihu6

docs/generated/settings/settings.html

@@ -86,7 +86,7 @@
 <tr><td><div id="setting-jobs-retention-time" class="anchored"><code>jobs.retention_time</code></div></td><td>duration</td><td><code>336h0m0s</code></td><td>the amount of time for which records for completed jobs are retained</td><td>Basic/Standard/Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-kv-allocator-lease-rebalance-threshold" class="anchored"><code>kv.allocator.lease_rebalance_threshold</code></div></td><td>float</td><td><code>0.05</code></td><td>minimum fraction away from the mean a store&#39;s lease count can be before it is considered for lease-transfers</td><td>Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-kv-allocator-load-based-lease-rebalancing-enabled" class="anchored"><code>kv.allocator.load_based_lease_rebalancing.enabled</code></div></td><td>boolean</td><td><code>true</code></td><td>set to enable rebalancing of range leases based on load and latency</td><td>Advanced/Self-Hosted</td></tr>
-<tr><td><div id="setting-kv-allocator-load-based-rebalancing" class="anchored"><code>kv.allocator.load_based_rebalancing</code></div></td><td>enumeration</td><td><code>leases and replicas</code></td><td>whether to rebalance based on the distribution of load across stores [off = 0, leases = 1, leases and replicas = 2, multi-metric = 3]</td><td>Advanced/Self-Hosted</td></tr>
+<tr><td><div id="setting-kv-allocator-load-based-rebalancing" class="anchored"><code>kv.allocator.load_based_rebalancing</code></div></td><td>enumeration</td><td><code>leases and replicas</code></td><td>whether to rebalance based on the distribution of load across stores [off = 0, leases = 1, leases and replicas = 2, multi-metric = 3, multi-metric and count = 4]</td><td>Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-kv-allocator-load-based-rebalancing-objective" class="anchored"><code>kv.allocator.load_based_rebalancing.objective</code></div></td><td>enumeration</td><td><code>cpu</code></td><td>what objective does the cluster use to rebalance; if set to `qps` the cluster will attempt to balance qps among stores, if set to `cpu` the cluster will attempt to balance cpu usage among stores [qps = 0, cpu = 1]</td><td>Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-kv-allocator-load-based-rebalancing-interval" class="anchored"><code>kv.allocator.load_based_rebalancing_interval</code></div></td><td>duration</td><td><code>1m0s</code></td><td>the rough interval at which each store will check for load-based lease / replica rebalancing opportunities</td><td>Advanced/Self-Hosted</td></tr>
 <tr><td><div id="setting-kv-allocator-qps-rebalance-threshold" class="anchored"><code>kv.allocator.qps_rebalance_threshold</code></div></td><td>float</td><td><code>0.1</code></td><td>minimum fraction away from the mean a store&#39;s QPS (such as queries per second) can be before it is considered overfull or underfull</td><td>Advanced/Self-Hosted</td></tr>

pkg/kv/kvserver/allocator/allocatorimpl/allocator.go

@@ -2744,12 +2744,12 @@ func (t TransferLeaseDecision) String() string {
 	}
 }
 
-// CountBasedRebalanceDisabled returns true if the count based rebalancing
-// should be disabled. Disabled DisableCountBasedRebalancingIfMMAEnabled is true
-// and mma is enabled (LBRebalancingMultiMetric).
-func (a *Allocator) CountBasedRebalanceDisabled() bool {
-	return kvserverbase.DisableCountBasedRebalancingIfMMAEnabled.Get(&a.st.SV) &&
-		kvserverbase.LoadBasedRebalancingMode.Get(&a.st.SV) == kvserverbase.LBRebalancingMultiMetric
+// CountBasedRebalancingDisabled returns true if count-based rebalancing should
+// be disabled. Count-based rebalancing is disabled only when
+// LBRebalancingMultiMetricOnly mode is active. To enable both multi-metric and
+// count-based rebalancing, use LBRebalancingMultiMetricAndCount mode instead.
+func (a *Allocator) CountBasedRebalancingDisabled() bool {
+	return kvserverbase.LoadBasedRebalancingMode.Get(&a.st.SV) == kvserverbase.LBRebalancingMultiMetricOnly
 }
 
 // ShouldTransferLease returns true if the specified store is overfull in terms
@@ -3052,7 +3052,7 @@ func (a Allocator) shouldTransferLeaseForLeaseCountConvergence(
 	existing []roachpb.ReplicaDescriptor,
 ) bool {
 	// Return false early if count based rebalancing is disabled.
-	if a.CountBasedRebalanceDisabled() {
+	if a.CountBasedRebalancingDisabled() {
 		return false
 	}
 	// TODO(a-robinson): Should we disable this behavior when load-based lease

pkg/kv/kvserver/allocator/plan/replicate.go

@@ -168,7 +168,7 @@ func (rp ReplicaPlanner) ShouldPlanChange(
 	nonVoterReplicas := desc.Replicas().NonVoterDescriptors()
 	if !rp.knobs.DisableReplicaRebalancing {
 		scorerOptions := allocatorimpl.ScorerOptions(rp.allocator.ScorerOptions(ctx))
-		if rp.allocator.CountBasedRebalanceDisabled() {
+		if rp.allocator.CountBasedRebalancingDisabled() {
 			scorerOptions = rp.allocator.BaseScorerOptionsWithNoConvergence()
 		}
 		rangeUsageInfo := repl.RangeUsageInfo()
@@ -793,7 +793,7 @@ func (rp ReplicaPlanner) considerRebalance(
 	if scatter {
 		scorerOpts = rp.allocator.ScorerOptionsForScatter(ctx)
 	}
-	if rp.allocator.CountBasedRebalanceDisabled() {
+	if rp.allocator.CountBasedRebalancingDisabled() {
 		scorerOpts = rp.allocator.BaseScorerOptionsWithNoConvergence()
 	}
 	rangeUsageInfo := repl.RangeUsageInfo()

pkg/kv/kvserver/asim/mmaintegration/mma_store_rebalancer.go

@@ -97,7 +97,7 @@ func (msr *MMAStoreRebalancer) Tick(ctx context.Context, tick time.Time, s state
 		return
 	}
 
-	if kvserverbase.LoadBasedRebalancingMode.Get(&msr.settings.ST.SV) != kvserverbase.LBRebalancingMultiMetric {
+	if !kvserverbase.LoadBasedRebalancingModeIsMMA(&msr.settings.ST.SV) {
 		// When the store rebalancer isn't set to use the multi-metric mode, the
 		// legacy store rebalancer is used.
 		return

pkg/kv/kvserver/asim/tests/datadriven_simulation_test.go

@@ -495,7 +495,7 @@ func TestDataDriven(t *testing.T) {
 							event.SetSimulationSettingsEvent{
 								IsClusterSetting: true,
 								Key:              "LBRebalancingMode",
-								Value:            int64(kvserverbase.LBRebalancingMultiMetric),
+								Value:            int64(kvserverbase.LBRebalancingMultiMetricAndCount),
 							})
 					},
 					// 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(kvserverbase.LBRebalancingMultiMetric),
+								Value:            int64(kvserverbase.LBRebalancingMultiMetricAndCount),
 							})
 					},
 				}

pkg/kv/kvserver/asim/tests/testdata/non_rand/load_distribution_movement_disabled_enable_later.txt

@@ -39,7 +39,7 @@ gen_load rate=10000 rw_ratio=0 min_block=1000 max_block=1000 min_key=10001 max_k
 ----
 
 # Two minutes into the test, we'll enable the MMA.
-setting rebalance_mode=3 delay=2m
+setting rebalance_mode=4 delay=2m
 ----
 
 eval duration=7m samples=1 seed=42 cfgs=(sma-only) metrics=(cpu,write_bytes_per_second,replicas,leases)

pkg/kv/kvserver/kvserverbase/base.go

@@ -126,10 +126,11 @@ var LoadBasedRebalancingMode = settings.RegisterEnumSetting(
 	"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",
+		LBRebalancingOff:                 "off",
+		LBRebalancingLeasesOnly:          "leases",
+		LBRebalancingLeasesAndReplicas:   "leases and replicas",
+		LBRebalancingMultiMetricOnly:     "multi-metric",
+		LBRebalancingMultiMetricAndCount: "multi-metric and count",
 	},
 	settings.WithPublic,
 	settings.WithValidateEnum(func(enumStr string) error {
@@ -141,6 +142,13 @@ var LoadBasedRebalancingMode = settings.RegisterEnumSetting(
 	}),
 )
 
+// LoadBasedRebalancingModeIsMMA returns true if the load-based rebalancing mode
+// uses the multi-metric store rebalancer.
+var LoadBasedRebalancingModeIsMMA = func(sv *settings.Values) bool {
+	mode := LoadBasedRebalancingMode.Get(sv)
+	return mode == LBRebalancingMultiMetricOnly || mode == LBRebalancingMultiMetricAndCount
+}
+
 // LBRebalancingMode controls if and when we do store-level rebalancing
 // based on load.
 type LBRebalancingMode int64
@@ -155,20 +163,17 @@ const (
 	// 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
+	// LBRebalancingMultiMetricOnly means that the store rebalancer yields to the
 	// multi-metric store rebalancer, balancing both leases and replicas based on
-	// store-level load imbalances.
-	LBRebalancingMultiMetric
-)
-
-// DisableCountBasedRebalancingIfMMAEnabled is a setting that controls whether
-// to disable replica and lease count based rebalancing if multi-metric
-// allocator is enabled.
-var DisableCountBasedRebalancingIfMMAEnabled = settings.RegisterBoolSetting(
-	settings.SystemOnly,
-	"kv.allocator.disable_count_based_rebalancing_with_mma.enabled",
-	"whether to disable replica and lease count based rebalancing if multi-metric allocator is enabled",
-	false,
+	// store-level load imbalances. Note that this disables replica-count and
+	// lease-count based rebalancing.
+	LBRebalancingMultiMetricOnly
+	// LBRebalancingMultiMetricAndCount means that both multi-metric store
+	// rebalancer and count based rebalancing via lease queue and replicate queue
+	// are enabled, balancing lease count, replica count, and store-level load
+	// across stores. Note that this might cause more thrashing since lease and
+	// replica counts goal may be in conflict with the store-level load goal.
+	LBRebalancingMultiMetricAndCount
 )
 
 // RangeFeedRefreshInterval is injected from kvserver to avoid import cycles

pkg/kv/kvserver/mma_store_rebalancer.go

@@ -81,7 +81,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 kvserverbase.LoadBasedRebalancingMode.Get(&m.st.SV) != kvserverbase.LBRebalancingMultiMetric {
+			if !kvserverbase.LoadBasedRebalancingModeIsMMA(&m.st.SV) {
 				continue
 			}
 

pkg/kv/kvserver/mmaintegration/allocator_sync.go

@@ -133,7 +133,7 @@ func (as *AllocatorSync) NonMMAPreTransferLease(
 	transferFrom, transferTo roachpb.ReplicationTarget,
 ) SyncChangeID {
 	var changeIDs []mmaprototype.ChangeID
-	if kvserverbase.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserverbase.LBRebalancingMultiMetric {
+	if kvserverbase.LoadBasedRebalancingModeIsMMA(&as.st.SV) {
 		changeIDs = as.mmaAllocator.RegisterExternalChanges(convertLeaseTransferToMMA(desc, usage, transferFrom, transferTo))
 	}
 	trackedChange := trackedAllocatorChange{
@@ -158,7 +158,7 @@ func (as *AllocatorSync) NonMMAPreChangeReplicas(
 	leaseholderStoreID roachpb.StoreID,
 ) SyncChangeID {
 	var changeIDs []mmaprototype.ChangeID
-	if kvserverbase.LoadBasedRebalancingMode.Get(&as.st.SV) == kvserverbase.LBRebalancingMultiMetric {
+	if kvserverbase.LoadBasedRebalancingModeIsMMA(&as.st.SV) {
 		changeIDs = as.mmaAllocator.RegisterExternalChanges(convertReplicaChangeToMMA(desc, usage, changes, leaseholderStoreID))
 	}
 	trackedChange := trackedAllocatorChange{

pkg/kv/kvserver/store.go

@@ -1183,7 +1183,8 @@ type StoreConfig struct {
 	StoreLiveness     *storeliveness.NodeContainer
 	StorePool         *storepool.StorePool
 	// One MMAllocator per node which guides mma store rebalancer to make
-	// allocation changes when LBRebalancingMultiMetric is enabled.
+	// allocation changes when
+	// LBRebalancingMultiMetricOnly/LBRebalancingMultiMetricAndCount is enabled.
 	MMAllocator          mmaprototype.Allocator
 	AllocatorSync        *mmaintegration.AllocatorSync
 	Transport            *RaftTransport