Merge #154688

154688: mmaprototype: pass []roachpb.StoreID for computeMeansForStoreSet r=tbg a=wenyihu6

Epic: CRDB-55052 
Release note: none

---
**mmaprototype: pass []roachpb.StoreID for computeMeansForStoreSet**

Previously, computeMeansForStoreSet took in a meansForStoreSet struct. Future
PRs will use this helper function to compute mean load summaries for stores
without having a storeIDPostingList handy. To avoid unnecessary construction of
a storeIDPostingList, the function signature now takes meansLoad and
[]roachpb.StoreID directly. This change lets future callers pass a simple slice
of stores. Since future callers may provide slices with duplicate store IDs, the
function now also de-duplicates them internally.

---
**mmaprototype: refactor computeMeansForStoreSet**

Previously, we updated computeMeansForStoreSet to take meansLoad and
[]roachpb.StoreID directly instead of a meansForStoreSet struct. This commit
updates the call sites that only require meansLoad to use it directly, removing
unnecessary use of meansForStoreSet.

---
**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.

---
**mmaprototype: return meansLoad directly for computeMeansForStoreSet**

Previously, computeMeansForStoreSet received *meansForStoreSet, which contained
information like stores needed to compute the mean. A recent commit changed it
to take *meansLoad, and computeMeansForStoreSet doesn’t actually use any fields
as provided information. This commit updates computeMeansForStoreSet to directly
construct and return a meanLoad instead.



Co-authored-by: wenyihu6 <wenyi@cockroachlabs.com>

Author: craig[bot]

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
@@ -473,10 +474,8 @@ func (a *allocatorState) rebalanceStores(
 				if len(candsPL) <= 1 {
 					continue // leaseholder is the only candidate
 				}
-				var means meansForStoreSet
 				clear(scratchNodes)
-				means.stores = candsPL
-				computeMeansForStoreSet(a.cs, &means, scratchNodes)
+				means := computeMeansForStoreSet(a.cs, 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 {
@@ -948,7 +947,7 @@ type candidateInfo struct {
 
 type candidateSet struct {
 	candidates []candidateInfo
-	means      *meansForStoreSet
+	means      *meansLoad
 }
 
 type ignoreLevel uint8
@@ -1390,7 +1389,7 @@ func (a *allocatorState) computeCandidatesForRange(
 			storeLoadSummary: csls,
 		})
 	}
-	cset.means = means
+	cset.means = &means.meansLoad
 	return cset, sheddingSLS
 }
 

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

@@ -1949,7 +1949,7 @@ func (cs *clusterState) canShedAndAddLoad(
 	srcSS *storeState,
 	targetSS *storeState,
 	delta LoadVector,
-	means *meansForStoreSet,
+	means *meansLoad,
 	onlyConsiderTargetCPUSummary bool,
 	overloadedDim LoadDimension,
 ) (canAddLoad bool) {

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

@@ -228,11 +228,14 @@ func (sls storeLoadSummary) SafeFormat(w redact.SafePrinter, _ rune) {
 // storeLoadSummary no longer matches that of storeState.loadSeqNum.
 type meansForStoreSet struct {
 	constraintsDisj
-	stores    storeIDPostingList
+	meansLoad
+	stores         storeIDPostingList
+	storeSummaries map[roachpb.StoreID]storeLoadSummary
+}
+
+type meansLoad struct {
 	storeLoad meanStoreLoad
 	nodeLoad  meanNodeLoad
-
-	storeSummaries map[roachpb.StoreID]storeLoadSummary
 }
 
 var _ mapEntry = &meansForStoreSet{}
@@ -277,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{
@@ -317,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{}{},
 	}
 }
 
@@ -339,7 +344,7 @@ func (mm *meansMemo) getMeans(expr constraintsDisj) *meansForStoreSet {
 	}
 	means.constraintsDisj = expr
 	mm.constraintMatcher.constrainStoresForExpr(expr, &means.stores)
-	computeMeansForStoreSet(mm.loadInfoProvider, means, mm.scratchNodes)
+	means.meansLoad = computeMeansForStoreSet(mm.loadInfoProvider, means.stores, mm.scratchNodes, mm.scratchStores)
 	return means
 }
 
@@ -362,18 +367,30 @@ func (mm *meansMemo) getStoreLoadSummary(
 // It does not do any filtering e.g. the stores can include fdDead stores. It
 // is up to the caller to adjust means.stores if it wants to do filtering.
 //
+// stores may contain duplicate storeIDs, in which case computeMeansForStoreSet
+// should deduplicate processing of the stores. stores should be immutable.
+//
 // TODO: fix callers to exclude stores based on node failure detection, from
 // the mean.
 func computeMeansForStoreSet(
-	loadProvider loadInfoProvider, means *meansForStoreSet, scratchNodes map[roachpb.NodeID]*NodeLoad,
-) {
-	n := len(means.stores)
-	if n == 0 {
-		panic(fmt.Sprintf("no stores for meansForStoreSet: %v", *means))
+	loadProvider loadInfoProvider,
+	stores []roachpb.StoreID,
+	scratchNodes map[roachpb.NodeID]*NodeLoad,
+	scratchStores map[roachpb.StoreID]struct{},
+) (means meansLoad) {
+	if len(stores) == 0 {
+		panic(fmt.Sprintf("no stores for meansForStoreSet: %v", stores))
 	}
 	clear(scratchNodes)
-	for _, storeID := range means.stores {
+	clear(scratchStores)
+	n := 0
+	for _, storeID := range stores {
 		nodeID, sload := loadProvider.getStoreReportedLoad(storeID)
+		if _, ok := scratchStores[storeID]; ok {
+			continue
+		}
+		n++
+		scratchStores[storeID] = struct{}{}
 		for j := range sload.reportedLoad {
 			means.storeLoad.load[j] += sload.reportedLoad[j]
 			if sload.capacity[j] == UnknownCapacity {
@@ -413,6 +430,7 @@ func computeMeansForStoreSet(
 		float64(means.nodeLoad.loadCPU) / float64(means.nodeLoad.capacityCPU)
 	means.nodeLoad.loadCPU /= LoadValue(n)
 	means.nodeLoad.capacityCPU /= LoadValue(n)
+	return means
 }
 
 // loadSummary aggregates across all load dimensions for a store, or a node.