kvserver: consolidate raft handling assertion checks

We now

- avoid queueing multiple assertion checks in an apply cycle
- queue the snapshot-induced assertion check at the end, giving us more
  flexibility over when the state is updated (and allowing future refactoring
  that streamlines these updates)

Author: tbg

pkg/kv/kvserver/replica_application_state_machine.go

@@ -206,13 +206,8 @@ func (sm *replicaStateMachine) ApplySideEffects(
 		// Some tests (TestRangeStatsInit) assumes that once the store has started
 		// and the first range has a lease that there will not be a later hard-state.
 		if shouldAssert {
-			// Assert that the on-disk state doesn't diverge from the in-memory
+			// Queue a check that the on-disk state doesn't diverge from the in-memory
 			// state as a result of the side effects.
-			sm.r.mu.RLock()
-			// TODO(sep-raft-log): either check only statemachine invariants or
-			// pass both engines in.
-			sm.r.assertStateRaftMuLockedReplicaMuRLocked(ctx, sm.r.store.TODOEngine())
-			sm.r.mu.RUnlock()
 			sm.applyStats.assertionsRequested++
 		}
 	} else if res := cmd.ReplicatedResult(); !res.IsZero() {

pkg/kv/kvserver/replica_raft.go

@@ -1082,6 +1082,11 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 		}
 	}
 
+	// If this field is set, by the end of the method (after snapshot, append,
+	// apply handling), we will verify invariants including checking that
+	// in-memory state is congruent with disk state.
+	var shouldAssert bool
+
 	// Grab the known leaseholder before applying to the state machine.
 	startingLeaseholderID := r.shMu.state.Lease.Replica.ReplicaID
 	refreshReason := noReason
@@ -1111,6 +1116,7 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 		}
 
 		if app.Snapshot != nil {
+			shouldAssert = true
 			if inSnap.Desc == nil {
 				// If we didn't expect Raft to have a snapshot but it has one
 				// regardless, that is unexpected and indicates a programming
@@ -1236,6 +1242,7 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 			// it is now marked as destroyed.
 			return stats, err
 		}
+		shouldAssert = shouldAssert || stats.apply.assertionsRequested > 0
 
 		if r.store.cfg.KVAdmissionController != nil &&
 			stats.apply.followerStoreWriteBytes.NumEntries > 0 {
@@ -1268,6 +1275,16 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 	if r.store.TestingKnobs().EnableUnconditionalRefreshesInRaftReady {
 		refreshReason = reasonNewLeaderOrConfigChange
 	}
+
+	if shouldAssert {
+		sm.r.mu.RLock()
+		// TODO(sep-raft-log): either check only statemachine invariants or
+		// pass both engines in.
+		sm.r.assertStateRaftMuLockedReplicaMuRLocked(ctx, sm.r.store.TODOEngine())
+		sm.r.mu.RUnlock()
+
+	}
+
 	if refreshReason != noReason {
 		r.mu.Lock()
 		r.refreshProposalsLocked(ctx, 0 /* refreshAtDelta */, refreshReason)

pkg/kv/kvserver/replica_raftstorage.go

@@ -775,14 +775,6 @@ func (r *Replica) applySnapshotRaftMuLocked(
 
 	r.mu.Unlock()
 
-	// Assert that the in-memory and on-disk states of the Replica are congruent
-	// after the application of the snapshot. Do so under a read lock, as this
-	// operation can be expensive. This is safe, as we hold the Replica.raftMu
-	// across both Replica.mu critical sections.
-	r.mu.RLock()
-	r.assertStateRaftMuLockedReplicaMuRLocked(ctx, r.store.TODOEngine())
-	r.mu.RUnlock()
-
 	// The rangefeed processor is listening for the logical ops attached to
 	// each raft command. These will be lost during a snapshot, so disconnect
 	// the rangefeed, if one exists.