kvserver: actuate load-based replica rebalancing under heterogeneous …

…localities

This commit teaches the `StoreRebalancer` to make load-based rebalancing
decisions that are meaningful within the context of the replication constraints
placed on the ranges being relocated and the set of stores that can legally
receive replicas for such ranges.

Previously, the `StoreRebalancer` would compute the QPS underfull and overfull
thresholds based on the overall average QPS being served by all stores in the
cluster. Notably, this included stores that were in replication zones that
would not satisfy required constraints for the range being considered for
rebalancing. This meant that the store rebalancer would effectively never be
able to rebalance ranges within the stores inside heavily loaded replication
zones (since all the _valid_ stores would be above the overfull thresholds).

This patch is a move away from the bespoke relocation logic in the
`StoreRebalancer`. Instead, we have the `StoreRebalancer` rely on the
rebalancing logic used by the `replicateQueue` that already has the machinery
to compute load based signals for candidates _relative to other comparable
stores_. The main difference here is that the `StoreRebalancer` uses this
machinery to promote convergence of QPS across stores, whereas the
`replicateQueue` uses it to promote convergence of range counts. A series of
preceeding commits in this patchset generalize the existing replica rebalancing
logic, and this commit teaches the `StoreRebalancer` to use it.

This generalization also addresses another key limitation (see cockroachdb#62922) of the
`StoreRebalancer` regarding its inability to make partial improvements to a
range. Previously, if the `StoreRebalancer` couldn't move a range _entirely_
off of overfull stores, it would give up and not even move the subset of
replicas it could. This is no longer the case.

Resolves cockroachdb#61883
Resolves cockroachdb#62992

Release justification: Fixes high priority bug

Release note (performance improvement): QPS-based replica rebalancing is now
aware of different constraints placed on different replication zones. This
means that heterogeneously loaded replication zones (for instance, regions)
will achieve a more even distribution of QPS within the stores inside each
such zone.

/cc @cockroachdb/kv

pkg/kv/kvserver/allocator.go

@@ -1083,7 +1083,7 @@ func (a Allocator) rebalanceTarget(
  return zero, zero, "", false
  }
 
- // Add a fake new replica to our copy of the range descriptor so that we can
+ // Add a fake new replica to our copy of the replica descriptor so that we can
  // simulate the removal logic. If we decide not to go with this target, note
  // that this needs to be removed from desc before we try any other target.
  newReplica := roachpb.ReplicaDescriptor{

pkg/kv/kvserver/allocator_scorer.go

@@ -86,10 +86,14 @@ var rangeRebalanceThreshold = func() *settings.FloatSetting {
  return s
 }()
 
+// TODO(aayush): Maybe turn this into an interface with one implementation
+// that cares about range count and another that cares about QPS, so its
+// impossible to misuse.
 type scorerOptions struct {
  deterministic bool
  rangeRebalanceThreshold float64
- qpsRebalanceThreshold float64 // only considered if non-zero
+ // Only used if `rangeRebalanceThreshold` is not set.
+ qpsRebalanceThreshold float64
 }
 
 // candidate store for allocation.
@@ -442,26 +446,11 @@ func rankedCandidateListForAllocation(
  }
  diversityScore := diversityAllocateScore(s, existingStoreLocalities)
  balanceScore := balanceScore(candidateStores, s.Capacity, options)
- var convergesScore int
- if options.qpsRebalanceThreshold > 0 {
- if s.Capacity.QueriesPerSecond < underfullQPSThreshold(
- options, candidateStores.candidateQueriesPerSecond.mean) {
- convergesScore = 1
- } else if s.Capacity.QueriesPerSecond < candidateStores.candidateQueriesPerSecond.mean {
- convergesScore = 0
- } else if s.Capacity.QueriesPerSecond < overfullQPSThreshold(
- options, candidateStores.candidateQueriesPerSecond.mean) {
- convergesScore = -1
- } else {
- convergesScore = -2
- }
- }
  candidates = append(candidates, candidate{
  store: s,
  valid: constraintsOK,
  necessary: necessary,
  diversityScore: diversityScore,
- convergesScore: convergesScore,
  balanceScore: balanceScore,
  rangeCount: int(s.Capacity.RangeCount),
  })
@@ -792,11 +781,11 @@ func rankedCandidateListForRebalancing(
  existingCandidates = append(existingCandidates, existing)
  continue
  }
- balanceScore := balanceScore(comparable.sl, existing.store.Capacity, options)
  // Similarly to in candidateListForRemoval, any replica whose
  // removal would not converge the range stats to their mean is given a
  // constraint score boost of 1 to make it less attractive for removal.
  convergesScore := rebalanceFromConvergesScore(comparable.sl, existing.store.Capacity, options)
+ balanceScore := balanceScore(comparable.sl, existing.store.Capacity, options)
  existing.convergesScore = convergesScore
  existing.balanceScore = balanceScore
  existing.rangeCount = int(existing.store.Capacity.RangeCount)
@@ -1314,16 +1303,17 @@ func diversityAllocateScore(
 }
 
 // diversityRemovalScore returns a value between 0 and 1 based on how desirable
-// it would be to remove a node's replica of a range. A higher score indicates
+// it would be to remove a store's replica of a range. A higher score indicates
 // that the node is a better fit (i.e. keeping it around is good for diversity).
 func diversityRemovalScore(
  storeID roachpb.StoreID, existingStoreLocalities map[roachpb.StoreID]roachpb.Locality,
 ) float64 {
  var sumScore float64
  var numSamples int
  locality := existingStoreLocalities[storeID]
- // We don't need to calculate the overall diversityScore for the range, because the original overall diversityScore
- // of this range is always the same.
+ // We don't need to calculate the overall diversityScore for the range,
+ // because the original overall diversityScore of this range is always the
+ // same.
  for otherStoreID, otherLocality := range existingStoreLocalities {
  if otherStoreID == storeID {
  continue
@@ -1339,11 +1329,10 @@ func diversityRemovalScore(
 }
 
 // diversityRebalanceScore returns a value between 0 and 1 based on how
-// desirable it would be to rebalance away from an existing store to the target
-// store. Because the store to be removed isn't specified, this assumes that
+// desirable it would be to rebalance to `store` from any of the existing
+// stores. Because the store to be removed isn't specified, this assumes that
 // the worst-fitting store from a diversity perspective will be removed. A
-// higher score indicates that the provided store is a better fit for the
-// range.
+// higher score indicates that the provided store is a better fit for the range.
 func diversityRebalanceScore(
  store roachpb.StoreDescriptor, existingStoreLocalities map[roachpb.StoreID]roachpb.Locality,
 ) float64 {
@@ -1371,20 +1360,20 @@ func diversityRebalanceScore(
 func diversityRebalanceFromScore(
  store roachpb.StoreDescriptor,
  fromStoreID roachpb.StoreID,
- existingNodeLocalities map[roachpb.StoreID]roachpb.Locality,
+ existingStoreLocalities map[roachpb.StoreID]roachpb.Locality,
 ) float64 {
  // Compute the pairwise diversity score of all replicas that will exist
  // after adding store and removing fromNodeID.
  var sumScore float64
  var numSamples int
- for storeID, locality := range existingNodeLocalities {
+ for storeID, locality := range existingStoreLocalities {
  if storeID == fromStoreID {
  continue
  }
  newScore := store.Locality().DiversityScore(locality)
  sumScore += newScore
  numSamples++
- for otherStoreID, otherLocality := range existingNodeLocalities {
+ for otherStoreID, otherLocality := range existingStoreLocalities {
  // Only compare pairs of replicas where otherNodeID > nodeID to avoid
  // computing the diversity score between each pair of localities twice.
  if otherStoreID <= storeID || otherStoreID == fromStoreID {

pkg/kv/kvserver/allocator_test.go

@@ -1525,7 +1525,7 @@ func TestAllocatorRemoveBasedOnQPS(t *testing.T) {
  }
  remove, _, err := a.RemoveVoter(
  ctx,
- zonepb.EmptyCompleteZoneConfig(),
+ emptySpanConfig(),
  subtest.existingRepls,
  subtest.existingRepls,
  nil,

pkg/kv/kvserver/store_pool.go

@@ -793,8 +793,9 @@ func (sl StoreList) String() string {
  return buf.String()
 }
 
-// filter takes a store list and filters it using the passed in constraints. It
-// maintains the original order of the passed in store list.
+// filter takes a store list and removes stores that would be explicitly invalid
+// under the given set of constraints. It maintains the original order of the
+// passed in store list.
 func (sl StoreList) filter(constraints []roachpb.ConstraintsConjunction) StoreList {
  if len(constraints) == 0 {
  return sl