storage: rebalance replicas based on load at the store level

Follow-up to cockroachdb#28340, which did this for just leases.

Fixes cockroachdb#17979

Release note (performance improvement): Range replicas will be
automatically rebalanced throughout the cluster to even out the amount
of QPS being handled by each node.

docs/generated/settings/settings.html

@@ -18,8 +18,8 @@
 <tr><td><code>kv.allocator.lease_rebalancing_aggressiveness</code></td><td>float</td><td><code>1</code></td><td>set greater than 1.0 to rebalance leases toward load more aggressively, or between 0 and 1.0 to be more conservative about rebalancing leases</td></tr>
 <tr><td><code>kv.allocator.load_based_lease_rebalancing.enabled</code></td><td>boolean</td><td><code>true</code></td><td>set to enable rebalancing of range leases based on load and latency</td></tr>
 <tr><td><code>kv.allocator.range_rebalance_threshold</code></td><td>float</td><td><code>0.05</code></td><td>minimum fraction away from the mean a store's range count can be before it is considered overfull or underfull</td></tr>
-<tr><td><code>kv.allocator.stat_based_rebalancing.enabled</code></td><td>boolean</td><td><code>false</code></td><td>set to enable rebalancing range replicas and leases to more evenly distribute read and write load across the stores in a cluster</td></tr>
-<tr><td><code>kv.allocator.stat_rebalance_threshold</code></td><td>float</td><td><code>0.2</code></td><td>minimum fraction away from the mean a store's stats (like disk usage or writes per second) can be before it is considered overfull or underfull</td></tr>
+<tr><td><code>kv.allocator.stat_based_rebalancing.enabled</code></td><td>boolean</td><td><code>true</code></td><td>set to enable rebalancing range replicas and leases to more evenly distribute read and write load across the stores in a cluster</td></tr>
+<tr><td><code>kv.allocator.stat_rebalance_threshold</code></td><td>float</td><td><code>0.2</code></td><td>minimum fraction away from the mean a store's stats (such as queries per second) can be before it is considered overfull or underfull</td></tr>
 <tr><td><code>kv.bulk_io_write.concurrent_export_requests</code></td><td>integer</td><td><code>5</code></td><td>number of export requests a store will handle concurrently before queuing</td></tr>
 <tr><td><code>kv.bulk_io_write.concurrent_import_requests</code></td><td>integer</td><td><code>1</code></td><td>number of import requests a store will handle concurrently before queuing</td></tr>
 <tr><td><code>kv.bulk_io_write.max_rate</code></td><td>byte size</td><td><code>8.0 EiB</code></td><td>the rate limit (bytes/sec) to use for writes to disk on behalf of bulk io ops</td></tr>

pkg/storage/allocator.go

