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.

Author: wenyihu6

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

@@ -57,7 +57,7 @@ func MakeAllocatorSync(
 ) *mmaintegration.AllocatorSync {
 	mmAllocator := mmaprototype.NewAllocatorState(timeutil.DefaultTimeSource{},
 		rand.New(rand.NewSource(timeutil.Now().UnixNano())))
-	return mmaintegration.NewAllocatorSync(sp, mmAllocator, st)
+	return mmaintegration.NewAllocatorSync(sp, mmAllocator, st, nil)
 }
 
 var simpleSpanConfig = &roachpb.SpanConfig{

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

@@ -28,24 +28,29 @@ import (
 func CreateTestAllocator(
 	ctx context.Context, numNodes int, deterministic bool,
 ) (*stop.Stopper, *gossip.Gossip, *storepool.StorePool, Allocator, *timeutil.ManualTime) {
-	return CreateTestAllocatorWithKnobs(ctx, numNodes, deterministic, nil /* knobs */)
+	return CreateTestAllocatorWithKnobs(ctx, numNodes, deterministic,
+		nil /* allocator.TestingKnobs */, nil /* mmaintegration.TestingKnobs */)
 }
 
 // CreateTestAllocatorWithKnobs is like `CreateTestAllocator`, but allows the
 // caller to pass in custom TestingKnobs. Stopper must be stopped by
 // the caller.
 func CreateTestAllocatorWithKnobs(
-	ctx context.Context, numNodes int, deterministic bool, knobs *allocator.TestingKnobs,
+	ctx context.Context,
+	numNodes int,
+	deterministic bool,
+	allocatorKnobs *allocator.TestingKnobs,
+	allocSyncKnobs *mmaintegration.TestingKnobs,
 ) (*stop.Stopper, *gossip.Gossip, *storepool.StorePool, Allocator, *timeutil.ManualTime) {
 	st := cluster.MakeTestingClusterSettings()
 	stopper, g, manual, storePool, _ := storepool.CreateTestStorePool(ctx, st,
 		liveness.TestTimeUntilNodeDeadOff, deterministic,
 		func() int { return numNodes },
 		livenesspb.NodeLivenessStatus_LIVE)
 	mmAllocator := mmaprototype.NewAllocatorState(timeutil.DefaultTimeSource{}, rand.New(rand.NewSource(timeutil.Now().UnixNano())))
-	as := mmaintegration.NewAllocatorSync(storePool, mmAllocator, st)
+	as := mmaintegration.NewAllocatorSync(storePool, mmAllocator, st, allocSyncKnobs)
 	a := MakeAllocator(st, as, deterministic, func(id roachpb.NodeID) (time.Duration, bool) {
 		return 0, true
-	}, knobs)
+	}, allocatorKnobs)
 	return stopper, g, storePool, a, manual
 }

pkg/kv/kvserver/asim/state/impl.go

@@ -433,7 +433,7 @@ func (s *state) AddNode(nodeCPUCapacity int64, locality roachpb.Locality) Node {
 		stores:      []StoreID{},
 		mmAllocator: mmAllocator,
 		storepool:   sp,
-		as:          mmaintegration.NewAllocatorSync(sp, mmAllocator, s.settings.ST),
+		as:          mmaintegration.NewAllocatorSync(sp, mmAllocator, s.settings.ST, nil),
 	}
 	s.nodes[nodeID] = node
 	s.SetNodeLiveness(nodeID, livenesspb.NodeLivenessStatus_LIVE)

pkg/kv/kvserver/mmaintegration/BUILD.bazel

@@ -7,6 +7,7 @@ go_library(
         "allocator_sync.go",
         "mma_conversion.go",
         "store_load_msg.go",
+        "testing_knobs.go",
         "thrashing.go",
     ],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/mmaintegration",

pkg/kv/kvserver/mmaintegration/allocator_sync.go

