This approach seems very reasonable.

Reviewed 3 of 3 files at r1, all commit messages.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @aayushshah15)

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

	ctx context.Context,
	conf roachpb.SpanConfig,
	leaseRepl *Replica,

It would be nice to test this behavior. To do that, you might want to replace the *Replica type with an interface. Even an inline interface like we have in findRemoveVoter would be enough to make this easy to test in unit tests.

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

	ctx context.Context,
	conf roachpb.SpanConfig,
	leaseRepl *Replica,

nit: move this next to leaseStoreID.

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

	ctx context.Context,
	conf roachpb.SpanConfig,
	leaseRepl *Replica,

nit: add an // optional comment.

pkg/kv/kvserver/allocator.go, line 1787 at r1 (raw file):

		// a snapshot iff it is in `StateReplicate`. However, even this is racey
		// because we can still possibly have an ill-timed log truncation between
		// when we make this determination and when we act on it.

Could you add to this comment why we don't make this as strict as replicaIsBehind? It has to do with needing to transfer leases to lagging followers in some cases, right? Didn't we have an issue about that at one point that led to the removal of this protection?

Edit: here's the commit that added the protection 83128c9, and here's the commit that removed the protection 371738c.

pkg/kv/kvserver/allocator.go, line 1802 at r1 (raw file):

	raftStatus *raft.Status, replicas []roachpb.ReplicaDescriptor,
) (result []roachpb.ReplicaDescriptor) {
	result = make([]roachpb.ReplicaDescriptor, 0, len(replicas))

nit: Does this need to re-allocate? Can't we re-use the provided slice?

pkg/kv/kvserver/allocator.go, line 1807 at r1 (raw file):

		if !replicaMayNeedSnapshot(raftStatus, repl.ReplicaID) {
			result = append(result, repl)
		}

Do you think it's worth logging on an else branch when we filter out a candidate?

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @nvanbenschoten)

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

thanks for the suggestion, done.

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

done.

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

done.

pkg/kv/kvserver/allocator.go, line 1787 at r1 (raw file):

Yeah, I found it a bit odd how we removed the protection from the replicateQueue but did not from the StoreRebalancer. The only rationale I can think of is that we'd likely hurt the range's QPS in the short term by transferring it to a lagging replica.

pkg/kv/kvserver/allocator.go, line 1802 at r1 (raw file):

done.

pkg/kv/kvserver/allocator.go, line 1807 at r1 (raw file):

done.

Reviewed 2 of 3 files at r2, 1 of 1 files at r3, all commit messages.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @aayushshah15)

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

I'm wondering whether it's worth keeping this optional just for the sake of a few tests. Allowing an input to be optional leads to more branching and harder code to reason about. And if we're going to change that, we might as well consolidate these two args to make it clear that they are in reference to the same object.

What do you think about replacing this with:

leaseRepl interface {
    StoreID() roachpb.StoreID
    RaftStatus() *raft.Status
},

Then in tests, you can pass some mock like:

type mockLeaseRepl roachpb.StoreID

func (m mockLeaseRepl) StoreID() roachpb.StoreID { return roachpb.StoreID(m) }
func (m mockLeaseRepl) RaftStatus() *raft.Status { return nil }

EDIT: or something like the mock you already added below.

pkg/kv/kvserver/allocator.go, line 1802 at r3 (raw file):

	ctx context.Context, raftStatus *raft.Status, replicas []roachpb.ReplicaDescriptor,
) []roachpb.ReplicaDescriptor {
	var filled = 0

nit: filled := 0 is slightly more idiomatic.

pkg/kv/kvserver/allocator.go, line 1814 at r3 (raw file):

			continue
		}
		replicas[filled] = replicas[i]

Why do you need i or replicas[i]? Isn't replicas[filled] = repl enough?

pkg/kv/kvserver/allocator_test.go, line 1387 at r3 (raw file):

}

