allocator: hoist some calls to Locality

store.Locality() allocates. Previously we were calling this inside an
double-nested loop in diversityRebalanceScore which itself is called in
a loop.

This is probably just a drop in the bucket given how much is going on in
the callers of this function, but perhaps it helps a bit:

    BenchmarkRebalanceToDiversityScore        427075              2734 ns/op            1920 B/op         12 allocs/op
    BenchmarkRebalanceToDiversityScore        904476              1443 ns/op             160 B/op          1 allocs/op

Informs #147800

Release note: None

Author: stevendanna

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

@@ -1159,7 +1159,7 @@ func rankedCandidateListForAllocation(
 			continue
 		}
 
-		diversityScore := diversityAllocateScore(s, existingStoreLocalities)
+		diversityScore := diversityAllocateScore(s.Locality(), existingStoreLocalities)
 		balanceScore := options.balanceScore(validStoreList, s.Capacity)
 		var hasNonVoter bool
 		if targetType == VoterTarget {
@@ -1656,9 +1656,10 @@ func rankedCandidateListForRebalancing(
 			// this stage, in additon to hard checks and validation.
 			// TODO(kvoli,ayushshah15): Refactor this to make it harder to
 			// inadvertently break the invariant above,
+			locality := store.Locality()
 			constraintsOK, necessary, voterNecessary := rebalanceConstraintsChecker(store, existing.store)
 			diversityScore := diversityRebalanceFromScore(
-				store, existing.store.StoreID, existingStoreLocalities)
+				locality, existing.store.StoreID, existingStoreLocalities)
 			cand := candidate{
 				store:          store,
 				valid:          constraintsOK,
@@ -1676,7 +1677,7 @@ func rankedCandidateListForRebalancing(
 						"s%d: should-rebalance(necessary/diversity=s%d): oldNecessary:%t, newNecessary:%t, "+
 							"oldDiversity:%f, newDiversity:%f, locality:%q",
 						existing.store.StoreID, store.StoreID, existing.necessary, cand.necessary,
-						existing.diversityScore, cand.diversityScore, store.Locality())
+						existing.diversityScore, cand.diversityScore, locality)
 				}
 			}
 		}
@@ -2236,7 +2237,7 @@ func RangeDiversityScore(existingStoreLocalities map[roachpb.StoreID]roachpb.Loc
 // desirable it would be to add a replica to store. A higher score means the
 // store is a better fit.
 func diversityAllocateScore(
-	store roachpb.StoreDescriptor, existingStoreLocalities map[roachpb.StoreID]roachpb.Locality,
+	storeLocality roachpb.Locality, existingStoreLocalities map[roachpb.StoreID]roachpb.Locality,
 ) float64 {
 	var sumScore float64
 	var numSamples int
@@ -2245,7 +2246,7 @@ func diversityAllocateScore(
 	// consider adding the pairwise average diversity of the existing replicas
 	// is the same.
 	for _, locality := range existingStoreLocalities {
-		newScore := store.Locality().DiversityScore(locality)
+		newScore := storeLocality.DiversityScore(locality)
 		sumScore += newScore
 		numSamples++
 	}
@@ -2296,8 +2297,9 @@ func diversityRebalanceScore(
 	var maxScore float64
 	// For every existing node, calculate what the diversity score would be if we
 	// remove that node's replica to replace it with one on the provided store.
+	storeLocality := store.Locality()
 	for removedStoreID := range existingStoreLocalities {
-		score := diversityRebalanceFromScore(store, removedStoreID, existingStoreLocalities)
+		score := diversityRebalanceFromScore(storeLocality, removedStoreID, existingStoreLocalities)
 		if score > maxScore {
 			maxScore = score
 		}
@@ -2312,7 +2314,7 @@ func diversityRebalanceScore(
 // A higher score indicates that the provided store is a better fit for the
 // range.
 func diversityRebalanceFromScore(
-	store roachpb.StoreDescriptor,
+	storeLocality roachpb.Locality,
 	fromStoreID roachpb.StoreID,
 	existingStoreLocalities map[roachpb.StoreID]roachpb.Locality,
 ) float64 {
@@ -2324,7 +2326,7 @@ func diversityRebalanceFromScore(
 		if storeID == fromStoreID {
 			continue
 		}
-		newScore := store.Locality().DiversityScore(locality)
+		newScore := storeLocality.DiversityScore(locality)
 		sumScore += newScore
 		numSamples++
 		for otherStoreID, otherLocality := range existingStoreLocalities {

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

@@ -1767,7 +1767,7 @@ func TestAllocateDiversityScore(t *testing.T) {
 					continue
 				}
 				var score storeScore
-				actualScore := diversityAllocateScore(s, existingStoreLocalities)
+				actualScore := diversityAllocateScore(s.Locality(), existingStoreLocalities)
 				score.storeID = s.StoreID
 				score.score = actualScore
 				scores = append(scores, score)
@@ -1858,6 +1858,23 @@ func TestRebalanceToDiversityScore(t *testing.T) {
 	}
 }
 
+func BenchmarkRebalanceToDiversityScore(b *testing.B) {
+	existingStoreLocalities := make(map[roachpb.StoreID]roachpb.Locality, len(testStores))
+	store := testStores[testStoreUSa15]
+	for sID := range testStores {
+		if sID == testStoreUSa15 {
+			continue
+		}
+		existingStoreLocalities[testStores[sID].StoreID] = testStores[sID].Locality()
+	}
+
+	b.ResetTimer()
+	sum := float64(0)
+	for i := 0; i < b.N; i++ {
+		sum += diversityRebalanceScore(store, existingStoreLocalities)
+	}
+}
+
 func TestRemovalDiversityScore(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -1973,7 +1990,7 @@ func TestDiversityScoreEquivalence(t *testing.T) {
 			s := testStores[storeID]
 			fromStoreID := s.StoreID
 			s.StoreID = 99
-			rebalanceScore := diversityRebalanceFromScore(s, fromStoreID, existingLocalities)
+			rebalanceScore := diversityRebalanceFromScore(s.Locality(), fromStoreID, existingLocalities)
 			if a, e := rebalanceScore, tc.expected; !scoresAlmostEqual(a, e) {
 				t.Errorf("diversityRebalanceFromScore(%v, %d, %v) got %f, want %f",
 					s, fromStoreID, existingLocalities, a, e)