From 01e7511ce756229d4c4f2ffde3ee0ba41cb5edba Mon Sep 17 00:00:00 2001 From: wenyihu6 Date: Thu, 20 Nov 2025 20:44:46 -0500 Subject: [PATCH 1/5] mmaprototype: extract analyzeFunc This commit moves analyzeFunc out of finishInit into its own helper improve readability. --- .../allocator/mmaprototype/constraint.go | 167 +++++++++--------- 1 file changed, 85 insertions(+), 82 deletions(-) diff --git a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go index e879ff4834c0..377c23bfd803 100644 --- a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go +++ b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go @@ -830,80 +830,65 @@ type storeMatchesConstraintInterface interface { storeMatches(storeID roachpb.StoreID, constraintConj constraintsConj) bool } -func (rac *rangeAnalyzedConstraints) finishInit( - spanConfig *normalizedSpanConfig, - constraintMatcher storeMatchesConstraintInterface, - leaseholder roachpb.StoreID, +func (ac *analyzedConstraints) analyzeFunc( + buf *analyzeConstraintsBuf, constraintMatcher storeMatchesConstraintInterface, ) { - rac.spanConfig = spanConfig - rac.numNeededReplicas[voterIndex] = spanConfig.numVoters - rac.numNeededReplicas[nonVoterIndex] = spanConfig.numReplicas - spanConfig.numVoters - rac.replicas = rac.buf.replicas - - analyzeFunc := func(ac *analyzedConstraints) { - if len(ac.constraints) == 0 { - // Nothing to do. - return - } - for i := 0; i < len(ac.constraints); i++ { - ac.satisfiedByReplica[voterIndex] = extend2DSlice(ac.satisfiedByReplica[voterIndex]) - ac.satisfiedByReplica[nonVoterIndex] = extend2DSlice(ac.satisfiedByReplica[nonVoterIndex]) - } - // Compute the list of all constraints satisfied by each store. - for kind := voterIndex; kind < numReplicaKinds; kind++ { - for i, store := range rac.buf.replicas[kind] { - rac.buf.replicaConstraintIndices[kind] = - extend2DSlice(rac.buf.replicaConstraintIndices[kind]) - for j, c := range ac.constraints { - if len(c.constraints) == 0 || constraintMatcher.storeMatches(store.StoreID, c.constraints) { - rac.buf.replicaConstraintIndices[kind][i] = - append(rac.buf.replicaConstraintIndices[kind][i], int32(j)) - } - } - n := len(rac.buf.replicaConstraintIndices[kind][i]) - if n == 0 { - ac.satisfiedNoConstraintReplica[kind] = - append(ac.satisfiedNoConstraintReplica[kind], store.StoreID) - } else if n == 1 { - // Satisfies exactly one constraint, so place it there. - constraintIndex := rac.buf.replicaConstraintIndices[kind][i][0] - ac.satisfiedByReplica[kind][constraintIndex] = - append(ac.satisfiedByReplica[kind][constraintIndex], store.StoreID) - rac.buf.replicaConstraintIndices[kind][i] = rac.buf.replicaConstraintIndices[kind][i][:0] + if len(ac.constraints) == 0 { + // Nothing to do. + return + } + for i := 0; i < len(ac.constraints); i++ { + ac.satisfiedByReplica[voterIndex] = extend2DSlice(ac.satisfiedByReplica[voterIndex]) + ac.satisfiedByReplica[nonVoterIndex] = extend2DSlice(ac.satisfiedByReplica[nonVoterIndex]) + } + // Compute the list of all constraints satisfied by each store. + for kind := voterIndex; kind < numReplicaKinds; kind++ { + for i, store := range buf.replicas[kind] { + buf.replicaConstraintIndices[kind] = + extend2DSlice(buf.replicaConstraintIndices[kind]) + for j, c := range ac.constraints { + if len(c.constraints) == 0 || constraintMatcher.storeMatches(store.StoreID, c.constraints) { + buf.replicaConstraintIndices[kind][i] = + append(buf.replicaConstraintIndices[kind][i], int32(j)) } - // Else, satisfied multiple constraints. Don't choose yet. - } - } - // The only stores not yet in ac are the ones that satisfy multiple - // constraints. For each store, the constraint indices it satisfies are in - // increasing order. Satisfy constraints in order, while not - // oversatisfying. - for j := range ac.constraints { - doneFunc := func() bool { - return len(ac.satisfiedByReplica[voterIndex][j])+ - len(ac.satisfiedByReplica[nonVoterIndex][j]) >= int(ac.constraints[j].numReplicas) } - done := doneFunc() - if done { - continue + n := len(buf.replicaConstraintIndices[kind][i]) + if n == 0 { + ac.satisfiedNoConstraintReplica[kind] = + append(ac.satisfiedNoConstraintReplica[kind], store.StoreID) + } else if n == 1 { + // Satisfies exactly one constraint, so place it there. + constraintIndex := buf.replicaConstraintIndices[kind][i][0] + ac.satisfiedByReplica[kind][constraintIndex] = + append(ac.satisfiedByReplica[kind][constraintIndex], store.StoreID) + buf.replicaConstraintIndices[kind][i] = buf.replicaConstraintIndices[kind][i][:0] } - for kind := voterIndex; kind < numReplicaKinds; kind++ { - for i := range rac.buf.replicaConstraintIndices[kind] { - constraintIndices := rac.buf.replicaConstraintIndices[kind][i] - for _, index := range constraintIndices { - if index == int32(j) { - ac.satisfiedByReplica[kind][j] = - append(ac.satisfiedByReplica[kind][j], rac.replicas[kind][i].StoreID) - rac.buf.replicaConstraintIndices[kind][i] = constraintIndices[:0] - done = doneFunc() - // This store is finished. - break - } - } - // done can be true if some store was appended to - // ac.satisfiedByReplica[kind][j] and made it fully satisfied. Don't - // need to look at other stores for this constraint. - if done { + // Else, satisfied multiple constraints. Don't choose yet. + } + } + // The only stores not yet in ac are the ones that satisfy multiple + // constraints. For each store, the constraint indices it satisfies are in + // increasing order. Satisfy constraints in order, while not + // oversatisfying. + for j := range ac.constraints { + doneFunc := func() bool { + return len(ac.satisfiedByReplica[voterIndex][j])+ + len(ac.satisfiedByReplica[nonVoterIndex][j]) >= int(ac.constraints[j].numReplicas) + } + done := doneFunc() + if done { + continue + } + for kind := voterIndex; kind < numReplicaKinds; kind++ { + for i := range buf.replicaConstraintIndices[kind] { + constraintIndices := buf.replicaConstraintIndices[kind][i] + for _, index := range constraintIndices { + if index == int32(j) { + ac.satisfiedByReplica[kind][j] = + append(ac.satisfiedByReplica[kind][j], buf.replicas[kind][i].StoreID) + buf.replicaConstraintIndices[kind][i] = constraintIndices[:0] + done = doneFunc() + // This store is finished. break } } @@ -914,27 +899,45 @@ func (rac *rangeAnalyzedConstraints) finishInit( break } } + // done can be true if some store was appended to + // ac.satisfiedByReplica[kind][j] and made it fully satisfied. Don't + // need to look at other stores for this constraint. + if done { + break + } } - // Nothing over-satisfied. Go and greedily assign. - for kind := voterIndex; kind < numReplicaKinds; kind++ { - for i := range rac.buf.replicaConstraintIndices[kind] { - constraintIndices := rac.buf.replicaConstraintIndices[kind][i] - for _, index := range constraintIndices { - ac.satisfiedByReplica[kind][index] = - append(ac.satisfiedByReplica[kind][index], rac.replicas[kind][i].StoreID) - rac.buf.replicaConstraintIndices[kind][i] = constraintIndices[:0] - break - } + } + // Nothing over-satisfied. Go and greedily assign. + for kind := voterIndex; kind < numReplicaKinds; kind++ { + for i := range buf.replicaConstraintIndices[kind] { + constraintIndices := buf.replicaConstraintIndices[kind][i] + for _, index := range constraintIndices { + ac.satisfiedByReplica[kind][index] = + append(ac.satisfiedByReplica[kind][index], buf.replicas[kind][i].StoreID) + buf.replicaConstraintIndices[kind][i] = constraintIndices[:0] + break } } } +} + +func (rac *rangeAnalyzedConstraints) finishInit( + spanConfig *normalizedSpanConfig, + constraintMatcher storeMatchesConstraintInterface, + leaseholder roachpb.StoreID, +) { + rac.spanConfig = spanConfig + rac.numNeededReplicas[voterIndex] = spanConfig.numVoters + rac.numNeededReplicas[nonVoterIndex] = spanConfig.numReplicas - spanConfig.numVoters + rac.replicas = rac.buf.replicas + if spanConfig.constraints != nil { rac.constraints.constraints = spanConfig.constraints - analyzeFunc(&rac.constraints) + rac.constraints.analyzeFunc(&rac.buf, constraintMatcher) } if spanConfig.voterConstraints != nil { rac.voterConstraints.constraints = spanConfig.voterConstraints - analyzeFunc(&rac.voterConstraints) + rac.voterConstraints.analyzeFunc(&rac.buf, constraintMatcher) } rac.leaseholderID = leaseholder From cef090fe7ba6889839dd42a3c2c28ff82aaba673 Mon Sep 17 00:00:00 2001 From: wenyihu6 Date: Thu, 20 Nov 2025 21:01:19 -0500 Subject: [PATCH 2/5] mmaprototype: extract doneFunc This commit moves doneFunc out of analyzeFunc into its own helper function to improve readability. --- .../allocator/mmaprototype/constraint.go | 37 +++++++++++-------- 1 file changed, 22 insertions(+), 15 deletions(-) diff --git a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go index 377c23bfd803..df10b0c0fc60 100644 --- a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go +++ b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go @@ -830,6 +830,14 @@ type storeMatchesConstraintInterface interface { storeMatches(storeID roachpb.StoreID, constraintConj constraintsConj) bool } +// isConstraintSatisfied checks if the given constraint index has been fully +// satisfied by the stores currently assigned to it. +// TODO(wenyihu6): voter constraint should only count voters here (#158109) +func (ac *analyzedConstraints) isConstraintSatisfied(constraintIndex int) bool { + return len(ac.satisfiedByReplica[voterIndex][constraintIndex])+ + len(ac.satisfiedByReplica[nonVoterIndex][constraintIndex]) >= int(ac.constraints[constraintIndex].numReplicas) +} + func (ac *analyzedConstraints) analyzeFunc( buf *analyzeConstraintsBuf, constraintMatcher storeMatchesConstraintInterface, ) { @@ -871,12 +879,8 @@ func (ac *analyzedConstraints) analyzeFunc( // increasing order. Satisfy constraints in order, while not // oversatisfying. for j := range ac.constraints { - doneFunc := func() bool { - return len(ac.satisfiedByReplica[voterIndex][j])+ - len(ac.satisfiedByReplica[nonVoterIndex][j]) >= int(ac.constraints[j].numReplicas) - } - done := doneFunc() - if done { + satisfied := ac.isConstraintSatisfied(j) + if satisfied { continue } for kind := voterIndex; kind < numReplicaKinds; kind++ { @@ -887,22 +891,25 @@ func (ac *analyzedConstraints) analyzeFunc( ac.satisfiedByReplica[kind][j] = append(ac.satisfiedByReplica[kind][j], buf.replicas[kind][i].StoreID) buf.replicaConstraintIndices[kind][i] = constraintIndices[:0] - done = doneFunc() + satisfied = ac.isConstraintSatisfied(j) // This store is finished. break } } - // done can be true if some store was appended to - // ac.satisfiedByReplica[kind][j] and made it fully satisfied. Don't - // need to look at other stores for this constraint. - if done { + + // NB: We check the `satisfied` flag directly instead of calling + // `isConstraintSatisfied` again, since this value should be set + // in the previous loop correctly when the constraint becomes + // satisfied. + // + // Check if constraint is now satisfied to avoid + // over-satisfaction. + if satisfied { break } } - // done can be true if some store was appended to - // ac.satisfiedByReplica[kind][j] and made it fully satisfied. Don't - // need to look at other stores for this constraint. - if done { + // Check if constraint is now satisfied to avoid over-satisfaction. + if satisfied { break } } From b9fb467e18e82c71fef98ac07cf42263e7cb725f Mon Sep 17 00:00:00 2001 From: wenyihu6 Date: Thu, 20 Nov 2025 23:48:25 -0500 Subject: [PATCH 3/5] mmaprototype: add more comments for finishInit This commit adds more comments for rangeAnalyzedConstraints.finishInit. --- .../allocator/mmaprototype/allocator_state.go | 8 ++ .../allocator/mmaprototype/constraint.go | 89 ++++++++++++++++++- 2 files changed, 94 insertions(+), 3 deletions(-) diff --git a/pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go b/pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go index a27832b4af9c..b7c316cdfd23 100644 --- a/pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go +++ b/pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go @@ -703,6 +703,14 @@ func sortTargetCandidateSetAndPick( return cands.candidates[j].StoreID } +// ensureAnalyzedConstraints populates the constraints for the given range. +// It should be called when computing lease-transfers / rebalancing candidates +// for a range. +// +// NB: given rstate.conf should already be normalized using +// makeNormalizedSpanConfig. The caller is responsible for calling +// clearAnalyzedConstraints when rstate or the rstate.constraints is no longer +// needed. func (cs *clusterState) ensureAnalyzedConstraints(rstate *rangeState) { if rstate.constraints != nil { return diff --git a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go index df10b0c0fc60..7c37af41585b 100644 --- a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go +++ b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go @@ -690,7 +690,19 @@ type analyzedConstraints struct { // we omit considering it for a later conjunction. That is, we satisfy in a // greedy manner instead of considering all possibilities. So all the // satisfiedBy slices represent sets that are non-intersecting. - + // + // satisfiedByReplica[kind][i] contains the set of storeIDs that satisfy + // constraints[i]. Populated by analyzeFunc and used later by mmma to + // determine candidates to satisfy constraints such as + // candidatesToReplaceVoterForRebalance. + // For example, satisfiedByReplica[voterIndex][0] = [1, 2, 3], means that + // the first constraint (constraints[0]) is satisfied by storeIDs 1, 2, and + // 3. + // + // NB: this does not mean the current replica set satisfies constraints[i] + // since we may populate satisfiedByReplica[nonVoterIndex][i] for voter + // constraints as well. This just means that the store can satisfy + // constraints[i] regardless of the replica type. satisfiedByReplica [numReplicaKinds][][]roachpb.StoreID // These are stores that satisfy no constraint. Even though we are strict @@ -838,6 +850,65 @@ func (ac *analyzedConstraints) isConstraintSatisfied(constraintIndex int) bool { len(ac.satisfiedByReplica[nonVoterIndex][constraintIndex]) >= int(ac.constraints[constraintIndex].numReplicas) } +// analyzeFunc analyzes the current replica set and determines which constraints +// replicas satisfy, populating ac.satisfiedByReplica and +// ac.satisfiedNoConstraintReplica (both empty at the beginning). They are later +// used by mma to compute lease-transfer and rebalancing candidates by functions +// like candidatesToReplaceVoterForRebalance. The given buf.replicas is already +// populated with the current replica set from buf.tryAddingStore. +// +// For stores that satisfy no constraint, +// ac.satisfiedNoConstraintReplica[kind][i] will be populated with the replica +// storeIDs that satisfy no constraint. +// +// For stores that satisfy at least one constraint, +// ac.satisfiedByReplica[kind][i] is filled with the storeIDs of replicas that +// satisfy ac.constraints[i]. Each store should appear exactly once in the 2D +// slice ac.satisfiedByReplica[kind]. A constraint can end up over-satisfied +// (len(ac.satisfiedByReplica[kind][i]) > ac.constraints[i].numReplicas) in +// phase 3. Phase 2 below attempts to correct this by prioritizing +// under-satisfied constraints and assigning stores to those first. +// +// NB: the given ac.constraints can represent either replica or voter +// constraints, but we still populate both voter and non-voter indices +// regardless because we need to know which constraints are satisfied by all +// replicas to determine when promotions or demotions should occur. +// +// It has three phases: 1. For each replica (voter and non-voter), determine +// which constraint indices in ac.constraints its store satisfies. We record +// these matches in buf.replicaConstraintIndices by iterating over all replicas, +// all constraint conjunctions, and checking store satisfaction for each using +// constraintMatcher.storeMatches. +// +// buf.replicaConstraintIndices[kind][i] will be populated with the constraint +// indices that the i+1-th voter/non-voter satisfies. This serves as a scratch +// space. Note that as we assign constraint to stores that satisfy it, we will +// be clearing the constraint indices for that store in +// buf.replicaConstraintIndices[kind][i] to indicate that we cannot reuse the +// store to satisfy a different constraint. +// +// During this phase, it also populates ac.satisfiedByReplica[kind][i] for +// stores that satisfy exactly one constraint and populates +// ac.satisfiedNoConstraintReplica[kind][i] for stores that satisfy no +// constraint. Since we will be assigning at least one constraint to each store, +// these stores are unambiguous. +// +// 2. For each constraint, iterate over all replicas that satisfy it and assign +// replicas to the constraint in a round-robin manner, skipping once the +// constraint is satisfied. This phase does not allow over-satisfaction. +// TODO(wenyihu6): Add an example for phase 2 - if s1 satisfies c1(num=1) and s2 +// satisfies both c1(num=1) and c2(num=1), we should prefer assigning s2 to c2 +// rather than consuming it again for c1. TODO(wenyihu6): Verify behavior in +// cases like: s1 satisfies c1(num=1) and c2(num=1), s2 satisfies only c1, and +// s1 is used for c1 and c2 cannot be satisfied although we could have. +// +// 3. For any remaining stores that satisfy multiple constraints, assign each to +// the first constraint it satisfies. This phase allows over-satisfaction, and +// some constraints may still end up under- or over-filled. + +// TODO(wenyihu6): document the lifecycle of scratch space +// replicaConstraintIndices once understood. TODO(wenyihu6): add more tests + +// examples here func (ac *analyzedConstraints) analyzeFunc( buf *analyzeConstraintsBuf, constraintMatcher storeMatchesConstraintInterface, ) { @@ -928,6 +999,9 @@ func (ac *analyzedConstraints) analyzeFunc( } } +// finishInit analyzes the span config and the range’s current replica set. It +// prepares the MMA to compute lease-transfer and rebalancing candidates while +// satisfying both constraint requirements and leaseholder preferences. func (rac *rangeAnalyzedConstraints) finishInit( spanConfig *normalizedSpanConfig, constraintMatcher storeMatchesConstraintInterface, @@ -1359,8 +1433,17 @@ func isNonVoter(typ roachpb.ReplicaType) bool { type analyzeConstraintsBuf struct { replicas [numReplicaKinds][]storeAndLocality - // Scratch space. replicaConstraintIndices[k][i] is the constraint matching - // state for replicas[k][i]. + // Scratch space. buf.replicaConstraintIndices[kind][i] contains the + // constraints indices that the buf.replicas[kind][i] replica satisfies. + // + // For example, buf.replicaConstraintIndices[voterIndex][0] = [0, 1, 2], + // means that the first voter replica (buf.replicas[voterIndex][0]) + // satisfies the 0th, 1st, and 2nd constraint conjunction in ac.constraints. + // + // NB: given ac.constraints can represent both replica and voter + // constraints. We still populate both voter and non-voter indexes because + // we need to know which constraints are satisfied by all replicas to + // determine when promotions or demotions should occur. replicaConstraintIndices [numReplicaKinds][][]int32 } From 93c9a0820b38fa35c0c71033c2fe338a45847b76 Mon Sep 17 00:00:00 2001 From: wenyihu6 Date: Fri, 21 Nov 2025 00:15:24 -0500 Subject: [PATCH 4/5] mmaprototype: extract diversity score computation This commit moves diversity score computation out of analyzeFunc into its own helper function to improve readability. --- .../allocator/mmaprototype/constraint.go | 106 ++++++++++++------ 1 file changed, 70 insertions(+), 36 deletions(-) diff --git a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go index 7c37af41585b..9b1b5b081873 100644 --- a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go +++ b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go @@ -999,6 +999,75 @@ func (ac *analyzedConstraints) analyzeFunc( } } +// diversityFunc computes the diversity score between two sets of stores. +// +// When sameStores is false, intersection between this and other is empty and no +// de-duplication is needed. For example, given two sets of stores [1, 2, 3] and +// [4, 5, 6], diversity score is computed among all pairs (1, 4), (1, 5), (1, +// 6), (2, 4), (2, 5), (2, 6), (3, 4), (3, 5), (3, 6). +// +// When sameStores is true, this and other are the same set and de-duplication +// is needed. For example, given two sets of stores [1, 2, 3] and [1, 2, 3], +// diversity score should be computed among pairs (1, 2), (1, 3), (2, 3) only. +func diversityFunc( + this []storeAndLocality, other []storeAndLocality, sameStores bool, +) (sumScore float64, numSamples int) { + for i := range this { + s1 := this[i] + for j := range other { + s2 := other[j] + // Only compare pairs of replicas where s2.StoreID > s1.StoreID to avoid + // computing the diversity score between each pair of stores twice. + if sameStores && s2.StoreID <= s1.StoreID { + continue + } + sumScore += s1.localityTiers.diversityScore(s2.localityTiers) + numSamples++ + } + } + return sumScore, numSamples +} + +// diversityScore returns the voter and replica diversity scores of the given replicas sets. +// +// voterDiversityScore: sum of diversity scores over all pairs from V×V(voter-voter pairs). +// replicaDiversityScore: sum of diversity scores over all pairs from: +// V×V(voter–voter pairs), N×N(non-voter–non-voter pairs), V×N(voter–non-voter +// pairs). +// +// diversity score of a single pair is computed using +// localityTiers.diversityScore which finds the longest common prefix of two +// localities. +func diversityScore( + replicas [numReplicaKinds][]storeAndLocality, +) (voterDiversityScore, replicaDiversityScore float64) { + // Helper to calculate average, or a max-value if no samples (represents lowest possible diversity) + scoreFromSumAndSamples := func(sumScore float64, numSamples int) float64 { + if numSamples == 0 { + return roachpb.MaxDiversityScore + } + return sumScore / float64(numSamples) + } + + // Calculate voter diversity (V × V, no self-pairs) + voterSum, voterSamples := diversityFunc(replicas[voterIndex], replicas[voterIndex], true) + // Calculate replica diversity over all pairs: + // (voter, voter), (nonvoter, nonvoter), (voter, nonvoter) + totalSum, totalSamples := 0.0, 0 + vs, ns := diversityFunc(replicas[voterIndex], replicas[voterIndex], true) + totalSum += vs + totalSamples += ns + + vs, ns = diversityFunc(replicas[nonVoterIndex], replicas[nonVoterIndex], true) + totalSum += vs + totalSamples += ns + + vs, ns = diversityFunc(replicas[voterIndex], replicas[nonVoterIndex], false) + totalSum += vs + totalSamples += ns + return scoreFromSumAndSamples(voterSum, voterSamples), scoreFromSumAndSamples(totalSum, totalSamples) +} + // finishInit analyzes the span config and the range’s current replica set. It // prepares the MMA to compute lease-transfer and rebalancing candidates while // satisfying both constraint requirements and leaseholder preferences. @@ -1058,42 +1127,7 @@ func (rac *rangeAnalyzedConstraints) finishInit( rac.replicaLocalityTiers = makeReplicasLocalityTiers(replicaLocalityTiers) rac.voterLocalityTiers = makeReplicasLocalityTiers(voterLocalityTiers) - diversityFunc := func( - stores1 []storeAndLocality, stores2 []storeAndLocality, sameStores bool, - ) (sumScore float64, numSamples int) { - for i := range stores1 { - s1 := stores1[i] - for j := range stores2 { - s2 := stores2[j] - // Only compare pairs of replicas where s2.StoreID > s1.StoreID to avoid - // computing the diversity score between each pair of stores twice. - if sameStores && s2.StoreID <= s1.StoreID { - continue - } - sumScore += s1.localityTiers.diversityScore(s2.localityTiers) - numSamples++ - } - } - return sumScore, numSamples - } - scoreFromSumAndSamples := func(sumScore float64, numSamples int) float64 { - if numSamples == 0 { - return roachpb.MaxDiversityScore - } - return sumScore / float64(numSamples) - } - sumVoterScore, numVoterSamples := diversityFunc( - rac.replicas[voterIndex], rac.replicas[voterIndex], true) - rac.votersDiversityScore = scoreFromSumAndSamples(sumVoterScore, numVoterSamples) - - sumReplicaScore, numReplicaSamples := sumVoterScore, numVoterSamples - srs, nrs := diversityFunc(rac.replicas[nonVoterIndex], rac.replicas[nonVoterIndex], true) - sumReplicaScore += srs - numReplicaSamples += nrs - srs, nrs = diversityFunc(rac.replicas[voterIndex], rac.replicas[nonVoterIndex], false) - sumReplicaScore += srs - numReplicaSamples += nrs - rac.replicasDiversityScore = scoreFromSumAndSamples(sumReplicaScore, numReplicaSamples) + rac.votersDiversityScore, rac.replicasDiversityScore = diversityScore(rac.replicas) } // Disjunction of conjunctions. From 83131a127cf62f3940b824063af1983f6b2a3a49 Mon Sep 17 00:00:00 2001 From: wenyihu6 Date: Fri, 21 Nov 2025 00:48:26 -0500 Subject: [PATCH 5/5] fixup! mmaprototype: add more comments for finishInit --- pkg/kv/kvserver/allocator/mmaprototype/constraint.go | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go index 9b1b5b081873..ee6c34b9dff4 100644 --- a/pkg/kv/kvserver/allocator/mmaprototype/constraint.go +++ b/pkg/kv/kvserver/allocator/mmaprototype/constraint.go @@ -864,10 +864,13 @@ func (ac *analyzedConstraints) isConstraintSatisfied(constraintIndex int) bool { // For stores that satisfy at least one constraint, // ac.satisfiedByReplica[kind][i] is filled with the storeIDs of replicas that // satisfy ac.constraints[i]. Each store should appear exactly once in the 2D -// slice ac.satisfiedByReplica[kind]. A constraint can end up over-satisfied +// slice ac.satisfiedByReplica[kind]. Phase 2 below attempts to be optimal by +// prioritizing under-satisfied constraints and assigning stores to those first. +// A constraint may end up being over-satisfied // (len(ac.satisfiedByReplica[kind][i]) > ac.constraints[i].numReplicas) in -// phase 3. Phase 2 below attempts to correct this by prioritizing -// under-satisfied constraints and assigning stores to those first. +// phase 3. A constraint may remain under-satisfied. If a store is +// under-satisfied in phase 2, it cannot be corrected in phase 3 and will +// continue to be under-satisfied. (TODO: Is this right during review) // // NB: the given ac.constraints can represent either replica or voter // constraints, but we still populate both voter and non-voter indices