Merge cockroachdb#106574

106574: kvserver: avoid leaked replica mutex during delegated snapshot r=erikgrinaker a=tbg

Prompted by cockroachdb#106568, I went
through all callers to `repl.mu.Lock` and `repl.mu.RLock` trying to find places
where we're leaking the mutex on certain return paths.

I found some, including one that at least seems to possibly explain what we saw in cockroachdb#106568.

I also looked into the possibility of panics under a held lock being recovered from
at higher levels in the stack, in particular this code:

https://github.com/cockroachdb/cockroach/blob/81569be4aeb61620d2fd51c41b416f7d0986698e/pkg/sql/colexecerror/error.go#L27-L30

However, it seems sufficiently restricted to recovering only from errors it knows to
be safe to recover from, and in particular it doesn't look like it will ever recover from
any panic in kvserver. I will say that this looks a little brittle, though I don't see a much
better way of doing what this code is trying to achieve.

Release note: None
Epic: None

Co-authored-by: Tobias Grieger <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/replica.go

@@ -181,6 +181,34 @@ type raftSparseStatus struct {
 	Progress map[uint64]tracker.Progress
 }
 
+// ReplicaMutex is an RWMutex. It has its own type to make it easier to look for
+// usages specific to the replica mutex.
+type ReplicaMutex syncutil.RWMutex
+
+func (mu *ReplicaMutex) Lock() {
+	(*syncutil.RWMutex)(mu).Lock()
+}
+
+func (mu *ReplicaMutex) Unlock() {
+	(*syncutil.RWMutex)(mu).Unlock()
+}
+
+func (mu *ReplicaMutex) RLock() {
+	(*syncutil.RWMutex)(mu).RLock()
+}
+
+func (mu *ReplicaMutex) AssertHeld() {
+	(*syncutil.RWMutex)(mu).AssertHeld()
+}
+
+func (mu *ReplicaMutex) AssertRHeld() {
+	(*syncutil.RWMutex)(mu).AssertRHeld()
+}
+
+func (mu *ReplicaMutex) RUnlock() {
+	(*syncutil.RWMutex)(mu).RUnlock()
+}
+
 // A Replica is a contiguous keyspace with writes managed via an
 // instance of the Raft consensus algorithm. Many ranges may exist
 // in a store and they are unlikely to be contiguous. Ranges are
@@ -386,7 +414,7 @@ type Replica struct {
 
 	mu struct {
 		// Protects all fields in the mu struct.
-		syncutil.RWMutex
+		ReplicaMutex
 		// The destroyed status of a replica indicating if it's alive, corrupt,
 		// scheduled for destruction or has been GCed.
 		// destroyStatus should only be set while also holding the raftMu and

pkg/kv/kvserver/replica_app_batch.go

@@ -580,6 +580,7 @@ func (b *replicaAppBatch) ApplyToStateMachine(ctx context.Context) error {
 	existingClosed := r.mu.state.RaftClosedTimestamp
 	newClosed := b.state.RaftClosedTimestamp
 	if !newClosed.IsEmpty() && newClosed.Less(existingClosed) && raftClosedTimestampAssertionsEnabled {
+		r.mu.Unlock()
 		return errors.AssertionFailedf(
 			"raft closed timestamp regression; replica has: %s, new batch has: %s.",
 			existingClosed.String(), newClosed.String())

pkg/kv/kvserver/replica_command.go

@@ -2953,15 +2953,15 @@ func (r *Replica) validateSnapshotDelegationRequest(
 	// that is also needs a snapshot, then any snapshot it sends will be useless.
 	r.mu.RLock()
 	replIdx := r.mu.state.RaftAppliedIndex + 1
-
 	status := r.raftStatusRLocked()
+	r.mu.RUnlock()
+
 	if status == nil {
 		// This code path is sometimes hit during scatter for replicas that
 		// haven't woken up yet.
 		return errors.Errorf("raft status not initialized")
 	}
 	replTerm := kvpb.RaftTerm(status.Term)
-	r.mu.RUnlock()
 
 	// Delegate has a lower term than the coordinator. This typically means the
 	// lease has been transferred, and we should not process this request. There

pkg/kv/kvserver/replica_proposal_buf.go

@@ -1277,11 +1277,11 @@ type replicaProposer Replica
 var _ proposer = &replicaProposer{}
 
 func (rp *replicaProposer) locker() sync.Locker {
-	return &rp.mu.RWMutex
+	return &rp.mu.ReplicaMutex
 }
 
 func (rp *replicaProposer) rlocker() sync.Locker {
-	return rp.mu.RWMutex.RLocker()
+	return &rp.mu.ReplicaMutex
 }
 
 func (rp *replicaProposer) getReplicaID() roachpb.ReplicaID {

pkg/kv/kvserver/replica_raftstorage.go

@@ -642,8 +642,10 @@ func (r *Replica) applySnapshot(
 	if isInitialSnap {
 		// NB: this will also call setDescLockedRaftMuLocked.
 		if err := r.initFromSnapshotLockedRaftMuLocked(ctx, desc); err != nil {
+			r.mu.Unlock()
 			log.Fatalf(ctx, "unable to initialize replica while applying snapshot: %+v", err)
 		} else if err := r.store.markReplicaInitializedLockedReplLocked(ctx, r); err != nil {
+			r.mu.Unlock()
 			log.Fatalf(ctx, "unable to mark replica initialized while applying snapshot: %+v", err)
 		}
 	} else {

pkg/kv/kvserver/stores_base.go

@@ -99,7 +99,7 @@ func (s *baseStore) SetQueueActive(active bool, queue string) error {
 func (s *baseStore) GetReplicaMutexForTesting(rangeID roachpb.RangeID) *syncutil.RWMutex {
 	store := (*Store)(s)
 	if repl := store.GetReplicaIfExists(rangeID); repl != nil {
-		return &repl.mu.RWMutex
+		return (*syncutil.RWMutex)(&repl.mu.ReplicaMutex)
 	}
 	return nil
 }