Merge #154499

154499: allocatorimpl: consult with mma for replicate queue rebalancing r=wenyihu6 a=wenyihu6

Part of: #153520
Epic: CRDB-55052 
Release note: none

---
**allocatorimpl: plumb allocatorSync to Allocator**

This commit adds allocatorSync to the Allocator struct. In the future, Allocator
will call mma directly to check if a change conflicts with mma's goals and could
cause thrashing, skipping candidates that do.

---
**allocatorimpl: consult with mma for replicate queue rebalancing**

Previously, LBRebalancingMultiMetricAndCount allowed both MMA- and count-based
rebalancing, which could lead to conflicting goals and thrashing. This commit
updates the replicate queue to consult MMA on the final rebalancing target
selected. Critical rebalances that repair constraints, disk fullness, or
diversity bypass the MMA check; all other rebalances consult MMA before
proceeding.

Some asim tests showed even more thrashing, possibly due to two factors:
1. The replicate queue may have a stale view of in-progress changes.
2. The mean calculation of replica count now includes overloaded stores that
   might not be viable targets, leading to suboptimal decisions. For example,
   the queue may rebalance from a slightly overfull store toward an
   around-the-mean store, but the mean can be deflated by overloaded stores with
   very low range counts. As a result, slightly overfull store might appear to
   deserve a rebalance even though it won't be able to pick a very good
   candidate in the end.
We are okay with this for now and will revisit if needed.

The lease queue does not yet apply this check; future commits will extend it.

I plan to add more data-driven tests to ensure that constraints are still
repaired even when MMA rejects certain actions. These will be included in
follow-up PRs.

---
**mmaintegration: add TestingKnobs to AllocatorSync**

This commit adds a TestingKnobs field to AllocatorSync and plumbs it through,
enabling future commits to add tests that override decisions in
InConflictWithMMA.

---
**allocatorimpl: add TestAllocatorRebalanceMMAConflict**

This commit adds TestAllocatorRebalanceMMAConflict, which verifies that
allocator.RebalanceTarget allows critical rebalances to proceed even when they
conflict with MMA's goals, and ensures it consults MMA correctly for
non-critical rebalances.

---
**mmaprototype: use meansLoad in MMARebalanceAdvisor**

This commit switches to using meansLoad and passes the stores slice directly to
computeMeansForStoreSet. Once the other PR is merged and this one is rebased, we
should fix-up this commit to allocatorimpl: consult with mma for replicate queue
rebalancing.

---
**allocatorimpl: remove leftover log lines**

This commit removes leftover fmt.Println debug logs.

---
**mmaprototype: return a pointer for BuildMMARebalanceAdvisor**

This commit updates BuildMMARebalanceAdvisor to return a pointer, since future
changes will require passing it around and modifying its fields.

---
**allocatorimpl: populate advisor under rebalanceOptions**

Previously, `bestRebalanceTarget` took a map and populated it with (option
index->advisor) during the loop. This commit changes it to populate the advisor
directly within the `rebalanceOptions` struct as a field.

---
**mmaprototype: lazily initialize existingStoreSLS**

Previously, `MMARebalanceAdvisor` built `existingStoreSLS` during
`BuildMMARebalanceAdvisor`. This commit changes it to store only
`existingStoreID` at `BuildMMARebalanceAdvisor`, and construct the actual
`existingStoreSLS` during the first call to `IsInConflictWithMMA`. This has two
advantages: (1) `existingStoreID` is now part of the advisor struct, making it
available for logging in `IsInConflictWithMMA`; and (2) only
`IsInConflictWithMMA` needs a context for logging, which makes sense for it to
have one.

---
**mmaprototype: pass context properly**

This commit ensures `IsInConflictWithMMA` receives a proper context so that
its logs appear in traces.

---
**mmaprototype: improve logs with IsInConflictWithMMA**

This commit enhances `IsInConflictWithMMA` logs so they appear in traces and
under vmodule, improving visibility into MMA rejections.

---
**allocatorimpl: improve comments for recent changes**

This commit updates comments to reflect changes made in previous commits.

---
**allocatorimpl: fixes TestBestRebalanceTarget**

This commit fixes a test that previously panic because `nil` was passed for
`allocatorSync` in `bestRebalanceTarget`.

---
**mmaprototype: move MMARebalanceAdvisor to rebalance_advisor.go**

