From f2c3e12c6363ef4cdc929fff5023f85802699b39 Mon Sep 17 00:00:00 2001
From: Austen McClernon <austen@cockroachlabs.com>
Date: Tue, 20 Dec 2022 18:31:22 +0000
Subject: [PATCH] kvserver: add shed lease target to repl queue

Previously, a call to `shedLease` was made within the process loop and
outside of planning, in the replicate queue. This patch moves the shed
lease into consider rebalance, where it originally was and converts it
into a plannable action.

Release note: None
---
 pkg/kv/kvserver/replicate_queue.go | 82 +++++++++++++++++++-----------
 1 file changed, 53 insertions(+), 29 deletions(-)

diff --git a/pkg/kv/kvserver/replicate_queue.go b/pkg/kv/kvserver/replicate_queue.go
index 8e8f03c3f4bc..0a10981125c9 100644
--- a/pkg/kv/kvserver/replicate_queue.go
+++ b/pkg/kv/kvserver/replicate_queue.go
@@ -736,34 +736,6 @@ func (rq *replicateQueue) process(
 			}
 		}
 
-		// After we made a replica change, make sure the lease is still on the
-		// correct store.
-		//
-		// TODO(kvoli): This lease transfer occurs outside of the planning
-		// ProcessOneChange planning and application logic. It should be moved
-		// into the planning phase and returned as a follow up change.
-		if rq.canTransferLeaseFrom(ctx, repl) {
-			transferStatus, err := rq.shedLease(
-				ctx,
-				repl,
-				repl.Desc(),
-				repl.SpanConfig(),
-				allocator.TransferLeaseOptions{
-					Goal:                   allocator.FollowTheWorkload,
-					ExcludeLeaseRepl:       false,
-					CheckCandidateFullness: true,
-				},
-			)
-			if err != nil {
-				return false, err
-			}
-			// If we successfully transferred the lease, we can't requeue, let the new
-			// leaseholder do it.
-			if transferStatus == allocator.TransferOK {
-				requeue = false
-			}
-		}
-
 		if requeue {
 			log.KvDistribution.VEventf(ctx, 1, "re-processing")
 			rq.maybeAdd(ctx, repl, rq.store.Clock().NowAsClockTimestamp())
@@ -1823,8 +1795,24 @@ func (rq *replicateQueue) considerRebalance(
 		}
 	}
 
+	// No rebalance target was found, check whether we are able and should
+	// transfer the lease away to another store.
 	if !ok {
-		return nil, nil
+		if !canTransferLeaseFrom(ctx, repl) {
+			return nil, nil
+		}
+		return rq.shedLeaseTarget(
+			ctx,
+			repl,
+			desc,
+			conf,
+			allocator.TransferLeaseOptions{
+				Goal:                   allocator.FollowTheWorkload,
+				ExcludeLeaseRepl:       false,
+				CheckCandidateFullness: true,
+			},
+		), nil
+
 	}
 
 	// If we have a valid rebalance action (ok == true) and we haven't
@@ -1962,6 +1950,42 @@ func replicationChangesForRebalance(
 	return chgs, performingSwap, nil
 }
 
+// shedLeaseTarget takes in a leaseholder replica, looks for a target for
+// transferring the lease and, if a suitable target is found (e.g. alive, not
+// draining), returns an allocation op to transfer the lease away.
+func (rq *replicateQueue) shedLeaseTarget(
+	ctx context.Context,
+	repl *Replica,
+	desc *roachpb.RangeDescriptor,
+	conf roachpb.SpanConfig,
+	opts allocator.TransferLeaseOptions,
+) (op AllocationOp) {
+	usage := RangeUsageInfoForRepl(repl)
+	// Learner replicas aren't allowed to become the leaseholder or raft leader,
+	// so only consider the `VoterDescriptors` replicas.
+	target := rq.allocator.TransferLeaseTarget(
+		ctx,
+		rq.storePool,
+		conf,
+		desc.Replicas().VoterDescriptors(),
+		repl,
+		usage,
+		false, /* forceDecisionWithoutStats */
+		opts,
+	)
+	if target == (roachpb.ReplicaDescriptor{}) {
+		return nil
+	}
+
+	op = AllocationTransferLeaseOp{
+		source:             repl.StoreID(),
+		target:             target.StoreID,
+		usage:              usage,
+		bypassSafetyChecks: false,
+	}
+	return op
+}
+
 // shedLease takes in a leaseholder replica, looks for a target for transferring
 // the lease and, if a suitable target is found (e.g. alive, not draining),
 // transfers the lease away.