Merge cockroachdb#124223

124223: kvserver: read lease under mutex when switching lease type r=nvanbenschoten a=kvoli A race could occur when a replica queue and post lease application both attempted to switch the lease type. This race would cause the queue to not process the replica because the lease type had already changed. As a result, lease preference violations might not have been quickly resolved by the lease queue. Read the lease under the same mutex used for requesting the lease, when possibly switching the lease type. Resolves: cockroachdb#123998 Release note: None Co-authored-by: Austen McClernon <[email protected]>
miyamo2 · May 17, 2024 · 3732a92 · 3732a92
2 parents 91fac21 + 574f667
commit 3732a92
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 12 deletions.
diff --git a/pkg/kv/kvserver/queue.go b/pkg/kv/kvserver/queue.go
@@ -1058,7 +1058,7 @@ func (bq *baseQueue) replicaCanBeProcessed(
 	// and renew or acquire if necessary.
 	if bq.needsLease {
 		if acquireLeaseIfNeeded {
-			leaseStatus, pErr := repl.redirectOnOrAcquireLease(ctx)
+			_, pErr := repl.redirectOnOrAcquireLease(ctx)
 			if pErr != nil {
 				switch v := pErr.GetDetail().(type) {
 				case *kvpb.NotLeaseHolderError, *kvpb.RangeNotFoundError:
@@ -1071,7 +1071,7 @@ func (bq *baseQueue) replicaCanBeProcessed(
 
 			// TODO(baptist): Should this be added to replicaInQueue?
 			realRepl, _ := repl.(*Replica)
-			pErr = realRepl.maybeSwitchLeaseType(ctx, leaseStatus)
+			pErr = realRepl.maybeSwitchLeaseType(ctx)
 			if pErr != nil {
 				return nil, pErr.GoError()
 			}

diff --git a/pkg/kv/kvserver/replica_range_lease.go b/pkg/kv/kvserver/replica_range_lease.go
@@ -1610,17 +1610,24 @@ func (r *Replica) shouldRequestLeaseRLocked(
 // maybeSwitchLeaseType will synchronously renew a lease using the appropriate
 // type if it is (or was) owned by this replica and has an incorrect type. This
 // typically happens when changing kv.expiration_leases_only.enabled.
-func (r *Replica) maybeSwitchLeaseType(ctx context.Context, st kvserverpb.LeaseStatus) *kvpb.Error {
-	if !st.OwnedBy(r.store.StoreID()) {
-		return nil
-	}
+func (r *Replica) maybeSwitchLeaseType(ctx context.Context) *kvpb.Error {
+	llHandle := func() *leaseRequestHandle {
+		now := r.store.Clock().NowAsClockTimestamp()
+		// The lease status needs to be checked and requested under the same lock,
+		// to avoid an interleaving lease request changing the lease between the
+		// two.
+		r.mu.Lock()
+		defer r.mu.Unlock()
 
-	var llHandle *leaseRequestHandle
-	r.mu.Lock()
-	if !r.hasCorrectLeaseTypeRLocked(st.Lease) {
-		llHandle = r.requestLeaseLocked(ctx, st, nil /* limiter */)
-	}
-	r.mu.Unlock()
+		st := r.leaseStatusAtRLocked(ctx, now)
+		if !st.OwnedBy(r.store.StoreID()) {
+			return nil
+		}
+		if r.hasCorrectLeaseTypeRLocked(st.Lease) {
+			return nil
+		}
+		return r.requestLeaseLocked(ctx, st, nil /* limiter */)
+	}()
 
 	if llHandle != nil {
 		select {