mmaprototype: move secondary load to mmaload

Author: tbg

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

@@ -175,7 +175,7 @@ type ReplicaChange struct {
 	// The load this change adds to a store. The values will be negative if the
 	// load is being removed.
 	loadDelta          mmaload.LoadVector
-	secondaryLoadDelta SecondaryLoadVector
+	secondaryLoadDelta mmaload.SecondaryLoadVector
 
 	// target is the target {store,node} for the change.
 	target roachpb.ReplicationTarget
@@ -328,8 +328,8 @@ func MakeLeaseTransferChanges(
 	}
 	removeLease.next.IsLeaseholder = false
 	addLease.next.IsLeaseholder = true
-	removeLease.secondaryLoadDelta[LeaseCount] = -1
-	addLease.secondaryLoadDelta[LeaseCount] = 1
+	removeLease.secondaryLoadDelta[mmaload.LeaseCount] = -1
+	addLease.secondaryLoadDelta[mmaload.LeaseCount] = 1
 
 	// Only account for the leaseholder CPU, all other primary load dimensions
 	// are ignored. Byte size and write bytes are not impacted by having a range
@@ -365,7 +365,7 @@ func MakeAddReplicaChange(
 	addReplica.next.ReplicaID = unknownReplicaID
 	addReplica.loadDelta.Add(loadVectorToAdd(rLoad.Load))
 	if replicaIDAndType.IsLeaseholder {
-		addReplica.secondaryLoadDelta[LeaseCount] = 1
+		addReplica.secondaryLoadDelta[mmaload.LeaseCount] = 1
 	} else {
 		// Set the load delta for CPU to be just the raft CPU. The non-raft CPU we
 		// assume is associated with the lease.
@@ -393,7 +393,7 @@ func MakeRemoveReplicaChange(
 	}
 	removeReplica.loadDelta.Subtract(rLoad.Load)
 	if replicaState.IsLeaseholder {
-		removeReplica.secondaryLoadDelta[LeaseCount] = -1
+		removeReplica.secondaryLoadDelta[mmaload.LeaseCount] = -1
 	} else {
 		// Set the load delta for CPU to be just the raft CPU. The non-raft CPU is
 		// associated with the lease.
@@ -421,10 +421,10 @@ func MakeReplicaTypeChange(
 		replicaChangeType: ChangeReplica,
 	}
 	if next.IsLeaseholder {
-		change.secondaryLoadDelta[LeaseCount] = 1
+		change.secondaryLoadDelta[mmaload.LeaseCount] = 1
 		change.loadDelta[mmaload.CPURate] = loadToAdd(rLoad.Load[mmaload.CPURate] - rLoad.RaftCPU)
 	} else if prev.IsLeaseholder {
-		change.secondaryLoadDelta[LeaseCount] = -1
+		change.secondaryLoadDelta[mmaload.LeaseCount] = -1
 		change.loadDelta[mmaload.CPURate] = rLoad.RaftCPU - rLoad.Load[mmaload.CPURate]
 	}
 	return change
@@ -1786,9 +1786,9 @@ func (cs *clusterState) processStoreLeaseholderMsgInternal(
 		// adjusted load being negative is very low.
 		adjustedStoreLoadValue := max(0, ss.adjusted.load[topk.dim])
 		threshold := mmaload.LoadValue(float64(adjustedStoreLoadValue) * fraction)
-		if ss.reportedSecondaryLoad[ReplicaCount] > 0 {
+		if ss.reportedSecondaryLoad[mmaload.ReplicaCount] > 0 {
 			// Allow all ranges above 90% of the mean. This is quite arbitrary.
-			meanLoad := (adjustedStoreLoadValue * 9) / (ss.reportedSecondaryLoad[ReplicaCount] * 10)
+			meanLoad := (adjustedStoreLoadValue * 9) / (ss.reportedSecondaryLoad[mmaload.ReplicaCount] * 10)
 			threshold = min(meanLoad, threshold)
 		}
 		topk.threshold = threshold
@@ -2219,7 +2219,7 @@ func (cs *clusterState) undoReplicaChange(change ReplicaChange) {
 func (cs *clusterState) applyChangeLoadDelta(change ReplicaChange) {
 	ss := cs.stores[change.target.StoreID]
 	ss.adjusted.load.Add(change.loadDelta)
-	ss.adjusted.secondaryLoad.add(change.secondaryLoadDelta)
+	ss.adjusted.secondaryLoad.Add(change.secondaryLoadDelta)
 	ss.loadSeqNum++
 	ss.computeMaxFractionPending()
 	cs.nodes[ss.NodeID].adjustedCPU += change.loadDelta[mmaload.CPURate]
@@ -2230,7 +2230,7 @@ func (cs *clusterState) applyChangeLoadDelta(change ReplicaChange) {
 func (cs *clusterState) undoChangeLoadDelta(change ReplicaChange) {
 	ss := cs.stores[change.target.StoreID]
 	ss.adjusted.load.Subtract(change.loadDelta)
-	ss.adjusted.secondaryLoad.subtract(change.secondaryLoadDelta)
+	ss.adjusted.secondaryLoad.Subtract(change.secondaryLoadDelta)
 	ss.loadSeqNum++
 	ss.computeMaxFractionPending()
 	cs.nodes[ss.NodeID].adjustedCPU -= change.loadDelta[mmaload.CPURate]

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

@@ -58,10 +58,10 @@ func parseLoadVector(t *testing.T, in string) mmaload.LoadVector {
 	return vec
 }
 
-func parseSecondaryLoadVector(t *testing.T, in string) SecondaryLoadVector {
-	var vec SecondaryLoadVector
+func parseSecondaryLoadVector(t *testing.T, in string) mmaload.SecondaryLoadVector {
+	var vec mmaload.SecondaryLoadVector
 	parts := strings.Split(stripBrackets(t, in), ",")
-	require.LessOrEqual(t, len(parts), int(NumSecondaryLoadDimensions))
+	require.LessOrEqual(t, len(parts), int(mmaload.NumSecondaryLoadDimensions))
 	for dim := range parts {
 		vec[dim] = mmaload.LoadValue(parseInt(t, parts[dim]))
 	}

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

@@ -26,57 +26,6 @@ const (
 	UnknownCapacity mmaload.LoadValue = math.MaxInt64
 )
 
-// SecondaryLoadDimension represents secondary load dimensions that should be
-// considered after we are done rebalancing using loadDimensions, since these
-// don't represent "real" resources. Currently, only lease and replica counts
-// are considered here. Lease rebalancing will see if there is scope to move
-// some leases between stores that do not have any pending changes and are not
-// overloaded (and will not get overloaded by the movement). This will happen
-// in a separate pass (i.e., not in clusterState.rebalanceStores) -- the
-// current plan is to continue using the leaseQueue and call from it into MMA.
-//
-// Note that lease rebalancing will only move leases and not replicas. Also,
-// the rebalancing will take into account the lease preferences, as discussed
-// in https://github.com/cockroachdb/cockroach/issues/93258, and the lease
-// counts among the current candidates (see
-// https://github.com/cockroachdb/cockroach/pull/98893).
-//
-// To use MMA for replica count rebalancing, done by the replicateQueue, we
-// also have a ReplicaCount load dimension.
-//
-// These are currently unused, since the initial integration of MMA is to
-// replace load-based rebalancing performed by the StoreRebalancer.
-type SecondaryLoadDimension uint8
-
-const (
-	LeaseCount SecondaryLoadDimension = iota
-	ReplicaCount
-	NumSecondaryLoadDimensions
-)
-
-type SecondaryLoadVector [NumSecondaryLoadDimensions]mmaload.LoadValue
-
-func (lv *SecondaryLoadVector) add(other SecondaryLoadVector) {
-	for i := range other {
-		(*lv)[i] += other[i]
-	}
-}
-
-func (lv *SecondaryLoadVector) subtract(other SecondaryLoadVector) {
-	for i := range other {
-		(*lv)[i] -= other[i]
-	}
-}
-
-func (lv SecondaryLoadVector) String() string {
-	return redact.StringWithoutMarkers(lv)
-}
-
-// SafeFormat implements the redact.SafeFormatter interface.
-func (lv SecondaryLoadVector) SafeFormat(w redact.SafePrinter, _ rune) {
-	w.Printf("[lease:%d, replica:%d]", lv[LeaseCount], lv[ReplicaCount])
-}
-
 type RangeLoad struct {
 	Load mmaload.LoadVector
 	// Nanos per second. RaftCPU <= Load[cpu]. Handling this as a special case,
@@ -111,7 +60,7 @@ type storeLoad struct {
 	// provisioned disk bandwidth in the near future.
 	capacity mmaload.LoadVector
 
-	reportedSecondaryLoad SecondaryLoadVector
+	reportedSecondaryLoad mmaload.SecondaryLoadVector
 }
 
 // NodeLoad is the load information for a node.
@@ -130,7 +79,7 @@ type meanStoreLoad struct {
 	// Util is 0 for CPURate, WriteBandwidth. Non-zero for ByteSize.
 	util [mmaload.NumLoadDimensions]float64
 
-	secondaryLoad SecondaryLoadVector
+	secondaryLoad mmaload.SecondaryLoadVector
 }
 
 // The mean node load for a set of NodeLoad.

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

@@ -89,7 +89,7 @@ func TestMeansMemo(t *testing.T) {
 					reportedLoad: mmaload.LoadVector{mmaload.LoadValue(cpuLoad), mmaload.LoadValue(wbLoad), mmaload.LoadValue(bsLoad)},
 					capacity: mmaload.LoadVector{
 						mmaload.LoadValue(cpuCapacity), mmaload.LoadValue(wbCapacity), mmaload.LoadValue(bsCapacity)},
-					reportedSecondaryLoad: SecondaryLoadVector{leaseCountLoad},
+					reportedSecondaryLoad: mmaload.SecondaryLoadVector{leaseCountLoad},
 				}
 				for i := range sLoad.capacity {
 					if sLoad.capacity[i] < 0 {

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

@@ -25,7 +25,7 @@ type StoreLoadMsg struct {
 	// derive a node level capacity, and then dividing that by the number of
 	// stores.
 	Capacity      mmaload.LoadVector
-	SecondaryLoad SecondaryLoadVector
+	SecondaryLoad mmaload.SecondaryLoadVector
 
 	LoadTime time.Time
 }

pkg/kv/kvserver/allocator/mmaprototype/mmaload/BUILD.bazel

@@ -3,6 +3,7 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library")
 go_library(
     name = "mmaload",
     srcs = [
+        "secondary.go",
         "signed.go",
         "unsigned.go",
     ],

pkg/kv/kvserver/allocator/mmaprototype/mmaload/secondary.go

@@ -0,0 +1,59 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package mmaload
+
+import "github.com/cockroachdb/redact"
+
+// SecondaryLoadDimension represents secondary load dimensions that should be
+// considered after we are done rebalancing using loadDimensions, since these
+// don't represent "real" resources. Currently, only lease and replica counts
+// are considered here. Lease rebalancing will see if there is scope to move
+// some leases between stores that do not have any pending changes and are not
+// overloaded (and will not get overloaded by the movement). This will happen
+// in a separate pass (i.e., not in clusterState.rebalanceStores) -- the
+// current plan is to continue using the leaseQueue and call from it into MMA.
+//
+// Note that lease rebalancing will only move leases and not replicas. Also,
+// the rebalancing will take into account the lease preferences, as discussed
+// in https://github.com/cockroachdb/cockroach/issues/93258, and the lease
+// counts among the current candidates (see
+// https://github.com/cockroachdb/cockroach/pull/98893).
+//
+// To use MMA for replica count rebalancing, done by the replicateQueue, we
+// also have a ReplicaCount load dimension.
+//
+// These are currently unused, since the initial integration of MMA is to
+// replace load-based rebalancing performed by the StoreRebalancer.
+type SecondaryLoadDimension uint8
+
+const (
+	LeaseCount SecondaryLoadDimension = iota
+	ReplicaCount
+	NumSecondaryLoadDimensions
+)
+
+type SecondaryLoadVector [NumSecondaryLoadDimensions]LoadValue
+
+func (lv *SecondaryLoadVector) Add(other SecondaryLoadVector) {
+	for i := range other {
+		(*lv)[i] += other[i]
+	}
+}
+
+func (lv *SecondaryLoadVector) Subtract(other SecondaryLoadVector) {
+	for i := range other {
+		(*lv)[i] -= other[i]
+	}
+}
+
+func (lv SecondaryLoadVector) String() string {
+	return redact.StringWithoutMarkers(lv)
+}
+
+// SafeFormat implements the redact.SafeFormatter interface.
+func (lv SecondaryLoadVector) SafeFormat(w redact.SafePrinter, _ rune) {
+	w.Printf("[lease:%d, replica:%d]", lv[LeaseCount], lv[ReplicaCount])
+}

pkg/kv/kvserver/mmaintegration/store_load_msg.go

@@ -118,9 +118,9 @@ func MakeStoreLoadMsg(
 		// which is desirably pessimistic.
 		capacity[mmaload.ByteSize] = mmaload.LoadValue(desc.Capacity.Available)
 	}
-	var secondaryLoad mmaprototype.SecondaryLoadVector
-	secondaryLoad[mmaprototype.LeaseCount] = mmaload.LoadValue(desc.Capacity.LeaseCount)
-	secondaryLoad[mmaprototype.ReplicaCount] = mmaload.LoadValue(desc.Capacity.RangeCount)
+	var secondaryLoad mmaload.SecondaryLoadVector
+	secondaryLoad[mmaload.LeaseCount] = mmaload.LoadValue(desc.Capacity.LeaseCount)
+	secondaryLoad[mmaload.ReplicaCount] = mmaload.LoadValue(desc.Capacity.RangeCount)
 	// TODO(tbg): this triggers early in tests, probably we're making load messages
 	// before having received the first capacity. Still, this is bad, should fix.
 	// or handle properly by communicating an unknown capacity.