mmaprototype: pass in scratch stores for computeMeansForStoreSet

Previously, computeMeansForStoreSet allocated a new map on every call to
deduplicate the provided stores list. This commit refactors it to reuse a
scratchStores map (similar to scratchNodes). The caller now allocates this map
once and stores it in the function scope or struct, reducing repeated
allocations.

Author: wenyihu6

pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go

@@ -356,6 +356,7 @@ func (a *allocatorState) rebalanceStores(
 	var storesToExclude storeIDPostingList
 	var storesToExcludeForRange storeIDPostingList
 	scratchNodes := map[roachpb.NodeID]*NodeLoad{}
+	scratchStores := map[roachpb.StoreID]struct{}{}
 	// The caller has a fixed concurrency limit it can move ranges at, when it
 	// is the sender of the snapshot. So we don't want to create too many
 	// changes, since then the allocator gets too far ahead of what has been
@@ -475,7 +476,7 @@ func (a *allocatorState) rebalanceStores(
 				}
 				var means meansLoad
 				clear(scratchNodes)
-				computeMeansForStoreSet(a.cs, &means, candsPL, scratchNodes)
+				computeMeansForStoreSet(a.cs, &means, candsPL, scratchNodes, scratchStores)
 				sls := a.cs.computeLoadSummary(ctx, store.StoreID, &means.storeLoad, &means.nodeLoad)
 				log.KvDistribution.VInfof(ctx, 2, "considering lease-transfer r%v from s%v: candidates are %v", rangeID, store.StoreID, candsPL)
 				if sls.dimSummary[CPURate] < overloadSlow {

pkg/kv/kvserver/allocator/mmaprototype/load.go

@@ -280,7 +280,8 @@ type meansMemo struct {
 	constraintMatcher *constraintMatcher
 	meansMap          *clearableMemoMap[constraintsDisj, *meansForStoreSet]
 
-	scratchNodes map[roachpb.NodeID]*NodeLoad
+	scratchNodes  map[roachpb.NodeID]*NodeLoad
+	scratchStores map[roachpb.StoreID]struct{}
 }
 
 var meansForStoreSetSlicePool = sync.Pool{
@@ -320,7 +321,8 @@ func newMeansMemo(
 		constraintMatcher: constraintMatcher,
 		meansMap: newClearableMapMemo[constraintsDisj, *meansForStoreSet](
 			meansForStoreSetAllocator{}, meansForStoreSetSlicePoolImpl{}),
-		scratchNodes: map[roachpb.NodeID]*NodeLoad{},
+		scratchNodes:  map[roachpb.NodeID]*NodeLoad{},
+		scratchStores: map[roachpb.StoreID]struct{}{},
 	}
 }
 
@@ -342,7 +344,7 @@ func (mm *meansMemo) getMeans(expr constraintsDisj) *meansForStoreSet {
 	}
 	means.constraintsDisj = expr
 	mm.constraintMatcher.constrainStoresForExpr(expr, &means.stores)
-	computeMeansForStoreSet(mm.loadInfoProvider, &means.meansLoad, means.stores, mm.scratchNodes)
+	computeMeansForStoreSet(mm.loadInfoProvider, &means.meansLoad, means.stores, mm.scratchNodes, mm.scratchStores)
 	return means
 }
 
@@ -375,20 +377,21 @@ func computeMeansForStoreSet(
 	means *meansLoad,
 	stores []roachpb.StoreID,
 	scratchNodes map[roachpb.NodeID]*NodeLoad,
+	scratchStores map[roachpb.StoreID]struct{},
 ) {
 	if len(stores) == 0 {
 		panic(fmt.Sprintf("no stores for meansForStoreSet: %v", *means))
 	}
 	clear(scratchNodes)
-	seen := make(map[roachpb.StoreID]struct{})
+	clear(scratchStores)
 	n := 0
 	for _, storeID := range stores {
 		nodeID, sload := loadProvider.getStoreReportedLoad(storeID)
-		if _, ok := seen[storeID]; ok {
+		if _, ok := scratchStores[storeID]; ok {
 			continue
 		}
 		n++
-		seen[storeID] = struct{}{}
+		scratchStores[storeID] = struct{}{}
 		for j := range sload.reportedLoad {
 			means.storeLoad.load[j] += sload.reportedLoad[j]
 			if sload.capacity[j] == UnknownCapacity {