diff --git a/pkg/kv/kvserver/client_replica_test.go b/pkg/kv/kvserver/client_replica_test.go index 4737cd86676f..befceeb18488 100644 --- a/pkg/kv/kvserver/client_replica_test.go +++ b/pkg/kv/kvserver/client_replica_test.go @@ -2933,6 +2933,9 @@ func TestLeaseTransferRejectedIfTargetNeedsSnapshot(t *testing.T) { ServerArgs: base.TestServerArgs{ Knobs: base.TestingKnobs{ Store: &kvserver.StoreTestingKnobs{ + // If we're testing the below-raft check, disable the above-raft check. + // See: https://github.com/cockroachdb/cockroach/pull/107526 + DisableAboveRaftLeaseTransferSafetyChecks: rejectAfterRevoke, TestingRequestFilter: func(ctx context.Context, ba *kvpb.BatchRequest) *kvpb.Error { if rejectAfterRevoke && ba.IsSingleTransferLeaseRequest() { transferLeaseReqBlockOnce.Do(func() { @@ -3003,10 +3006,17 @@ func TestLeaseTransferRejectedIfTargetNeedsSnapshot(t *testing.T) { // Replica.AdminTransferLease. transferErrC := make(chan error, 1) if rejectAfterRevoke { - _ = tc.Stopper().RunAsyncTask(ctx, "transfer lease", func(ctx context.Context) { + require.NoError(t, tc.Stopper().RunAsyncTask(ctx, "transfer lease", func(ctx context.Context) { transferErrC <- tc.TransferRangeLease(*repl0.Desc(), tc.Target(2)) - }) - <-transferLeaseReqBlockedC + })) + select { + case <-transferLeaseReqBlockedC: + // Expected case: lease transfer triggered our interceptor and is now + // waiting there for transferLeaseReqUnblockedCh. + case err := <-transferErrC: + // Unexpected case: lease transfer errored out before making it into the filter. + t.Fatalf("transferErrC unexpectedly signaled: %v", err) + } } // Truncate the log at index+1 (log entries < N are removed, so this diff --git a/pkg/kv/kvserver/replica_range_lease.go b/pkg/kv/kvserver/replica_range_lease.go index bd98eca5b655..b3615eac3244 100644 --- a/pkg/kv/kvserver/replica_range_lease.go +++ b/pkg/kv/kvserver/replica_range_lease.go @@ -949,7 +949,7 @@ func (r *Replica) AdminTransferLease( raftStatus := r.raftStatusRLocked() raftFirstIndex := r.raftFirstIndexRLocked() snapStatus := raftutil.ReplicaMayNeedSnapshot(raftStatus, raftFirstIndex, nextLeaseHolder.ReplicaID) - if snapStatus != raftutil.NoSnapshotNeeded && !bypassSafetyChecks { + if snapStatus != raftutil.NoSnapshotNeeded && !bypassSafetyChecks && !r.store.cfg.TestingKnobs.DisableAboveRaftLeaseTransferSafetyChecks { r.store.metrics.LeaseTransferErrorCount.Inc(1) log.VEventf(ctx, 2, "not initiating lease transfer because the target %s may "+ "need a snapshot: %s", nextLeaseHolder, snapStatus) diff --git a/pkg/kv/kvserver/testing_knobs.go b/pkg/kv/kvserver/testing_knobs.go index d8462940073b..7cb72b992956 100644 --- a/pkg/kv/kvserver/testing_knobs.go +++ b/pkg/kv/kvserver/testing_knobs.go @@ -225,6 +225,10 @@ type StoreTestingKnobs struct { // leadership when it diverges from the range's leaseholder. This can // also be set via COCKROACH_DISABLE_LEADER_FOLLOWS_LEASEHOLDER. DisableLeaderFollowsLeaseholder bool + // If set, the above-raft lease transfer safety checks (that verify that + // we don't transfer leases to followers that need a snapshot, etc) are + // disabled. The proposal-time checks are not affected by this knob. + DisableAboveRaftLeaseTransferSafetyChecks bool // DisableRefreshReasonNewLeader disables refreshing pending commands when a new // leader is discovered. DisableRefreshReasonNewLeader bool