allocatorimpl: populate advisor under rebalanceOptions

Previously, `bestRebalanceTarget` took a map and populated it with (option
index->advisor) during the loop. This commit changes it to populate the advisor
directly within the `rebalanceOptions` struct as a field.

Author: wenyihu6

pkg/kv/kvserver/allocator/allocatorimpl/allocator.go

@@ -17,7 +17,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/load"
-	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/mmaprototype"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/constraint"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvflowcontrol/rac2"
@@ -1890,11 +1889,10 @@ func (a Allocator) RebalanceTarget(
 	// to compute the means for the original set {s1, s2, s3} ∪ {existing} once
 	// and then use it for all calls to MMA, and bestRebalanceTarget does not need
 	// to copy the candidate set.
-	advisors := make(map[int]*mmaprototype.MMARebalanceAdvisor, len(results))
 	var bestIdx int
 
 	for {
-		target, existingCandidate, bestIdx = bestRebalanceTarget(a.randGen, results, a.as, advisors)
+		target, existingCandidate, bestIdx = bestRebalanceTarget(a.randGen, results, a.as)
 		if target == nil {
 			return zero, zero, "", false
 		}
@@ -1912,7 +1910,7 @@ func (a Allocator) RebalanceTarget(
 			// rebalance, we will continue to the next target. Note that this target
 			// would have been deleted from the candidates set in bestRebalanceTarget,
 			// so we will not select it again.
-			if advisor, ok := advisors[bestIdx]; ok {
+			if advisor := results[bestIdx].advisor; advisor != nil {
 				if a.as.IsInConflictWithMMA(target.store.StoreID, advisor, false) {
 					continue
 				}

pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer.go

@@ -1420,6 +1420,7 @@ func candidateListForRemoval(
 type rebalanceOptions struct {
 	existing   candidate
 	candidates candidateList
+	advisor    *mmaprototype.MMARebalanceAdvisor
 }
 
 // equivalenceClass captures the set of "equivalent" replacement candidates
@@ -1929,10 +1930,7 @@ func rankedCandidateListForRebalancing(
 // Contract: responsible for making sure that the returned bestIdx has the
 // corresponding MMA advisor in advisors.
 func bestRebalanceTarget(
-	randGen allocatorRand,
-	options []rebalanceOptions,
-	as *mmaintegration.AllocatorSync,
-	advisors map[int]*mmaprototype.MMARebalanceAdvisor,
+	randGen allocatorRand, options []rebalanceOptions, as *mmaintegration.AllocatorSync,
 ) (target, existingCandidate *candidate, bestIdx int) {
 	bestIdx = -1
 	var bestTarget *candidate
@@ -1957,12 +1955,12 @@ func bestRebalanceTarget(
 	}
 	// For the first time an option in options is selected, build and cache the
 	// corresponding MMA advisor in advisors[bestIdx].
-	if _, ok := advisors[bestIdx]; !ok {
+	if options[bestIdx].advisor == nil {
 		stores := make([]roachpb.StoreID, 0, len(options[bestIdx].candidates))
 		for _, cand := range options[bestIdx].candidates {
 			stores = append(stores, cand.store.StoreID)
 		}
-		advisors[bestIdx] = as.BuildMMARebalanceAdvisor(options[bestIdx].existing.store.StoreID, stores)
+		options[bestIdx].advisor = as.BuildMMARebalanceAdvisor(options[bestIdx].existing.store.StoreID, stores)
 	}
 
 	// Copy the selected target out of the candidates slice before modifying