mmaprototype: extract newRebalanceEnv

This commit

- creates a proper constructor for rebalanceEnv
- moves rebalanceStores from the clusterState receiver onto rebalanceEnv
- moves some consts and improves a number of comments.

This all prepares for more testing of rebalanceStores via TestClusterState,
which will require fiddling with the limits such as maxRangeMoveCount.
Having these consts only inform the default configuration of a rebalanceEnv
will allow the test to update these values on the create rebalanceEnv before
calling into rebalanceStores.

Epic: CRDB-55052

Author: tbg

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

@@ -213,6 +213,8 @@ type allocatorState struct {
 var _ Allocator = &allocatorState{}
 
 // TODO(sumeer): temporary constants.
+// TODO(tbg): avoid direct access to this constant so that it can be configured
+// in tests. As is, it can be overridden in some places but not others.
 const (
 	maxFractionPendingThreshold = 0.1
 )
@@ -333,7 +335,8 @@ func (a *allocatorState) ComputeChanges(
 		panic(fmt.Sprintf("ComputeChanges: expected StoreID %d, got %d", opts.LocalStoreID, msg.StoreID))
 	}
 	a.cs.processStoreLeaseholderMsg(ctx, msg, a.mmaMetrics)
-	return a.cs.rebalanceStores(ctx, opts.LocalStoreID, a.rand, a.diversityScoringMemo)
+	re := newRebalanceEnv(a.cs, a.rand, a.diversityScoringMemo, a.cs.ts.Now())
+	return re.rebalanceStores(ctx, opts.LocalStoreID)
 }
 
 // AdminRelocateOne implements the Allocator interface.

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

