Merge #158280

158280: kvserver: split state/raft engines in testContext r=arulajmani a=pav-kv

This PR makes tests using `testContext` aware of the separated engines.

Part of #97627

Co-authored-by: Pavel Kalinnikov <pavel@cockroachlabs.com>

Author: craig[bot]

pkg/kv/kvserver/client_merge_test.go

@@ -5499,7 +5499,7 @@ func TestStoreMergeGCHint(t *testing.T) {
 				require.Greater(t, gcHint.LatestRangeDeleteTimestamp.WallTime, beforeSecondDel,
 					"highest timestamp wasn't picked up")
 			}
-			repl.AssertState(ctx, store.TODOEngine())
+			repl.AssertState(ctx, store.StateEngine(), store.LogEngine())
 		})
 	}
 }

pkg/kv/kvserver/client_split_test.go

@@ -2389,7 +2389,7 @@ func TestStoreSplitGCThreshold(t *testing.T) {
 		t.Fatalf("expected RHS's GCThreshold is equal to %v, but got %v", specifiedGCThreshold, gcThreshold)
 	}
 
-	repl.AssertState(ctx, store.TODOEngine())
+	repl.AssertState(ctx, store.StateEngine(), store.LogEngine())
 }
 
 func TestStoreSplitGCHint(t *testing.T) {
@@ -2451,7 +2451,7 @@ func TestStoreSplitGCHint(t *testing.T) {
 	gcHint = repl.GetGCHint()
 	require.False(t, gcHint.IsEmpty(), "GC hint is empty after range delete")
 
-	repl.AssertState(ctx, store.TODOEngine())
+	repl.AssertState(ctx, store.StateEngine(), store.LogEngine())
 }
 
 // TestStoreRangeSplitRaceUninitializedRHS reproduces #7600 (before it was

pkg/kv/kvserver/helpers_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval/result"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/intentresolver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
@@ -318,12 +319,14 @@ func (r *Replica) Breaker() *circuit.Breaker {
 	return r.breaker.wrapped
 }
 
-func (r *Replica) AssertState(ctx context.Context, reader storage.Reader) {
+func (r *Replica) AssertState(
+	ctx context.Context, stateRO kvstorage.StateRO, raftRO kvstorage.RaftRO,
+) {
 	r.raftMu.Lock()
 	defer r.raftMu.Unlock()
 	r.mu.RLock()
 	defer r.mu.RUnlock()
-	r.assertStateRaftMuLockedReplicaMuRLocked(ctx, reader)
+	r.assertStateRaftMuLockedReplicaMuRLocked(ctx, stateRO, raftRO)
 }
 
 func (r *Replica) RaftLock() {

pkg/kv/kvserver/mvcc_gc_queue_test.go

@@ -1145,14 +1145,12 @@ func TestMVCCGCQueueTransactionTable(t *testing.T) {
 		txns[strKey] = *txn
 		for _, addrKey := range []roachpb.Key{baseKey, outsideKey} {
 			key := keys.TransactionKey(addrKey, txn.ID)
-			if err := storage.MVCCPutProto(ctx, tc.engine, key, hlc.Timestamp{}, txn, storage.MVCCWriteOptions{}); err != nil {
-				t.Fatal(err)
-			}
+			require.NoError(t, storage.MVCCPutProto(
+				ctx, tc.stateEng, key, hlc.Timestamp{}, txn, storage.MVCCWriteOptions{},
+			))
 		}
 		entry := roachpb.AbortSpanEntry{Key: txn.Key, Timestamp: txn.LastActive()}
-		if err := tc.repl.abortSpan.Put(ctx, tc.engine, nil, txn.ID, &entry); err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, tc.repl.abortSpan.Put(ctx, tc.stateEng, nil, txn.ID, &entry))
 	}
 
 	// Run GC.
@@ -1174,7 +1172,7 @@ func TestMVCCGCQueueTransactionTable(t *testing.T) {
 			txnKey := keys.TransactionKey(roachpb.Key(strKey), txns[strKey].ID)
 			txnTombstoneTSCacheKey := transactionTombstoneMarker(
 				roachpb.Key(strKey), txns[strKey].ID)
-			ok, err := storage.MVCCGetProto(ctx, tc.engine, txnKey, hlc.Timestamp{}, txn,
+			ok, err := storage.MVCCGetProto(ctx, tc.stateEng, txnKey, hlc.Timestamp{}, txn,
 				storage.MVCCGetOptions{})
 			if err != nil {
 				return err
@@ -1238,11 +1236,9 @@ func TestMVCCGCQueueTransactionTable(t *testing.T) {
 			"but only %d are left", exp, count)
 	}
 
-	batch := tc.engine.NewSnapshot()
-	defer batch.Close()
 	tc.repl.raftMu.Lock()
 	tc.repl.mu.RLock()
-	tc.repl.assertStateRaftMuLockedReplicaMuRLocked(ctx, batch) // check that in-mem and on-disk state were updated
+	tc.repl.assertStateRaftMuLockedReplicaMuRLocked(ctx, tc.stateEng, tc.raftEng) // check that in-mem and on-disk state were updated
 	tc.repl.mu.RUnlock()
 	tc.repl.raftMu.Unlock()
 }
@@ -1349,9 +1345,9 @@ func TestMVCCGCQueueLastProcessedTimestamps(t *testing.T) {
 
 	ts := tc.Clock().Now()
 	for _, lpv := range lastProcessedVals {
-		if err := storage.MVCCPutProto(ctx, tc.engine, lpv.key, hlc.Timestamp{}, &ts, storage.MVCCWriteOptions{}); err != nil {
-			t.Fatal(err)
-		}
+		require.NoError(t, storage.MVCCPutProto(
+			ctx, tc.stateEng, lpv.key, hlc.Timestamp{}, &ts, storage.MVCCWriteOptions{},
+		))
 	}
 
 	confReader, err := tc.store.GetConfReader(ctx)
@@ -1370,7 +1366,7 @@ func TestMVCCGCQueueLastProcessedTimestamps(t *testing.T) {
 	// Verify GC.
 	testutils.SucceedsSoon(t, func() error {
 		for _, lpv := range lastProcessedVals {
-			ok, err := storage.MVCCGetProto(ctx, tc.engine, lpv.key, hlc.Timestamp{}, &ts,
+			ok, err := storage.MVCCGetProto(ctx, tc.stateEng, lpv.key, hlc.Timestamp{}, &ts,
 				storage.MVCCGetOptions{})
 			if err != nil {
 				return err

pkg/kv/kvserver/replica.go

@@ -1973,18 +1973,18 @@ func (r *Replica) State(ctx context.Context) kvserverpb.RangeInfo {
 // to check that the in-memory and on-disk states of the Replica are congruent.
 // Requires that r.raftMu is locked and r.mu is read locked.
 func (r *Replica) assertStateRaftMuLockedReplicaMuRLocked(
-	ctx context.Context, reader storage.Reader,
+	ctx context.Context, stateRO kvstorage.StateRO, raftRO kvstorage.RaftRO,
 ) {
 	if ts := r.shMu.state.TruncatedState; ts != nil {
 		log.KvExec.Fatalf(ctx, "non-empty RaftTruncatedState in ReplicaState: %+v", ts)
-	} else if loaded, err := r.raftMu.stateLoader.LoadRaftTruncatedState(ctx, reader); err != nil {
+	} else if loaded, err := r.raftMu.stateLoader.LoadRaftTruncatedState(ctx, raftRO); err != nil {
 		log.KvExec.Fatalf(ctx, "%s", err)
 	} else if ts := r.asLogStorage().shMu.trunc; loaded != ts {
 		log.KvExec.Fatalf(ctx, "on-disk and in-memory RaftTruncatedState diverged: %s",
 			redact.Safe(pretty.Diff(loaded, ts)))
 	}
 
-	diskState, err := r.raftMu.stateLoader.Load(ctx, reader, r.shMu.state.Desc)
+	diskState, err := r.raftMu.stateLoader.Load(ctx, stateRO, r.shMu.state.Desc)
 	if err != nil {
 		log.KvExec.Fatalf(ctx, "%v", err)
 	}
@@ -2041,7 +2041,7 @@ func (r *Replica) assertStateRaftMuLockedReplicaMuRLocked(
 			log.KvExec.Fatalf(ctx, "replica's replicaID %d diverges from descriptor %+v", r.replicaID, r.shMu.state.Desc)
 		}
 	}
-	diskReplID, err := r.raftMu.stateLoader.LoadRaftReplicaID(ctx, reader)
+	diskReplID, err := r.raftMu.stateLoader.LoadRaftReplicaID(ctx, stateRO)
 	if err != nil {
 		log.KvExec.Fatalf(ctx, "%s", err)
 	}

pkg/kv/kvserver/replica_application_state_machine_test.go

@@ -255,7 +255,7 @@ func TestReplicaStateMachineRaftLogTruncationStronglyCoupled(t *testing.T) {
 			require.Equal(t, expectedSize, ls.shMu.size)
 			require.Equal(t, accurate, ls.shMu.sizeTrusted)
 			truncState, err := logstore.NewStateLoader(r.RangeID).LoadRaftTruncatedState(
-				context.Background(), tc.engine)
+				context.Background(), tc.raftEng)
 			require.NoError(t, err)
 			require.Equal(t, ls.shMu.trunc.Index, truncState.Index)
 		}()

pkg/kv/kvserver/replica_raft.go

@@ -1283,9 +1283,8 @@ func (r *Replica) handleRaftReadyRaftMuLocked(
 
 	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.assertStateRaftMuLockedReplicaMuRLocked(
+			ctx, sm.r.store.StateEngine(), sm.r.store.LogEngine())
 		sm.r.mu.RUnlock()
 	}