type mockRepl struct {

Let's give this a comment hinting at the fact that this is a mock to satisfy the interface of TransferLeaseTarget.

pkg/kv/kvserver/allocator_test.go, line 1389 at r3 (raw file):

type mockRepl struct {
	replicationFactor     uint64
	replsInNeedOfSnapshot atomic.Value

Why the use of atomics here? Isn't the new test single-threaded? If this was just a []roachpb.ReplicaID, the code below would be simpler.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @nvanbenschoten)

pkg/kv/kvserver/allocator.go, line 1257 at r1 (raw file):

done.

pkg/kv/kvserver/allocator.go, line 1802 at r3 (raw file):

done.

pkg/kv/kvserver/allocator.go, line 1814 at r3 (raw file):

no reason, done.

pkg/kv/kvserver/allocator_test.go, line 1387 at r3 (raw file):

done.

pkg/kv/kvserver/allocator_test.go, line 1389 at r3 (raw file):

you're right, done.

but get @nvanbenschoten's take as well.

Reviewed 2 of 3 files at r4, 1 of 1 files at r5, all commit messages.
Reviewable status: complete! 1 of 0 LGTMs obtained (waiting on @aayushshah15 and @nvanbenschoten)

pkg/kv/kvserver/allocator.go, line 1320 at r5 (raw file):

	existing, _ = a.storePool.liveAndDeadReplicas(existing, false /* includeSuspectStores */)

	// Only proceed with the lease transfer if we are also the raft leader (we

Are there any meaningful failure scenarios where we may need to transfer the lease without being able to get Raft leadership? I can't think of any off the bat, but this seems like the most risky aspect of this change.

pkg/kv/kvserver/allocator.go, line 1805 at r5 (raw file):

				"not considering [n%d, s%d] as a potential candidate for a lease transfer"+
					" because the replica may be waiting for a snapshot",

Maybe we want to vary this message depending on whether we're the leader or not? I could see it being confusing to see this message when in fact the problem is that we're not the leader, and the replica is replicating just fine.

Reviewed 2 of 3 files at r4, 1 of 1 files at r5, all commit messages.
Reviewable status: complete! 2 of 0 LGTMs obtained (waiting on @aayushshah15)

pkg/kv/kvserver/allocator_test.go, line 1335 at r5 (raw file):

	}
	for i := uint64(1); i <= r.replicationFactor; i++ {
		raftStatus.Progress[i] = tracker.Progress{State: tracker.StateReplicate}

nit: this would be easier to read if you switched replsInNeedOfSnapshot to a set (map[roachpb.ReplicaID]struct{}) and wrote the logic here as:

for i := uint64(1); i <= r.replicationFactor; i++ {
    state := tracker.StateReplicate
    if _, ok := r.replsInNeedOfSnapshot[roachpb.ReplicaID(i)]; ok {
        state = tracker.StateSnapshot
    }
    raftStatus.Progress[i] = tracker.Progress{State: state}
}

A linear flow like this avoids mutability, which makes it easier to reason about and more immediately obvious what effect the presence of a replica in the replsInNeedOfSnapshot has.

Reviewable status: complete! 0 of 0 LGTMs obtained (and 2 stale) (waiting on @erikgrinaker and @nvanbenschoten)

pkg/kv/kvserver/allocator.go, line 1320 at r5 (raw file):

maybeTransferRaftLeadershipToLeaseholderLocked() sends a MsgTransferLeader to the range's raft group (and is called on every raft tick), so my understanding is that if we're persistently in a scenario where leaseholdership and leadership are uncollocated, it must imply that the underlying raft group is actually wedged. In such cases, the range should already be kaput and any lease transfer should also just fail / get stuck.

Does this line up with your / @nvanbenschoten's intuition?

pkg/kv/kvserver/allocator.go, line 1805 at r5 (raw file):

Done.

pkg/kv/kvserver/allocator_test.go, line 1335 at r5 (raw file):

done.

Reviewed 1 of 2 files at r6, all commit messages.
Reviewable status: complete! 0 of 0 LGTMs obtained (and 2 stale) (waiting on @erikgrinaker and @nvanbenschoten)

pkg/kv/kvserver/allocator.go, line 1320 at r5 (raw file):

Yes, this matches my intuition. maybeTransferRaftLeadershipToLeaseholderLocked does require the leaseholder to be caught up on the Raft log to be able to assume leadership, but I can't see that causing the kinds of problems we saw in 371738c, because the leaseholder is the only replica able to propose (a serious amount of) writes. So it seems unlikely, even under sustained load, that it would be persistently behind on the log.

But I think it's worth posing this question to the #kv slack channel, just in case anyone spots anything that we're missing. Basically just whether anyone notices any undesirable consequences of creating this dependency and preventing follower replicas that hold a range lease from transferring that lease away.

Reviewed 2 of 2 files at r6, all commit messages.
Reviewable status: complete! 0 of 0 LGTMs obtained (and 2 stale) (waiting on @aayushshah15 and @nvanbenschoten)

pkg/kv/kvserver/allocator.go, line 1320 at r5 (raw file):

Yes, that's how I understand it as well, and I can't think of any meaningful exceptions. This is probably the sort of problem we'll have to rely on randomized tests to catch, but worth asking around as @nvanbenschoten suggests.