Merge #156482

156482: kvserver: annotate engines in splitPreApply r=arulajmani a=pav-kv

This PR annotates all storage operations with raft/state engine types, down the stack from `splitPreApply`.

Epic: CRDB-55220

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

Author: craig[bot]

pkg/kv/kvserver/kvstorage/destroy.go

@@ -195,16 +195,12 @@ func SubsumeReplica(
 	}, destroyReplicaImpl(ctx, reader, writer, info, MergedTombstoneReplicaID)
 }
 
-// RemoveStaleRHSFromSplit removes all data for the RHS replica of a split. This
-// is used in a situation when the RHS replica is already known to have been
-// removed from our store, so any pending writes that were supposed to
+// RemoveStaleRHSFromSplit removes all replicated data for the RHS replica of a
+// split. This is used in a situation when the RHS replica is already known to
+// have been removed from our store, so any pending writes that were supposed to
 // initialize the RHS replica should be dropped from the write batch.
 func RemoveStaleRHSFromSplit(
-	ctx context.Context,
-	reader storage.Reader,
-	writer storage.Writer,
-	rangeID roachpb.RangeID,
-	keys roachpb.RSpan,
+	ctx context.Context, stateRW State, rangeID roachpb.RangeID, keys roachpb.RSpan,
 ) error {
 	for _, span := range rditer.Select(rangeID, rditer.SelectOpts{
 		// Since the RHS replica is uninitalized, we know there isn't anything in
@@ -219,7 +215,7 @@ func RemoveStaleRHSFromSplit(
 		UnreplicatedByRangeID: false,
 	}) {
 		if err := storage.ClearRangeWithHeuristic(
-			ctx, reader, writer, span.Key, span.EndKey, ClearRangeThresholdPointKeys(),
+			ctx, stateRW.RO, stateRW.WO, span.Key, span.EndKey, ClearRangeThresholdPointKeys(),
 		); err != nil {
 			return err
 		}

pkg/kv/kvserver/kvstorage/storage.go

@@ -55,11 +55,17 @@ type ReadWriter struct {
 
 // TODOState interprets the provided storage accessor as the State engine.
 //
-// TODO(pav-kv): remove when all callers have clarified their access patterns.
+// TODO(pav-kv): remove when all callers have clarified their access patterns,
+// and switched to WrapState() or other explicitly defined engine types.
 func TODOState(rw storage.ReadWriter) State {
 	return State{RO: rw, WO: rw}
 }
 
+// WrapState interprets the provided storage accessor as the State engine.
+func WrapState(rw StateRW) State {
+	return State{RO: rw, WO: rw}
+}
+
 // TODORaft interprets the provided storage accessor as the Raft engine.
 //
 // TODO(pav-kv): remove when all callers have clarified their access patterns.

pkg/kv/kvserver/replica_app_batch.go

@@ -321,7 +321,10 @@ func (b *replicaAppBatch) runPostAddTriggersReplicaOnly(
 		// NB: another reason why we shouldn't write HardState at evaluation time is
 		// that it belongs to the log engine, whereas the evaluated batch must
 		// contain only state machine updates.
-		splitPreApply(ctx, b.r, b.batch, res.Split.SplitTrigger, cmd.Cmd.ClosedTimestamp)
+		splitPreApply(
+			ctx, b.r, kvstorage.StateRW(b.batch), kvstorage.TODORaft(b.batch),
+			res.Split.SplitTrigger, cmd.Cmd.ClosedTimestamp,
+		)
 
 		// The rangefeed processor will no longer be provided logical ops for
 		// its entire range, so it needs to be shut down and all registrations

pkg/kv/kvserver/store_split.go

@@ -13,7 +13,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/load"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
-	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -29,7 +28,8 @@ import (
 func splitPreApply(
 	ctx context.Context,
 	r *Replica,
-	readWriter storage.ReadWriter,
+	stateRW kvstorage.StateRW,
+	raftRW kvstorage.Raft,
 	split roachpb.SplitTrigger,
 	initClosedTS *hlc.Timestamp,
 ) {
@@ -79,13 +79,17 @@ func splitPreApply(
 	// on this Store, it doesn't have a RaftTruncatedState (which only initialized
 	// replicas can have), so this deletion will not conflict with or corrupt it.
 	//
+	// NB: the key is cleared in stateRW rather than raftRW, deliberately. It
+	// lives in the raft engine, but here we want to clear it from the state
+	// engine batch, so that it doesn't make it to the state engine.
+	//
 	// TODO(#152847): remove this workaround when there are no historical
 	// proposals with RaftTruncatedState, e.g. after a below-raft migration.
-	if ts, err := rsl.LoadRaftTruncatedState(ctx, readWriter); err != nil {
+	if ts, err := rsl.LoadRaftTruncatedState(ctx, stateRW); err != nil {
 		log.KvExec.Fatalf(ctx, "cannot load RaftTruncatedState: %v", err)
 	} else if ts == (kvserverpb.RaftTruncatedState{}) {
 		// Common case. Do nothing.
-	} else if err := rsl.ClearRaftTruncatedState(readWriter); err != nil {
+	} else if err := rsl.ClearRaftTruncatedState(stateRW); err != nil {
 		log.KvExec.Fatalf(ctx, "cannot clear RaftTruncatedState: %v", err)
 	}
 
@@ -122,7 +126,7 @@ func splitPreApply(
 			}
 		}
 		if err := kvstorage.RemoveStaleRHSFromSplit(
-			ctx, readWriter, readWriter, split.RightDesc.RangeID, split.RightDesc.RSpan(),
+			ctx, kvstorage.WrapState(stateRW), split.RightDesc.RangeID, split.RightDesc.RSpan(),
 		); err != nil {
 			log.KvExec.Fatalf(ctx, "failed to clear range data for removed rhs: %v", err)
 		}
@@ -138,10 +142,10 @@ func splitPreApply(
 	//
 	// [*] Note that uninitialized replicas may cast votes, and if they have, we
 	// can't load the default Term and Vote values.
-	if err := rsl.SynthesizeRaftState(ctx, readWriter, kvstorage.TODORaft(readWriter)); err != nil {
+	if err := rsl.SynthesizeRaftState(ctx, stateRW, raftRW); err != nil {
 		log.KvExec.Fatalf(ctx, "%v", err)
 	}
-	if err := rsl.SetRaftTruncatedState(ctx, readWriter, &kvserverpb.RaftTruncatedState{
+	if err := rsl.SetRaftTruncatedState(ctx, raftRW.WO, &kvserverpb.RaftTruncatedState{
 		Index: kvstorage.RaftInitialLogIndex,
 		Term:  kvstorage.RaftInitialLogTerm,
 	}); err != nil {
@@ -159,7 +163,7 @@ func splitPreApply(
 		initClosedTS = &hlc.Timestamp{}
 	}
 	initClosedTS.Forward(r.GetCurrentClosedTimestamp(ctx))
-	if err := rsl.SetClosedTimestamp(ctx, readWriter, *initClosedTS); err != nil {
+	if err := rsl.SetClosedTimestamp(ctx, stateRW, *initClosedTS); err != nil {
 		log.KvExec.Fatalf(ctx, "%s", err)
 	}
 }

pkg/kv/kvserver/store_test.go

@@ -4043,12 +4043,10 @@ func TestSplitPreApplyInitializesTruncatedState(t *testing.T) {
 	}, &rightDesc.InternalReplicas[0])
 	require.NoError(t, err)
 
-	split := roachpb.SplitTrigger{
-		LeftDesc:  leftDesc,
-		RightDesc: rightDesc,
-	}
-
-	splitPreApply(ctx, lhsRepl, batch, split, nil)
+	splitPreApply(
+		ctx, lhsRepl, kvstorage.StateRW(batch), kvstorage.TODORaft(batch),
+		roachpb.SplitTrigger{LeftDesc: leftDesc, RightDesc: rightDesc}, nil,
+	)
 
 	// Verify that the RHS truncated state is initialized as expected.
 	rsl := kvstorage.MakeStateLoader(rightDesc.RangeID)