@@ -66,6 +66,7 @@ type mmaState interface {
 // pool. When mma is disabled, its sole purpose is to track and apply changes
 // to the store pool upon success.
 type AllocatorSync struct {
+	knobs        *TestingKnobs
 	sp           storePool
 	st           *cluster.Settings
 	mmaAllocator mmaState
@@ -82,11 +83,14 @@ type AllocatorSync struct {
 	}
 }
 
-func NewAllocatorSync(sp storePool, mmaAllocator mmaState, st *cluster.Settings) *AllocatorSync {
+func NewAllocatorSync(
+	sp storePool, mmaAllocator mmaState, st *cluster.Settings, knobs *TestingKnobs,
+) *AllocatorSync {
 	as := &AllocatorSync{
 		sp:           sp,
 		st:           st,
 		mmaAllocator: mmaAllocator,
+		knobs:        knobs,
 	}
 	as.mu.trackedChanges = make(map[SyncChangeID]trackedAllocatorChange)
 	return as

pkg/kv/kvserver/mmaintegration/testing_knobs.go

@@ -0,0 +1,15 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package mmaintegration
+
+import "github.com/cockroachdb/cockroach/pkg/roachpb"
+
+// TestingKnobs is used for testing purposes to tune the behavior of mma.
+type TestingKnobs struct {
+	// OverrideIsInConflictWithMMA is used to override the value returned by
+	// isInConflictWithMMA.
+	OverrideIsInConflictWithMMA func(cand roachpb.StoreID) bool
+}

pkg/kv/kvserver/mmaintegration/thrashing.go

@@ -103,5 +103,8 @@ func (as *AllocatorSync) BuildMMARebalanceAdvisor(
 func (as *AllocatorSync) IsInConflictWithMMA(
 	cand roachpb.StoreID, advisor mmaprototype.MMARebalanceAdvisor, cpuOnly bool,
 ) bool {
-	return as.mmaAllocator.IsInConflictWithMMA(cand, advisor, cpuOnly)
+	if as.knobs != nil && as.knobs.OverrideIsInConflictWithMMA != nil {
+		return as.knobs.OverrideIsInConflictWithMMA(cand)
+	}
+    return as.mmaAllocator.IsInConflictWithMMA(cand, advisor, cpuOnly)
 }

pkg/kv/kvserver/store_rebalancer_test.go

@@ -582,7 +582,7 @@ func TestChooseLeaseToTransfer(t *testing.T) {
 			// of snapshots, in order to avoid mocking out a fake raft group for the
 			// `replicaMayNeedSnapshot` checks inside `TransferLeaseTarget`.
 			AllowLeaseTransfersToReplicasNeedingSnapshots: true,
-		},
+		}, nil, /*allocSyncKnobs*/
 	)
 	defer stopper.Stop(context.Background())
 	objectiveProvider := &testRebalanceObjectiveProvider{}
@@ -1230,7 +1230,8 @@ func TestChooseRangeToRebalanceAcrossHeterogeneousZones(t *testing.T) {
 		t.Run(tc.name, withQPSCPU(t, objectiveProvider, func(t *testing.T) {
 			// Boilerplate for test setup.
 			testingKnobs := allocator.TestingKnobs{RaftStatusFn: TestingRaftStatusFn}
-			stopper, g, sp, a, _ := allocatorimpl.CreateTestAllocatorWithKnobs(ctx, 10, false /* deterministic */, &testingKnobs)
+			stopper, g, sp, a, _ := allocatorimpl.CreateTestAllocatorWithKnobs(
+				ctx, 10, false /* deterministic */, &testingKnobs, nil /*allocSyncKnobs*/)
 			defer stopper.Stop(context.Background())
 			gossiputil.NewStoreGossiper(g).GossipStores(multiRegionStores, t)
 
@@ -1329,6 +1330,7 @@ func TestChooseRangeToRebalanceIgnoresRangeOnBestStores(t *testing.T) {
 		10,
 		false, /* deterministic */
 		&allocator.TestingKnobs{AllowLeaseTransfersToReplicasNeedingSnapshots: true},
+		nil, /*allocSyncKnobs*/
 	)
 	defer stopper.Stop(context.Background())
 
@@ -1592,6 +1594,7 @@ func TestNoLeaseTransferToBehindReplicas(t *testing.T) {
 			AllowLeaseTransfersToReplicasNeedingSnapshots: false,
 			RaftStatusFn: behindTestingRaftStatusFn,
 		},
+		nil, /*allocSyncKnobs*/
 	)
 	defer stopper.Stop(context.Background())
 

pkg/server/server.go

@@ -904,7 +904,7 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 
 	mmaAllocator := mmaprototype.NewAllocatorState(timeutil.DefaultTimeSource{},
 		rand.New(rand.NewSource(timeutil.Now().UnixNano())))
-	allocatorSync := mmaintegration.NewAllocatorSync(storePool, mmaAllocator, st)
+	allocatorSync := mmaintegration.NewAllocatorSync(storePool, mmaAllocator, st, nil)
 	g.RegisterCallback(
 		gossip.MakePrefixPattern(gossip.KeyStoreDescPrefix),
 		func(_ string, content roachpb.Value, origTimestampNanos int64) {