@@ -880,7 +880,7 @@ type storeState struct {
 	localityTiers
 
 	// Time when this store started to be observed as overloaded. Set by
-	// clusterState.rebalanceStores.
+	// rebalanceStores.
 	overloadStartTime time.Time
 	// When overloaded this is equal to time.Time{}.
 	overloadEndTime time.Time

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

@@ -24,6 +24,13 @@ import (
 var mmaid = atomic.Int64{}
 
 // rebalanceEnv tracks the state and outcomes of a rebalanceStores invocation.
+// Recall that such an invocation is on behalf of a local store, but will
+// iterate over a slice of shedding stores - these are not necessarily local
+// but are instead stores which are overloaded and need to shed load. Typically,
+// the particular calling local store has at least some leases (i.e. is
+// capable of making decisions), though in edge cases this may not be true,
+// in which case the rebalanceStores invocation will simply not be able to
+// do anything.
 type rebalanceEnv struct {
 	*clusterState
 	// rng is the random number generator for non-deterministic decisions.
@@ -36,13 +43,20 @@ type rebalanceEnv struct {
 	rangeMoveCount int
 	// leaseTransferCount tracks the number of lease transfers made.
 	leaseTransferCount int
-	// maxRangeMoveCount is the maximum number of range moves allowed.
+	// maxRangeMoveCount is the maximum number of range moves allowed in the
+	// context of the current rebalanceStores invocation (across multiple
+	// shedding stores).
 	maxRangeMoveCount int
-	// maxLeaseTransferCount is the maximum number of lease transfers allowed.
+	// maxLeaseTransferCount is the maximum number of lease transfers allowed in
+	// the context of the current rebalanceStores invocation. Note that because
+	// leases can only be shed from the particular local store on whose behalf
+	// rebalanceStores was called, this limit applies to this particular one
+	// store.
 	maxLeaseTransferCount int
 	// lastFailedChangeDelayDuration is the delay after a failed change before retrying.
 	lastFailedChangeDelayDuration time.Duration
-	// now is the timestamp when rebalancing started.
+	// now is the timestamp representing the start time of the current
+	// rebalanceStores invocation.
 	now time.Time
 	// Scratch variables reused across iterations.
 	scratch struct {
@@ -54,6 +68,43 @@ type rebalanceEnv struct {
 	}
 }
 
+func newRebalanceEnv(
+	cs *clusterState, rng *rand.Rand, dsm *diversityScoringMemo, now time.Time,
+) *rebalanceEnv {
+
+	// NB: these consts are intentionally local to the constructor, proving
+	// that they're not accessed anywhere else.
+	const (
+		// The current usage of the multi-metric allocator has rebalanceStores called
+		// in a loop while it emits work, and on a timer otherwise. This means that
+		// we don't need to emit many changes per invocation. Especially for range
+		// moves, which require moving significant amounts of data, emitting them one
+		// by one is fine and by doing so we operate on more recent information at each
+		// turn. Even if the caller carried out multiple changes concurrently, we'd
+		// want to be careful not to emit too many range moves at once, since they are
+		// expensive.
+		// Lease transfers are cheap and fast, so we emit a few more to avoid frequent
+		// trips through rebalanceStores which could cause some logging noise.
+		maxRangeMoveCount     = 1
+		maxLeaseTransferCount = 8
+		// See the long comment where rangeState.lastFailedChange is declared.
+		lastFailedChangeDelayDuration time.Duration = 60 * time.Second
+	)
+
+	re := &rebalanceEnv{
+		clusterState:                  cs,
+		rng:                           rng,
+		dsm:                           dsm,
+		now:                           now,
+		maxRangeMoveCount:             maxRangeMoveCount,
+		maxLeaseTransferCount:         maxLeaseTransferCount,
+		lastFailedChangeDelayDuration: lastFailedChangeDelayDuration,
+	}
+	re.scratch.nodes = map[roachpb.NodeID]*NodeLoad{}
+	re.scratch.stores = map[roachpb.StoreID]struct{}{}
+	return re
+}
+
 type sheddingStore struct {
 	roachpb.StoreID
 	storeLoadSummary
@@ -63,11 +114,11 @@ type sheddingStore struct {
 //
 // We do not want to shed replicas for CPU from a remote store until its had a
 // chance to shed leases.
-func (cs *clusterState) rebalanceStores(
-	ctx context.Context, localStoreID roachpb.StoreID, rng *rand.Rand, dsm *diversityScoringMemo,
+func (re *rebalanceEnv) rebalanceStores(
+	ctx context.Context, localStoreID roachpb.StoreID,
 ) []PendingRangeChange {
-	now := cs.ts.Now()
 	ctx = logtags.AddTag(ctx, "mmaid", mmaid.Add(1))
+
 	log.KvDistribution.VEventf(ctx, 2, "rebalanceStores begins")
 	// To select which stores are overloaded, we use a notion of overload that
 	// is based on cluster means (and of course individual store/node
@@ -87,7 +138,7 @@ func (cs *clusterState) rebalanceStores(
 	// responsible for equalizing load across two nodes that have 30% and 50%
 	// cpu utilization while the cluster mean is 70% utilization (as an
 	// example).
-	clusterMeans := cs.meansMemo.getMeans(nil)
+	clusterMeans := re.meansMemo.getMeans(nil)
 	var sheddingStores []sheddingStore
 	log.KvDistribution.Infof(ctx,
 		"cluster means: (stores-load %s) (stores-capacity %s) (nodes-cpu-load %d) (nodes-cpu-capacity %d)",
@@ -97,15 +148,15 @@ func (cs *clusterState) rebalanceStores(
 	// fdDrain or fdDead, nor do we attempt to shed replicas from a store which
 	// is storeMembershipRemoving (decommissioning). These are currently handled
 	// via replicate_queue.go.
-	for storeID, ss := range cs.stores {
-		sls := cs.meansMemo.getStoreLoadSummary(ctx, clusterMeans, storeID, ss.loadSeqNum)
+	for storeID, ss := range re.stores {
+		sls := re.meansMemo.getStoreLoadSummary(ctx, clusterMeans, storeID, ss.loadSeqNum)
 		log.KvDistribution.VEventf(ctx, 2, "evaluating s%d: node load %s, store load %s, worst dim %s",
 			storeID, sls.nls, sls.sls, sls.worstDim)
 
 		if sls.sls >= overloadSlow {
 			if ss.overloadEndTime != (time.Time{}) {
-				if now.Sub(ss.overloadEndTime) > overloadGracePeriod {
-					ss.overloadStartTime = now
+				if re.now.Sub(ss.overloadEndTime) > overloadGracePeriod {
+					ss.overloadStartTime = re.now
 					log.KvDistribution.Infof(ctx, "overload-start s%v (%v) - grace period expired", storeID, sls)
 				} else {
 					// Else, extend the previous overload interval.
@@ -128,7 +179,7 @@ func (cs *clusterState) rebalanceStores(
 			// NB: we don't stop the overloaded interval if the store is at
 			// loadNoChange, since a store can hover at the border of the two.
 			log.KvDistribution.Infof(ctx, "overload-end s%v (%v) - load dropped below no-change threshold", storeID, sls)
-			ss.overloadEndTime = now
+			ss.overloadEndTime = re.now
 		}
 	}
 
@@ -158,32 +209,6 @@ func (cs *clusterState) rebalanceStores(
 		}
 	}
 
-	// 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
-	// enacted, and its decision-making is no longer based on recent
-	// information. We don't have this issue with lease transfers since they are
-	// very fast, so we set a much higher limit.
-	//
-	// TODO: revisit these constants.
-	const maxRangeMoveCount = 1
-	const maxLeaseTransferCount = 8
-	// See the long comment where rangeState.lastFailedChange is declared.
-	const lastFailedChangeDelayDuration time.Duration = 60 * time.Second
-	re := &rebalanceEnv{
-		clusterState:                  cs,
-		rng:                           rng,
-		dsm:                           dsm,
-		changes:                       []PendingRangeChange{},
-		rangeMoveCount:                0,
-		leaseTransferCount:            0,
-		maxRangeMoveCount:             maxRangeMoveCount,
-		maxLeaseTransferCount:         maxLeaseTransferCount,
-		lastFailedChangeDelayDuration: lastFailedChangeDelayDuration,
-		now:                           now,
-	}
-	re.scratch.nodes = map[roachpb.NodeID]*NodeLoad{}
-	re.scratch.stores = map[roachpb.StoreID]struct{}{}
 	for _, store := range sheddingStores {
 		if re.rangeMoveCount >= re.maxRangeMoveCount || re.leaseTransferCount >= re.maxLeaseTransferCount {
 			break

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

@@ -523,7 +523,8 @@ func TestClusterState(t *testing.T) {
 					tr.SetRedactable(true)
 					defer tr.Close()
 					ctx, finishAndGet := tracing.ContextWithRecordingSpan(context.Background(), tr, "rebalance-stores")
-					cs.rebalanceStores(ctx, storeID, rng, dsm)
+					re := newRebalanceEnv(cs, rng, dsm, cs.ts.Now())
+					re.rebalanceStores(ctx, storeID)
 					rec := finishAndGet()
 					var sb redact.StringBuilder
 					rec.SafeFormatMinimal(&sb)