Merge #152205

152205: kvserver: rm unnecessary Replica.mu locking r=arulajmani a=pav-kv

This PR removes the unnecessary `Replica.mu.RLock()` in `Store.tryGetReplica()`. This function is on the critical path - every raft message [processing](https://github.com/cockroachdb/cockroach/blob/d1af030c8b9894306710132a253bd7c2c9145623/pkg/kv/kvserver/store_raft.go#L699) calls it. So this should have a non-zero performance impact.

Presumably, `Replica.mu` used to be locked for reading `destroyStatus`, but the latter can be read with only `raftMu` held. To emphasize that this field can be read under either mutex, this PR also moves `destroyStatus` from the `Replica.mu` to `Replica.shMu` section.

Technically, `destroyStatus` doesn't fit either of `Replica.{mu,shMu}` sections, because it must be mutated with also `readOnlyCmdMu` locked. But `shMu` is strictly closer by semantics.

Epic: none
Release note: none

Co-authored-by: Pavel Kalinnikov <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/replica.go

@@ -571,6 +571,14 @@ type Replica struct {
 	//
 	// TODO(pav-kv): audit all other fields and include here.
 	shMu struct {
+		// 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 readOnlyCmdMu.
+		//
+		// When this replica is being removed, the destroyStatus is updated and
+		// RangeTombstone is written in the same raftMu critical section.
+		destroyStatus
+
 		// The state of the Raft state machine.
 		// Invariant: state.TruncatedState == nil. The field is being phased out in
 		// favour of the one contained in logStorage.
@@ -587,14 +595,6 @@ type Replica struct {
 	mu struct {
 		// Protects all fields in the mu struct.
 		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
-		// readOnlyCmdMu.
-		//
-		// When this replica is being removed, the destroyStatus is updated and
-		// RangeTombstone is written in the same raftMu critical section.
-		destroyStatus
 		// Is the range quiescent? Quiescent ranges are not Tick()'d and unquiesce
 		// whenever a Raft operation is performed.
 		//
@@ -1250,7 +1250,7 @@ func (r *Replica) IsDestroyed() (DestroyReason, error) {
 }
 
 func (r *Replica) isDestroyedRLocked() (DestroyReason, error) {
-	return r.mu.destroyStatus.reason, r.mu.destroyStatus.err
+	return r.shMu.destroyStatus.reason, r.shMu.destroyStatus.err
 }
 
 // IsQuiescent returns whether the replica is quiescent or not.
@@ -2672,11 +2672,11 @@ func (r *Replica) maybeWatchForMergeLocked(ctx context.Context) (bool, error) {
 		r.raftMu.Lock()
 		r.readOnlyCmdMu.Lock()
 		r.mu.Lock()
-		if mergeCommitted && r.mu.destroyStatus.IsAlive() {
+		if mergeCommitted && r.shMu.destroyStatus.IsAlive() {
 			// The merge committed but the left-hand replica on this store hasn't
 			// subsumed this replica yet. Mark this replica as destroyed so it
 			// doesn't serve requests when we close the mergeCompleteCh below.
-			r.mu.destroyStatus.Set(kvpb.NewRangeNotFoundError(r.RangeID, r.store.StoreID()), destroyReasonMergePending)
+			r.shMu.destroyStatus.Set(kvpb.NewRangeNotFoundError(r.RangeID, r.store.StoreID()), destroyReasonMergePending)
 		}
 		// Unblock pending requests. If the merge committed, the requests will
 		// notice that the replica has been destroyed and return an appropriate

pkg/kv/kvserver/replica_app_batch.go

@@ -353,7 +353,7 @@ func (b *replicaAppBatch) runPostAddTriggersReplicaOnly(
 		// commits by handleMergeResult() to finish the removal.
 		rhsRepl.readOnlyCmdMu.Lock()
 		rhsRepl.mu.Lock()
-		rhsRepl.mu.destroyStatus.Set(
+		rhsRepl.shMu.destroyStatus.Set(
 			kvpb.NewRangeNotFoundError(rhsRepl.RangeID, rhsRepl.store.StoreID()),
 			destroyReasonRemoved)
 		rhsRepl.mu.Unlock()
@@ -443,7 +443,7 @@ func (b *replicaAppBatch) runPostAddTriggersReplicaOnly(
 		// application.
 		b.r.readOnlyCmdMu.Lock()
 		b.r.mu.Lock()
-		b.r.mu.destroyStatus.Set(
+		b.r.shMu.destroyStatus.Set(
 			kvpb.NewRangeNotFoundError(b.r.RangeID, b.r.store.StoreID()),
 			destroyReasonRemoved)
 		span := b.r.descRLocked().RSpan()

pkg/kv/kvserver/replica_application_state_machine_test.go

@@ -149,9 +149,11 @@ func TestReplicaStateMachineChangeReplicas(t *testing.T) {
 		// should there be a command in the raft log (i.e. some errant lease request
 		// or whatnot) this will fire assertions because it will conflict with the
 		// log index that we pulled out of thin air above.
+		r.readOnlyCmdMu.Lock()
+		defer r.readOnlyCmdMu.Unlock()
 		r.mu.Lock()
 		defer r.mu.Unlock()
-		r.mu.destroyStatus.Set(errors.New("test done"), destroyReasonRemoved)
+		r.shMu.destroyStatus.Set(errors.New("test done"), destroyReasonRemoved)
 	})
 }
 
@@ -231,14 +233,16 @@ func TestReplicaStateMachineRaftLogTruncationStronglyCoupled(t *testing.T) {
 		_, err = sm.ApplySideEffects(checkedCmd.Ctx(), checkedCmd)
 		require.NoError(t, err)
 		func() {
-			r.mu.Lock()
-			defer r.mu.Unlock()
 			// Set a destroyStatus to make sure there won't be any raft processing once
 			// we release raftMu. We applied a command but not one from the raft log, so
 			// should there be a command in the raft log (i.e. some errant lease request
 			// or whatnot) this will fire assertions because it will conflict with the
 			// log index that we pulled out of thin air above.
-			r.mu.destroyStatus.Set(errors.New("test done"), destroyReasonRemoved)
+			r.readOnlyCmdMu.Lock()
+			r.mu.Lock()
+			defer r.mu.Unlock()
+			r.shMu.destroyStatus.Set(errors.New("test done"), destroyReasonRemoved)
+			r.readOnlyCmdMu.Unlock()
 
 			require.Equal(t, raftAppliedIndex+1, r.shMu.state.RaftAppliedIndex)
 			require.Equal(t, truncatedIndex+1, ls.shMu.trunc.Index)
@@ -427,7 +431,13 @@ func TestReplicaStateMachineEphemeralAppBatchRejection(t *testing.T) {
 	defer r.raftMu.Unlock()
 	// Avoid additional raft processing after we're done with this replica because
 	// we've applied entries that aren't in the log.
-	defer r.mu.destroyStatus.Set(errors.New("boom"), destroyReasonRemoved)
+	defer func() {
+		r.readOnlyCmdMu.Lock()
+		defer r.readOnlyCmdMu.Unlock()
+		r.mu.Lock()
+		defer r.mu.Unlock()
+		r.shMu.destroyStatus.Set(errors.New("boom"), destroyReasonRemoved)
+	}()
 
 	sm := r.getStateMachine()
 

pkg/kv/kvserver/replica_command_test.go

@@ -677,9 +677,13 @@ func TestWaitForLeaseAppliedIndex(t *testing.T) {
 	stopper.Stop(ctx)
 
 	destroyErr := errors.New("destroyed")
-	tc.repl.mu.Lock()
-	tc.repl.mu.destroyStatus.Set(destroyErr, destroyReasonRemoved)
-	tc.repl.mu.Unlock()
+	func() {
+		tc.repl.readOnlyCmdMu.Lock()
+		defer tc.repl.readOnlyCmdMu.Unlock()
+		tc.repl.mu.Lock()
+		defer tc.repl.mu.Unlock()
+		tc.repl.shMu.destroyStatus.Set(destroyErr, destroyReasonRemoved)
+	}()
 
 	_, err = tc.repl.WaitForLeaseAppliedIndex(ctx, maxLAI)
 	require.Error(t, err)

pkg/kv/kvserver/replica_corruption.go

@@ -43,7 +43,7 @@ func (r *Replica) setCorruptRaftMuLocked(
 
 	log.Dev.ErrorfDepth(ctx, 1, "stalling replica due to: %s", cErr.ErrorMsg)
 	cErr.Processed = true
-	r.mu.destroyStatus.Set(cErr, destroyReasonRemoved)
+	r.shMu.destroyStatus.Set(cErr, destroyReasonRemoved)
 
 	auxDir := r.store.TODOEngine().GetAuxiliaryDir()
 	_ = r.store.TODOEngine().Env().MkdirAll(auxDir, os.ModePerm)

pkg/kv/kvserver/replica_destroy.go

@@ -167,7 +167,7 @@ func (r *Replica) disconnectReplicationRaftMuLocked(ctx context.Context) {
 		p.finishApplication(ctx, makeProposalResultErr(kvpb.NewAmbiguousResultError(apply.ErrRemoved)))
 	}
 
-	if !r.mu.destroyStatus.Removed() {
+	if !r.shMu.destroyStatus.Removed() {
 		log.Dev.Fatalf(ctx, "removing raft group before destroying replica %s", r)
 	}
 	r.mu.internalRaftGroup = nil

pkg/kv/kvserver/replica_proposal_buf.go

@@ -1116,7 +1116,7 @@ func (rp *replicaProposer) getReplicaID() roachpb.ReplicaID {
 }
 
 func (rp *replicaProposer) destroyed() destroyStatus {
-	return rp.mu.destroyStatus
+	return rp.shMu.destroyStatus
 }
 
 func (rp *replicaProposer) leaseAppliedIndex() kvpb.LeaseAppliedIndex {

pkg/kv/kvserver/replica_raft.go

@@ -1648,7 +1648,7 @@ func (r *Replica) refreshProposalsLocked(
 	}
 
 	r.mu.slowProposalCount = slowProposalCount
-	destroyed := r.mu.destroyStatus.Removed()
+	destroyed := r.shMu.destroyStatus.Removed()
 
 	// If the breaker isn't tripped yet but we've detected commands that have
 	// taken too long to replicate, and the replica is not destroyed, trip the
@@ -2222,7 +2222,7 @@ func (s pendingCmdSlice) Less(i, j int) bool {
 func (r *Replica) withRaftGroupLocked(
 	f func(r *raft.RawNode) (unquiesceAndWakeLeader bool, _ error),
 ) error {
-	if r.mu.destroyStatus.Removed() {
+	if r.shMu.destroyStatus.Removed() {
 		// Callers know to detect errRemoved as non-fatal.
 		return errRemoved
 	}

pkg/kv/kvserver/replica_raftstorage.go

@@ -577,7 +577,7 @@ func (r *Replica) applySnapshotRaftMuLocked(
 		// erroneously return empty data.
 		sr.readOnlyCmdMu.Lock()
 		sr.mu.Lock()
-		sr.mu.destroyStatus.Set(
+		sr.shMu.destroyStatus.Set(
 			kvpb.NewRangeNotFoundError(sr.RangeID, sr.store.StoreID()),
 			destroyReasonRemoved)
 		sr.mu.Unlock()

pkg/kv/kvserver/replica_test.go

@@ -13781,9 +13781,13 @@ func TestAdminScatterDestroyedReplica(t *testing.T) {
 	tc.Start(ctx, t, stopper)
 
 	errBoom := errors.New("boom")
-	tc.repl.mu.Lock()
-	tc.repl.mu.destroyStatus.Set(errBoom, destroyReasonMergePending)
-	tc.repl.mu.Unlock()
+	func() {
+		tc.repl.readOnlyCmdMu.Lock()
+		defer tc.repl.readOnlyCmdMu.Unlock()
+		tc.repl.mu.Lock()
+		defer tc.repl.mu.Unlock()
+		tc.repl.shMu.destroyStatus.Set(errBoom, destroyReasonMergePending)
+	}()
 
 	desc := tc.repl.Desc()
 	resp, err := tc.repl.adminScatter(ctx, kvpb.AdminScatterRequest{