@@ -361,9 +361,36 @@ func (a *Allocator) AllocateTarget(
 ) (*roachpb.StoreDescriptor, string, error) {
 	sl, aliveStoreCount, throttledStoreCount := a.storePool.getStoreList(rangeInfo.Desc.RangeID, storeFilterThrottled)
 
+	target, details := a.allocateTargetFromList(
+		ctx, sl, zone, existing, rangeInfo, a.scorerOptions(disableStatsBasedRebalancing))
+
+	if target != nil {
+		return target, details, nil
+	}
+
+	// When there are throttled stores that do match, we shouldn't send
+	// the replica to purgatory.
+	if throttledStoreCount > 0 {
+		return nil, "", errors.Errorf("%d matching stores are currently throttled", throttledStoreCount)
+	}
+	return nil, "", &allocatorError{
+		constraints:      zone.Constraints,
+		existingReplicas: len(existing),
+		aliveStores:      aliveStoreCount,
+		throttledStores:  throttledStoreCount,
+	}
+}
+
+func (a *Allocator) allocateTargetFromList(
+	ctx context.Context,
+	sl StoreList,
+	zone config.ZoneConfig,
+	existing []roachpb.ReplicaDescriptor,
+	rangeInfo RangeInfo,
+	options scorerOptions,
+) (*roachpb.StoreDescriptor, string) {
 	analyzedConstraints := analyzeConstraints(
 		ctx, a.storePool.getStoreDescriptor, existing, zone)
-	options := a.scorerOptions(disableStatsBasedRebalancing)
 	candidates := allocateCandidates(
 		sl, analyzedConstraints, existing, rangeInfo, a.storePool.getLocalities(existing), options,
 	)
@@ -379,20 +406,10 @@ func (a *Allocator) AllocateTarget(
 		if err != nil {
 			log.Warningf(ctx, "failed to marshal details for choosing allocate target: %s", err)
 		}
-		return &target.store, string(detailsBytes), nil
+		return &target.store, string(detailsBytes)
 	}
 
-	// When there are throttled stores that do match, we shouldn't send
-	// the replica to purgatory.
-	if throttledStoreCount > 0 {
-		return nil, "", errors.Errorf("%d matching stores are currently throttled", throttledStoreCount)
-	}
-	return nil, "", &allocatorError{
-		constraints:      zone.Constraints,
-		existingReplicas: len(existing),
-		aliveStores:      aliveStoreCount,
-		throttledStores:  throttledStoreCount,
-	}
+	return nil, ""
 }
 
 func (a Allocator) simulateRemoveTarget(

pkg/storage/allocator_scorer.go

@@ -60,7 +60,7 @@ const (
 var EnableStatsBasedRebalancing = settings.RegisterBoolSetting(
 	"kv.allocator.stat_based_rebalancing.enabled",
 	"set to enable rebalancing range replicas and leases to more evenly distribute read and write load across the stores in a cluster",
-	false, // TODO(a-robinson): switch to true for v2.1 once the store-rebalancer is done
+	true,
 )
 
 // rangeRebalanceThreshold is the minimum ratio of a store's range count to
@@ -85,12 +85,13 @@ var rangeRebalanceThreshold = settings.RegisterNonNegativeFloatSetting(
 // TODO(a-robinson): Should disk usage be held to a higher standard than this?
 var statRebalanceThreshold = settings.RegisterNonNegativeFloatSetting(
 	"kv.allocator.stat_rebalance_threshold",
-	"minimum fraction away from the mean a store's stats (like disk usage or writes per second) can be before it is considered overfull or underfull",
+	"minimum fraction away from the mean a store's stats (such as queries per second) can be before it is considered overfull or underfull",
 	0.20,
 )
 
 type scorerOptions struct {
 	deterministic                bool
+	balanceQPSInsteadOfCount     bool
 	statsBasedRebalancingEnabled bool
 	rangeRebalanceThreshold      float64
 	statRebalanceThreshold       float64
@@ -442,11 +443,22 @@ func allocateCandidates(
 		}
 		diversityScore := diversityAllocateScore(s, existingNodeLocalities)
 		balanceScore := balanceScore(sl, s.Capacity, rangeInfo, options)
+		var convergesScore int
+		if options.balanceQPSInsteadOfCount {
+			if s.Capacity.QueriesPerSecond < underfullStatThreshold(options, sl.candidateQueriesPerSecond.mean) {
+				convergesScore = 1
+			} else if s.Capacity.QueriesPerSecond < sl.candidateQueriesPerSecond.mean {
+				convergesScore = 0
+			} else {
+				convergesScore = -1
+			}
+		}
 		candidates = append(candidates, candidate{
 			store:          s,
 			valid:          constraintsOK,
 			necessary:      necessary,
 			diversityScore: diversityScore,
+			convergesScore: convergesScore,
 			balanceScore:   balanceScore,
 			rangeCount:     int(s.Capacity.RangeCount),
 		})

pkg/storage/replica_rankings.go

@@ -28,6 +28,7 @@ const (
 type replicaWithStats struct {
 	repl *Replica
 	qps  float64
+	// TODO(a-robinson): Include writes-per-second and logicalBytes of storage?
 }
 
 // replicaRankings maintains top-k orderings of the replicas in a store along

pkg/storage/store.go

@@ -1560,7 +1560,7 @@ func (s *Store) Start(ctx context.Context, stopper *stop.Stopper) error {
 	// Connect rangefeeds to closed timestamp updates.
 	s.startClosedTimestampRangefeedSubscriber(ctx)
 
-	s.storeRebalancer.Start(ctx, s.stopper, s.StoreID())
+	s.storeRebalancer.Start(ctx, s.stopper)
 
 	// Start the storage engine compactor.
 	if envutil.EnvOrDefaultBool("COCKROACH_ENABLE_COMPACTOR", true) {