mmaprototype: introduce mmaload package

I want to avoid footguns related to negative load values. Having a package for
the load primitives allows hiding the naked slice, so that use from the
mmaprototype package can be enforced to be automatically correct.

This first commit creates the package and moves the LoadValue LoadVector types.
It also creates a SignledLoad{Value,Vector} type that will be adopted as
representation for the adjusted load (which can become negative).

Epic: CRDB-55052

Author: tbg

pkg/BUILD.bazel

@@ -1495,6 +1495,7 @@ GO_TARGETS = [
     "//pkg/kv/kvserver/allocator/allocatorimpl:allocatorimpl_test",
     "//pkg/kv/kvserver/allocator/load:load",
     "//pkg/kv/kvserver/allocator/load:load_test",
+    "//pkg/kv/kvserver/allocator/mmaprototype/mmaload:mmaload",
     "//pkg/kv/kvserver/allocator/mmaprototype:mmaprototype",
     "//pkg/kv/kvserver/allocator/mmaprototype:mmaprototype_test",
     "//pkg/kv/kvserver/allocator/plan:plan",

pkg/kv/kvserver/BUILD.bazel

@@ -132,6 +132,7 @@ go_library(
         "//pkg/kv/kvserver/allocator/allocatorimpl",
         "//pkg/kv/kvserver/allocator/load",
         "//pkg/kv/kvserver/allocator/mmaprototype",
+        "//pkg/kv/kvserver/allocator/mmaprototype/mmaload",
         "//pkg/kv/kvserver/allocator/plan",
         "//pkg/kv/kvserver/allocator/storepool",
         "//pkg/kv/kvserver/apply",

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

@@ -23,6 +23,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver/allocator/mmaprototype/mmaload",
         "//pkg/roachpb",
         "//pkg/util/log",
         "//pkg/util/metric",
@@ -48,6 +49,7 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":mmaprototype"],
     deps = [
+        "//pkg/kv/kvserver/allocator/mmaprototype/mmaload",
         "//pkg/roachpb",
         "//pkg/spanconfig/spanconfigtestutils",
         "//pkg/testutils/datapathutils",

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

@@ -16,6 +16,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype/mmaload"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/metric"
@@ -492,7 +493,7 @@ func sortTargetCandidateSetAndPick(
 	cands candidateSet,
 	loadThreshold loadSummary,
 	ignoreLevel ignoreLevel,
-	overloadedDim LoadDimension,
+	overloadedDim mmaload.LoadDimension,
 	rng *rand.Rand,
 ) roachpb.StoreID {
 	var b strings.Builder
@@ -570,7 +571,7 @@ func sortTargetCandidateSetAndPick(
 			if j == 0 {
 				// This is the lowestLoad set being considered now.
 				lowestLoadSet = cand.sls
-			} else if ignoreLevel < ignoreHigherThanLoadThreshold || overloadedDim == NumLoadDimensions {
+			} else if ignoreLevel < ignoreHigherThanLoadThreshold || overloadedDim == mmaload.NumLoadDimensions {
 				// Past the lowestLoad set. We don't care about these.
 				break
 			}
@@ -583,7 +584,7 @@ func sortTargetCandidateSetAndPick(
 		}
 		candDiscardedByNLS := cand.nls > loadThreshold ||
 			(cand.nls == loadThreshold && ignoreLevel < ignoreHigherThanLoadThreshold)
-		candDiscardedByOverloadDim := overloadedDim != NumLoadDimensions &&
+		candDiscardedByOverloadDim := overloadedDim != mmaload.NumLoadDimensions &&
 			cand.dimSummary[overloadedDim] >= loadNoChange
 		if candDiscardedByNLS || candDiscardedByOverloadDim ||
 			cand.maxFractionPendingIncrease >= maxFractionPendingThreshold {
@@ -664,7 +665,7 @@ func sortTargetCandidateSetAndPick(
 		return 0
 	}
 	cands.candidates = cands.candidates[:j]
-	if lowestLoadSet != highestLoadSet || (lowestLoadSet >= loadNoChange && overloadedDim != NumLoadDimensions) {
+	if lowestLoadSet != highestLoadSet || (lowestLoadSet >= loadNoChange && overloadedDim != mmaload.NumLoadDimensions) {
 		// Sort candidates from lowest to highest along overloaded dimension. We
 		// limit when we do this, since this will further restrict the pool of
 		// candidates and in general we don't want to restrict the pool.

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

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype/mmaload"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -173,7 +174,7 @@ func (s ReplicaChangeType) String() string {
 type ReplicaChange struct {
 	// The load this change adds to a store. The values will be negative if the
 	// load is being removed.
-	loadDelta          LoadVector
+	loadDelta          mmaload.LoadVector
 	secondaryLoadDelta SecondaryLoadVector
 
 	// target is the target {store,node} for the change.
@@ -333,9 +334,9 @@ func MakeLeaseTransferChanges(
 	// 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
 	// lease.
-	nonRaftCPU := rLoad.Load[CPURate] - rLoad.RaftCPU
-	removeLease.loadDelta[CPURate] = -nonRaftCPU
-	addLease.loadDelta[CPURate] = loadToAdd(nonRaftCPU)
+	nonRaftCPU := rLoad.Load[mmaload.CPURate] - rLoad.RaftCPU
+	removeLease.loadDelta[mmaload.CPURate] = -nonRaftCPU
+	addLease.loadDelta[mmaload.CPURate] = loadToAdd(nonRaftCPU)
 	return [2]ReplicaChange{removeLease, addLease}
 }
 
@@ -362,13 +363,13 @@ func MakeAddReplicaChange(
 		replicaChangeType: AddReplica,
 	}
 	addReplica.next.ReplicaID = unknownReplicaID
-	addReplica.loadDelta.add(loadVectorToAdd(rLoad.Load))
+	addReplica.loadDelta.Add(loadVectorToAdd(rLoad.Load))
 	if replicaIDAndType.IsLeaseholder {
 		addReplica.secondaryLoadDelta[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.
-		addReplica.loadDelta[CPURate] = loadToAdd(rLoad.RaftCPU)
+		addReplica.loadDelta[mmaload.CPURate] = loadToAdd(rLoad.RaftCPU)
 	}
 	return addReplica
 }
@@ -390,13 +391,13 @@ func MakeRemoveReplicaChange(
 		},
 		replicaChangeType: RemoveReplica,
 	}
-	removeReplica.loadDelta.subtract(rLoad.Load)
+	removeReplica.loadDelta.Subtract(rLoad.Load)
 	if replicaState.IsLeaseholder {
 		removeReplica.secondaryLoadDelta[LeaseCount] = -1
 	} else {
 		// Set the load delta for CPU to be just the raft CPU. The non-raft CPU is
 		// associated with the lease.
-		removeReplica.loadDelta[CPURate] = -rLoad.RaftCPU
+		removeReplica.loadDelta[mmaload.CPURate] = -rLoad.RaftCPU
 	}
 	return removeReplica
 }
@@ -421,10 +422,10 @@ func MakeReplicaTypeChange(
 	}
 	if next.IsLeaseholder {
 		change.secondaryLoadDelta[LeaseCount] = 1
-		change.loadDelta[CPURate] = loadToAdd(rLoad.Load[CPURate] - rLoad.RaftCPU)
+		change.loadDelta[mmaload.CPURate] = loadToAdd(rLoad.Load[mmaload.CPURate] - rLoad.RaftCPU)
 	} else if prev.IsLeaseholder {
 		change.secondaryLoadDelta[LeaseCount] = -1
-		change.loadDelta[CPURate] = rLoad.RaftCPU - rLoad.Load[CPURate]
+		change.loadDelta[mmaload.CPURate] = rLoad.RaftCPU - rLoad.Load[mmaload.CPURate]
 	}
 	return change
 }
@@ -775,7 +776,7 @@ type storeState struct {
 		// NB: these load values can become negative due to applying pending
 		// changes. We need to let them be negative to retain the ability to undo
 		// pending changes.
-		load          LoadVector
+		load          mmaload.LoadVector
 		secondaryLoad SecondaryLoadVector
 		// Pending changes for computing loadReplicas and load.
 		// Added to at the same time as clusterState.pendingChanges.
@@ -977,7 +978,7 @@ type nodeState struct {
 	stores []roachpb.StoreID
 	NodeLoad
 	// NB: adjustedCPU can be negative.
-	adjustedCPU LoadValue
+	adjustedCPU mmaload.LoadValue
 }
 
 func newNodeState(nodeID roachpb.NodeID) *nodeState {
@@ -1400,11 +1401,11 @@ func (cs *clusterState) processStoreLoadMsg(ctx context.Context, storeMsg *Store
 	if ss == nil {
 		panic(fmt.Sprintf("store %d not found", storeMsg.StoreID))
 	}
-	ns.ReportedCPU += storeMsg.Load[CPURate] - ss.reportedLoad[CPURate]
-	ns.CapacityCPU += storeMsg.Capacity[CPURate] - ss.capacity[CPURate]
+	ns.ReportedCPU += storeMsg.Load[mmaload.CPURate] - ss.reportedLoad[mmaload.CPURate]
+	ns.CapacityCPU += storeMsg.Capacity[mmaload.CPURate] - ss.capacity[mmaload.CPURate]
 	// Undo the adjustment for the store. We will apply the adjustment again
 	// below.
-	ns.adjustedCPU += storeMsg.Load[CPURate] - ss.adjusted.load[CPURate]
+	ns.adjustedCPU += storeMsg.Load[mmaload.CPURate] - ss.adjusted.load[mmaload.CPURate]
 
 	// The store's load sequence number is incremented on each load change. The
 	// store's load is updated below.
@@ -1702,18 +1703,18 @@ func (cs *clusterState) processStoreLeaseholderMsgInternal(
 		topk.startInit()
 		sls := cs.computeLoadSummary(ctx, ss.StoreID, &clusterMeans.storeLoad, &clusterMeans.nodeLoad)
 		if ss.StoreID == localss.StoreID {
-			topk.dim = CPURate
+			topk.dim = mmaload.CPURate
 		} else {
-			topk.dim = WriteBandwidth
+			topk.dim = mmaload.WriteBandwidth
 		}
 		if sls.highDiskSpaceUtilization {
-			topk.dim = ByteSize
+			topk.dim = mmaload.ByteSize
 		} else if sls.sls > loadNoChange {
 			// If multiple dimensions are contributing the same loadSummary, we will pick
 			// CPURate before WriteBandwidth before ByteSize.
 			for i := range sls.dimSummary {
 				if sls.dimSummary[i] == sls.sls {
-					topk.dim = LoadDimension(i)
+					topk.dim = mmaload.LoadDimension(i)
 					break
 				}
 			}
@@ -1775,7 +1776,7 @@ func (cs *clusterState) processStoreLeaseholderMsgInternal(
 			minReplicaLoadFraction = 0.02
 		)
 		fraction := minReplicaLoadFraction
-		if ss.StoreID == msg.StoreID && topk.dim == CPURate {
+		if ss.StoreID == msg.StoreID && topk.dim == mmaload.CPURate {
 			// We are assuming we will be able to shed leases, but if we can't we
 			// will start shedding replicas, so this is just a heuristic.
 			fraction = minLeaseLoadFraction
@@ -1784,7 +1785,7 @@ func (cs *clusterState) processStoreLeaseholderMsgInternal(
 		// Given that this is a most overloaded dim, the likelihood of the
 		// adjusted load being negative is very low.
 		adjustedStoreLoadValue := max(0, ss.adjusted.load[topk.dim])
-		threshold := LoadValue(float64(adjustedStoreLoadValue) * fraction)
+		threshold := mmaload.LoadValue(float64(adjustedStoreLoadValue) * fraction)
 		if ss.reportedSecondaryLoad[ReplicaCount] > 0 {
 			// Allow all ranges above 90% of the mean. This is quite arbitrary.
 			meanLoad := (adjustedStoreLoadValue * 9) / (ss.reportedSecondaryLoad[ReplicaCount] * 10)
@@ -1818,16 +1819,16 @@ func (cs *clusterState) processStoreLeaseholderMsgInternal(
 				ss := cs.stores[replica.StoreID]
 				topk := ss.adjusted.topKRanges[msg.StoreID]
 				switch topk.dim {
-				case CPURate:
-					l := rs.load.Load[CPURate]
+				case mmaload.CPURate:
+					l := rs.load.Load[mmaload.CPURate]
 					if !replica.ReplicaState.IsLeaseholder {
 						l = rs.load.RaftCPU
 					}
 					topk.addReplica(ctx, rangeID, l, replica.StoreID, msg.StoreID)
-				case WriteBandwidth:
-					topk.addReplica(ctx, rangeID, rs.load.Load[WriteBandwidth], replica.StoreID, msg.StoreID)
-				case ByteSize:
-					topk.addReplica(ctx, rangeID, rs.load.Load[ByteSize], replica.StoreID, msg.StoreID)
+				case mmaload.WriteBandwidth:
+					topk.addReplica(ctx, rangeID, rs.load.Load[mmaload.WriteBandwidth], replica.StoreID, msg.StoreID)
+				case mmaload.ByteSize:
+					topk.addReplica(ctx, rangeID, rs.load.Load[mmaload.ByteSize], replica.StoreID, msg.StoreID)
 				}
 			}
 		}
@@ -2217,22 +2218,22 @@ func (cs *clusterState) undoReplicaChange(change ReplicaChange) {
 // store and node affected.
 func (cs *clusterState) applyChangeLoadDelta(change ReplicaChange) {
 	ss := cs.stores[change.target.StoreID]
-	ss.adjusted.load.add(change.loadDelta)
+	ss.adjusted.load.Add(change.loadDelta)
 	ss.adjusted.secondaryLoad.add(change.secondaryLoadDelta)
 	ss.loadSeqNum++
 	ss.computeMaxFractionPending()
-	cs.nodes[ss.NodeID].adjustedCPU += change.loadDelta[CPURate]
+	cs.nodes[ss.NodeID].adjustedCPU += change.loadDelta[mmaload.CPURate]
 }
 
 // undoChangeLoadDelta subtracts the change load delta from the adjusted load
 // of the store and node affected.
 func (cs *clusterState) undoChangeLoadDelta(change ReplicaChange) {
 	ss := cs.stores[change.target.StoreID]
-	ss.adjusted.load.subtract(change.loadDelta)
+	ss.adjusted.load.Subtract(change.loadDelta)
 	ss.adjusted.secondaryLoad.subtract(change.secondaryLoadDelta)
 	ss.loadSeqNum++
 	ss.computeMaxFractionPending()
-	cs.nodes[ss.NodeID].adjustedCPU -= change.loadDelta[CPURate]
+	cs.nodes[ss.NodeID].adjustedCPU -= change.loadDelta[mmaload.CPURate]
 }
 
 // setStore updates the store attributes and locality in the cluster state. If
@@ -2303,12 +2304,12 @@ func (cs *clusterState) canShedAndAddLoad(
 	ctx context.Context,
 	srcSS *storeState,
 	targetSS *storeState,
-	delta LoadVector,
+	delta mmaload.LoadVector,
 	means *meansLoad,
 	onlyConsiderTargetCPUSummary bool,
-	overloadedDim LoadDimension,
+	overloadedDim mmaload.LoadDimension,
 ) (canAddLoad bool) {
-	if overloadedDim == NumLoadDimensions {
+	if overloadedDim == mmaload.NumLoadDimensions {
 		panic("overloadedDim must not be NumLoadDimensions")
 	}
 	// TODO(tbg): in experiments, we often see interesting behavior right when
@@ -2318,26 +2319,26 @@ func (cs *clusterState) canShedAndAddLoad(
 	targetNS := cs.nodes[targetSS.NodeID]
 	// Add the delta.
 	deltaToAdd := loadVectorToAdd(delta)
-	targetSS.adjusted.load.add(deltaToAdd)
+	targetSS.adjusted.load.Add(deltaToAdd)
 	// TODO(tbg): why does NodeLoad have an adjustedCPU field but not fields for
 	// the other load dimensions? We just added deltaToAdd to targetSS.adjusted,
 	// shouldn't this be wholly reflected in targetNS as well, not just for CPU?
 	// Or maybe CPU is the only dimension that matters at the node level. It feels
 	// sloppy/confusing though.
-	targetNS.adjustedCPU += deltaToAdd[CPURate]
+	targetNS.adjustedCPU += deltaToAdd[mmaload.CPURate]
 	targetSLS := computeLoadSummary(ctx, targetSS, targetNS, &means.storeLoad, &means.nodeLoad)
 	// Undo the addition.
-	targetSS.adjusted.load.subtract(deltaToAdd)
-	targetNS.adjustedCPU -= deltaToAdd[CPURate]
+	targetSS.adjusted.load.Subtract(deltaToAdd)
+	targetNS.adjustedCPU -= deltaToAdd[mmaload.CPURate]
 
 	// Remove the delta.
 	srcNS := cs.nodes[srcSS.NodeID]
-	srcSS.adjusted.load.subtract(delta)
-	srcNS.adjustedCPU -= delta[CPURate]
+	srcSS.adjusted.load.Subtract(delta)
+	srcNS.adjustedCPU -= delta[mmaload.CPURate]
 	srcSLS := computeLoadSummary(ctx, srcSS, srcNS, &means.storeLoad, &means.nodeLoad)
 	// Undo the removal.
-	srcSS.adjusted.load.add(delta)
-	srcNS.adjustedCPU += delta[CPURate]
+	srcSS.adjusted.load.Add(delta)
+	srcNS.adjustedCPU += delta[mmaload.CPURate]
 
 	var reason strings.Builder
 	defer func() {
@@ -2365,7 +2366,7 @@ func (cs *clusterState) canShedAndAddLoad(
 	// loadNoChange, and the source must have been > loadNoChange.
 	var targetSummary loadSummary
 	if onlyConsiderTargetCPUSummary {
-		targetSummary = targetSLS.dimSummary[CPURate]
+		targetSummary = targetSLS.dimSummary[mmaload.CPURate]
 		if targetSummary < targetSLS.nls {
 			targetSummary = targetSLS.nls
 		}
@@ -2423,7 +2424,7 @@ func (cs *clusterState) canShedAndAddLoad(
 	}
 	otherDimensionsBecameWorseInTarget := false
 	for i := range targetSLS.dimSummary {
-		dim := LoadDimension(i)
+		dim := mmaload.LoadDimension(i)
 		if dim == overloadedDim {
 			continue
 		}
@@ -2530,24 +2531,24 @@ func computeLoadSummary(
 ) storeLoadSummary {
 	sls := loadLow
 	var highDiskSpaceUtil bool
-	var dimSummary [NumLoadDimensions]loadSummary
-	var worstDim LoadDimension
+	var dimSummary [mmaload.NumLoadDimensions]loadSummary
+	var worstDim mmaload.LoadDimension
 	for i := range msl.load {
 		const nodeIDForLogging = 0
-		ls := loadSummaryForDimension(ctx, ss.StoreID, nodeIDForLogging, LoadDimension(i), ss.adjusted.load[i], ss.capacity[i],
+		ls := loadSummaryForDimension(ctx, ss.StoreID, nodeIDForLogging, mmaload.LoadDimension(i), ss.adjusted.load[i], ss.capacity[i],
 			msl.load[i], msl.util[i])
 		if ls > sls {
 			sls = ls
-			worstDim = LoadDimension(i)
+			worstDim = mmaload.LoadDimension(i)
 		}
 		dimSummary[i] = ls
-		switch LoadDimension(i) {
-		case ByteSize:
+		switch mmaload.LoadDimension(i) {
+		case mmaload.ByteSize:
 			highDiskSpaceUtil = highDiskSpaceUtilization(ss.adjusted.load[i], ss.capacity[i])
 		}
 	}
 	const storeIDForLogging = 0
-	nls := loadSummaryForDimension(ctx, storeIDForLogging, ns.NodeID, CPURate, ns.adjustedCPU, ns.CapacityCPU, mnl.loadCPU, mnl.utilCPU)
+	nls := loadSummaryForDimension(ctx, storeIDForLogging, ns.NodeID, mmaload.CPURate, ns.adjustedCPU, ns.CapacityCPU, mnl.loadCPU, mnl.utilCPU)
 	return storeLoadSummary{
 		worstDim:                   worstDim,
 		sls:                        sls,