This commit moves `MMARebalanceAdvisor` related code from `allocator_state.go` to
`rebalance_advisor.go`.

---
**allocatorimpl: remove CountBasedRebalancingOnlyEnabledByMMA**

This commit removes `CountBasedRebalancingOnlyEnabledByMMA`, which was left
unused.

---
**mmaprototype: improve comments and logs**

This commit resolves misc comments from PR review.




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

Author: craig[bot]

pkg/kv/kvserver/allocator/allocatorimpl/BUILD.bazel

@@ -15,12 +15,14 @@ go_library(
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/allocator",
         "//pkg/kv/kvserver/allocator/load",
+        "//pkg/kv/kvserver/allocator/mmaprototype",
         "//pkg/kv/kvserver/allocator/storepool",
         "//pkg/kv/kvserver/constraint",
         "//pkg/kv/kvserver/kvflowcontrol/rac2",
         "//pkg/kv/kvserver/kvserverbase",
         "//pkg/kv/kvserver/liveness",
         "//pkg/kv/kvserver/liveness/livenesspb",
+        "//pkg/kv/kvserver/mmaintegration",
         "//pkg/kv/kvserver/raftutil",
         "//pkg/raft",
         "//pkg/raft/raftpb",
@@ -53,11 +55,13 @@ go_test(
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/allocator",
         "//pkg/kv/kvserver/allocator/load",
+        "//pkg/kv/kvserver/allocator/mmaprototype",
         "//pkg/kv/kvserver/allocator/storepool",
         "//pkg/kv/kvserver/constraint",
         "//pkg/kv/kvserver/kvflowcontrol/rac2",
         "//pkg/kv/kvserver/liveness",
         "//pkg/kv/kvserver/liveness/livenesspb",
+        "//pkg/kv/kvserver/mmaintegration",
         "//pkg/kv/kvserver/replicastats",
         "//pkg/raft",
         "//pkg/raft/raftpb",
@@ -72,6 +76,7 @@ go_test(
         "//pkg/util/log",
         "//pkg/util/metric",
         "//pkg/util/stop",
+        "//pkg/util/timeutil",
         "//pkg/util/tracing",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_kr_pretty//:pretty",

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

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/constraint"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvflowcontrol/rac2"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/mmaintegration"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/raftutil"
 	"github.com/cockroachdb/cockroach/pkg/raft"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
@@ -619,6 +620,7 @@ type AllocatorMetrics struct {
 // in the cluster.
 type Allocator struct {
 	st            *cluster.Settings
+	as            *mmaintegration.AllocatorSync
 	deterministic bool
 	nodeLatencyFn func(nodeID roachpb.NodeID) (time.Duration, bool)
 	// TODO(aayush): Let's replace this with a *rand.Rand that has a rand.Source
@@ -658,6 +660,7 @@ func makeAllocatorMetrics() AllocatorMetrics {
 // close coupling with the StorePool.
 func MakeAllocator(
 	st *cluster.Settings,
+	as *mmaintegration.AllocatorSync,
 	deterministic bool,
 	nodeLatencyFn func(nodeID roachpb.NodeID) (time.Duration, bool),
 	knobs *allocator.TestingKnobs,
@@ -670,6 +673,7 @@ func MakeAllocator(
 	}
 	allocator := Allocator{
 		st:            st,
+		as:            as,
 		deterministic: deterministic,
 		nodeLatencyFn: nodeLatencyFn,
 		randGen:       makeAllocatorRand(randSource),
@@ -1868,11 +1872,76 @@ func (a Allocator) RebalanceTarget(
 	// would, we don't want to actually rebalance to that target.
 	var target, existingCandidate *candidate
 	var removeReplica roachpb.ReplicationTarget
+
+	// The loop below can iterate multiple times. This is because the
+	// (source,target) pair chosen by bestRebalanceTarget may be rejected either
+	// by the multi-metric allocator or by the check that a moved replica wouldn't
+	// immediately be removable. bestRebalanceTarget mutates the candidate slice
+	// for the given option (=source) to exclude each considered candidate. For
+	// example, the loop may proceed as follows:
+	// - initially, we may consider rebalancing from s1 to either of s2, s3, or
+	// s5, or rebalancing from s6 to either of s2 or s4: options = [s1 ->
+	// [s2,s3,s5], s6 -> [s2,s4]]. Each option here is considered as an
+	// equivalence class.
+	// - bestRebalanceTarget might pick s6->s4, and removes this choice from
+	// `options`. options now becomes [s1->[s2,s3,s5], s6 -> [s2]].
+	// - mma might reject s6->s4, so we loop around.
+	// - next, s1->s3 might be chosen, but fail the removable replica check.
+	// - so we'll begin a third loop with: options is now [s1->[s2,s5], s6 ->
+	// [s2]].
+	// - s6->s2 might be chosen and might succeed, terminating the loop and
+	// proceeding to make the change.
+	//
+	// Note that in general (and in the example) a source store can be considered
+	// multiple times (s6 is considered twice), so we cache the corresponding MMA
+	// advisor to avoid potentially expensive O(store) recomputations. The
+	// corresponding advisor is constructed only once and cached in
+	// results[bestIdx].advisor when the the source store is selected as the best
+	// rebalance target for the first time. After that, bestRebalanceTarget is
+	// free to mutate the cands set of the option. However, MMARebalancerAdvisor
+	// should use the original candidate set union the existing store to compute
+	// the load summary when calling IsInConflictWithMMA. It does so by using the
+	// computed meansLoad summary cached when this option was selected as the best
+	// rebalance target for the first time.
+	var bestIdx int
+
+	// NB: bestRebalanceTarget may modify the candidate set (cands) within each
+	// option in results. However, for each option, the associated source store,
+	// MMARebalanceAdvisor, and their index in results must remain unchanged
+	// throughout the process. This ensures that any cached MMARebalanceAdvisor
+	// continues to correspond to the original candidate set and source store,
+	// even as candidates are removed.
 	for {
-		target, existingCandidate = bestRebalanceTarget(a.randGen, results)
+		target, existingCandidate, bestIdx = bestRebalanceTarget(a.randGen, results, a.as)
 		if target == nil {
 			return zero, zero, "", false
 		}
+		if bestIdx == -1 {
+			log.KvDistribution.Fatalf(ctx, "programmer error: bestIdx is -1 when target is not nil")
+		}
+
+		// Skip mma conflict checks for critical rebalances, which repairs a bad
+		// state such as constraint violation, disk-fullness, and diversity
+		// improvements.
+		if !existingCandidate.isCriticalRebalance(target) {
+			// If the rebalance is not critical, we check if it conflicts with mma's
+			// goal. advisor for bestIdx should always be cached by
+			// bestRebalanceTarget. If mma rejects the rebalance, we will continue to
+			// the next target. Note that bestRebalanceTarget would delete this target
+			// from the candidates set when being selected, so this target will not be
+			// selected again.
+			if advisor := results[bestIdx].advisor; advisor != nil {
+				if a.as.IsInConflictWithMMA(ctx, target.store.StoreID, advisor, false) {
+					continue
+				}
+			} else {
+				if buildutil.CrdbTestBuild {
+					log.KvDistribution.Fatalf(ctx, "expected to find MMA handle for idx %d", bestIdx)
+				} else {
+					log.KvDistribution.Errorf(ctx, "expected to find MMA handle for idx %d", bestIdx)
+				}
+			}
+		}
 
 		// Add a fake new replica to our copy of the replica descriptor so that we can
 		// simulate the removal logic. If we decide not to go with this target, note

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

@@ -14,8 +14,10 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/load"
+	"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/constraint"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/mmaintegration"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
@@ -853,6 +855,35 @@ func (c candidate) less(o candidate) bool {
 	return c.compare(o) < 0
 }
 
+// isCriticalRebalance returns true if the rebalance from source to target is
+// considered as a critical rebalance to repair constraints, disk fullness, or
+// diversity improvements.
+func (source candidate) isCriticalRebalance(target *candidate) bool {
+	// valid is better.
+	if !source.valid && target.valid {
+		return true
+	}
+	// !fullDisk is better.
+	if source.fullDisk && !target.fullDisk {
+		return true
+	}
+	// necessary is better.
+	if !source.necessary && target.necessary {
+		return true
+	}
+	// voterNecessary is better.
+	if !source.voterNecessary && target.voterNecessary {
+		return true
+	}
+	// higher diversityScore is better.
+	if !scoresAlmostEqual(source.diversityScore, target.diversityScore) {
+		if target.diversityScore > source.diversityScore {
+			return true
+		}
+	}
+	return false
+}
+
 // compare is analogous to strcmp in C or string::compare in C++ -- it returns
 // a positive result if c is a better fit for the range than o, 0 if they're
 // equivalent, or a negative result if o is a better fit than c. The magnitude
@@ -1389,6 +1420,11 @@ func candidateListForRemoval(
 type rebalanceOptions struct {
 	existing   candidate
 	candidates candidateList
+	// advisor is lazily initialized by bestRebalanceTarget when this option is
+	// selected as best rebalance target. It is used to determine if a candidate
+	// is in conflict with mma's goals when LBRebalancingMultiMetricAndCount mode
+	// is enabled.
+	advisor *mmaprototype.MMARebalanceAdvisor
 }
 
 // equivalenceClass captures the set of "equivalent" replacement candidates
@@ -1892,18 +1928,22 @@ func rankedCandidateListForRebalancing(
 // bestRebalanceTarget returns the best target to try to rebalance to out of
 // the provided options, and removes it from the relevant candidate list.
 // Also returns the existing replicas that the chosen candidate was compared to.
+// Also returns the index of the best target in the options slice.
 // Returns nil if there are no more targets worth rebalancing to.
+//
+// Contract: responsible for making sure that the returned bestIdx has the
+// corresponding MMA advisor in advisors.
 func bestRebalanceTarget(
-	randGen allocatorRand, options []rebalanceOptions,
-) (target, existingCandidate *candidate) {
-	bestIdx := -1
+	randGen allocatorRand, options []rebalanceOptions, as *mmaintegration.AllocatorSync,
+) (target, existingCandidate *candidate, bestIdx int) {
+	bestIdx = -1
 	var bestTarget *candidate
 	var replaces candidate
 	for i, option := range options {
 		if len(option.candidates) == 0 {
 			continue
 		}
-		target := option.candidates.selectBest(randGen)
+		target = option.candidates.selectBest(randGen)
 		if target == nil {
 			continue
 		}
@@ -1915,14 +1955,24 @@ func bestRebalanceTarget(
 		}
 	}
 	if bestIdx == -1 {
-		return nil, nil
+		return nil, nil, -1 /*bestIdx*/
+	}
+	// For the first time an option in options is selected, build and cache the
+	// corresponding MMA advisor in advisors[bestIdx].
+	if options[bestIdx].advisor == nil {
+		stores := make([]roachpb.StoreID, 0, len(options[bestIdx].candidates))
+		for _, cand := range options[bestIdx].candidates {
+			stores = append(stores, cand.store.StoreID)
+		}
+		options[bestIdx].advisor = as.BuildMMARebalanceAdvisor(options[bestIdx].existing.store.StoreID, stores)
 	}
+
 	// Copy the selected target out of the candidates slice before modifying
 	// the slice. Without this, the returned pointer likely will be pointing
 	// to a different candidate than intended due to movement within the slice.
 	copiedTarget := *bestTarget
 	options[bestIdx].candidates = options[bestIdx].candidates.removeCandidate(copiedTarget)
-	return &copiedTarget, &options[bestIdx].existing
+	return &copiedTarget, &options[bestIdx].existing, bestIdx
 }
 
 // betterRebalanceTarget returns whichever of target1 or target2 is a larger

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

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/constraint"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/mmaintegration"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -436,10 +437,18 @@ func TestBestRebalanceTarget(t *testing.T) {
 	expectedTargets := []roachpb.StoreID{13, 13, 11, 12}
 	expectedExistingRepls := []roachpb.StoreID{3, 2, 1, 1}
 	allocRand := makeAllocatorRand(rand.NewSource(0))
+	ctx := context.Background()
+	stopper, _, _, a, _ := CreateTestAllocatorWithKnobs(ctx, 10, false, /* deterministic */
+		nil /* allocator.TestingKnobs */, &mmaintegration.TestingKnobs{
+			OverrideIsInConflictWithMMA: func(cand roachpb.StoreID) bool {
+				return false
+			},
+		} /* mmaintegration.TestingKnobs */)
+	defer stopper.Stop(ctx)
 	var i int
 	for {
 		i++
-		target, existing := bestRebalanceTarget(allocRand, candidates)
+		target, existing, _ := bestRebalanceTarget(allocRand, candidates, a.as)
 		if len(expectedTargets) == 0 {
 			if target == nil {
 				break