storage: Avoid using different qps measurements in StoreRebalancer

It's better to use the same measurement for everything, especially
because it's the one that was used to compute the store's total QPS.

Fixes cockroachdb#31303

Release note: None

pkg/storage/replicate_queue.go

@@ -590,24 +590,25 @@ func (rq *replicateQueue) findTargetAndTransferLease(
  return false, nil
  }
 
- err := rq.transferLease(ctx, repl, target)
+ avgQPS, qpsMeasurementDur := repl.leaseholderStats.avgQPS()
+ if qpsMeasurementDur < MinStatsDuration {
+ avgQPS = 0
+ }
+ err := rq.transferLease(ctx, repl, target, avgQPS)
  return err == nil, err
 }
 
 func (rq *replicateQueue) transferLease(
- ctx context.Context, repl *Replica, target roachpb.ReplicaDescriptor,
+ ctx context.Context, repl *Replica, target roachpb.ReplicaDescriptor, rangeQPS float64,
 ) error {
  rq.metrics.TransferLeaseCount.Inc(1)
  log.VEventf(ctx, 1, "transferring lease to s%d", target.StoreID)
- avgQPS, qpsMeasurementDur := repl.leaseholderStats.avgQPS()
  if err := repl.AdminTransferLease(ctx, target.StoreID); err != nil {
  return errors.Wrapf(err, "%s: unable to transfer lease to s%d", repl, target.StoreID)
  }
  rq.lastLeaseTransfer.Store(timeutil.Now())
- if qpsMeasurementDur >= MinStatsDuration {
- rq.allocator.storePool.updateLocalStoresAfterLeaseTransfer(
- repl.store.StoreID(), target.StoreID, avgQPS)
- }
+ rq.allocator.storePool.updateLocalStoresAfterLeaseTransfer(
+ repl.store.StoreID(), target.StoreID, rangeQPS)
  return nil
 }
 

pkg/storage/store_rebalancer.go

@@ -253,7 +253,9 @@ func (sr *StoreRebalancer) rebalanceStore(
  log.VEventf(ctx, 1, "transferring r%d (%.2f qps) to s%d to better balance load",
  replWithStats.repl.RangeID, replWithStats.qps, target.StoreID)
  replCtx, cancel := context.WithTimeout(replWithStats.repl.AnnotateCtx(ctx), sr.rq.processTimeout)
- if err := sr.rq.transferLease(replCtx, replWithStats.repl, target); err != nil {
+ if err := sr.rq.transferLease(
+ replCtx, replWithStats.repl, target, replWithStats.qps,
+ ); err != nil {
  cancel()
  log.Errorf(replCtx, "unable to transfer lease to s%d: %v", target.StoreID, err)
  continue
@@ -551,6 +553,8 @@ func (sr *StoreRebalancer) chooseReplicaToRebalance(
 
  // Then pick out which new stores to add the remaining replicas to.
  rangeInfo := rangeInfoForRepl(replWithStats.repl, desc)
+ // Make sure to use the same qps measurement throughout everything we do.
+ rangeInfo.QueriesPerSecond = replWithStats.qps
  options := sr.rq.allocator.scorerOptions()
  options.qpsRebalanceThreshold = qpsRebalanceThreshold.Get(&sr.st.SV)
  for len(targets) < desiredReplicas {