kvstorage: delete raft state when deleting destroy/subsume events

When deleting a WAG node that has the events destroy or subsume, we also
need to delete the raft log (up until and including the addr.index), and
any sideloaded files that are part of that raft log prefix.

We do this because when a replica delete or subsume is applied, we cannot
just yet delete the replica state (log and sideloaded files) because we
might crash before the state engine has flushed, and we might need the
raft log to replay the event.

Release note: None

Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

Author: iskettaneh

pkg/kv/kvserver/kvstorage/BUILD.bazel

@@ -69,6 +69,7 @@ go_test(
         "//pkg/roachpb",
         "//pkg/settings/cluster",
         "//pkg/storage",
+        "//pkg/storage/fs",
         "//pkg/testutils",
         "//pkg/testutils/datapathutils",
         "//pkg/testutils/dd",
@@ -82,5 +83,6 @@ go_test(
         "//pkg/util/uuid",
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_stretchr_testify//require",
+        "@org_golang_x_time//rate",
     ],
 )

pkg/kv/kvserver/kvstorage/wag_truncator.go

@@ -8,10 +8,12 @@ package kvstorage
 import (
 	"context"
 
+	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag/wagpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/errors"
 )
 
@@ -102,33 +104,71 @@ func isNodeApplied(ctx context.Context, stateRO StateRO, node wagpb.Node) (bool,
 	return true, nil
 }
 
-// truncateAppliedNode checks the first WAG node and deletes it if all of its
-// events have been applied to the state engine. Returns true if a node was
-// deleted.
+// SideloadClearer truncates sideloaded files for a given range up to and
+// including the specified raft log index. Files beyond this index may belong
+// to a newer replica and must be preserved.
+// The SideloadClearer must also sync the sideloaded storage after truncation to
+// avoid leaking the files in case of a node crash.
+type SideloadClearer func(ctx context.Context, rangeID roachpb.RangeID, lastIndex kvpb.RaftIndex) error
+
+// clearReplicaRaftLog clears raft log entries at or below the given index for
+// a destroyed or subsumed replica.
+func clearReplicaRaftLog(
+	ctx context.Context, raft Raft, rangeID roachpb.RangeID, lastIndex kvpb.RaftIndex,
+) error {
+	return storage.ClearRangeWithHeuristic(
+		ctx, raft.RO, raft.WO,
+		keys.RaftLogPrefix(rangeID),           /* start */
+		keys.RaftLogKey(rangeID, lastIndex+1), /* end */
+		ClearRangeThresholdPointKeys(),
+	)
+}
+
+// truncateAppliedWAGNodeAndClearRaftState checks the first WAG node and
+// deletes it if all of its events have been applied to the state engine. For
+// nodes containing EventDestroy or EventSubsume events, it also clears the
+// raft log entries and sideloaded files for those replicas.
 //
 // The caller must provide a stateRO reader with GuaranteedDurability so that
 // only state confirmed flushed to persistent storage is visible. This ensures
 // we never delete a WAG node whose mutations aren't flushed yet.
-func truncateAppliedNode(ctx context.Context, raft Raft, stateRO StateRO) (bool, error) {
+// TODO(ibrahim): Support deleting multiple WAG nodes within the same batch.
+func truncateAppliedWAGNodeAndClearRaftState(
+	ctx context.Context, raft Raft, stateRO StateRO, clearSideloaded SideloadClearer,
+) (bool, error) {
 	var iter wag.Iterator
 	for index, node := range iter.Iter(ctx, raft.RO) {
 		applied, err := isNodeApplied(ctx, stateRO, node)
-		if err != nil {
+		if err != nil || !applied {
 			return false, err
 		}
-		if !applied {
-			return false, nil
-		}
 		if err := wag.Delete(raft.WO, index); err != nil {
 			return false, err
 		}
-		// TODO(Ibrahim): Add logic to clear raft state (log entries, HardState,
-		// TruncatedState) for destroyed/subsumed replicas.
-		// TODO(ibrahim): Support deleting multiple WAG nodes within the same batch.
+
+		// Clean up the raft log prefix of a destroyed/subsumed replica.
+		for _, event := range node.Events {
+			if event.Type != wagpb.EventDestroy && event.Type != wagpb.EventSubsume {
+				continue
+			}
+			if err := clearReplicaRaftLog(ctx, raft, event.Addr.RangeID, event.Addr.Index); err != nil {
+				return false, errors.Wrapf(err, "clearing raft log for r%d", event.Addr.RangeID)
+			}
+			if clearSideloaded != nil {
+				// In general, we shouldn't delete sideloaded files before committing the
+				// batch that deletes the Raft log entries. However, in this case, we
+				// know that the destroy/subsume event has already been flushed to disk,
+				// and the replica is destroyed. We will not need to reference the
+				// sideloaded files anymore. If we were to delay deleting them till
+				// after the batch is committed, we risk leaking the sideloaded files
+				// if a node crash happens after the batch is committed (and synced) and
+				// before the sideloaded files are deleted.
+				if err := clearSideloaded(ctx, event.Addr.RangeID, event.Addr.Index); err != nil {
+					return false, errors.Wrapf(err, "clearing sideloaded files for r%d", event.Addr.RangeID)
+				}
+			}
+		}
 		return true, nil
 	}
-	if err := iter.Error(); err != nil {
-		return false, err
-	}
-	return false, nil
+	return false, iter.Error()
 }

pkg/kv/kvserver/kvstorage/wag_truncator_test.go

@@ -7,16 +7,23 @@ package kvstorage
 
 import (
 	"context"
+	"math"
 	"testing"
 
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag/wagpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/stretchr/testify/require"
+	"golang.org/x/time/rate"
 )
 
 // testEngines holds the separated log and state engines used by truncator
@@ -45,25 +52,44 @@ func (e *testEngines) writeWAGNode(t *testing.T, event wagpb.Event) {
 	require.NoError(t, wag.Write(e.LogEngine(), index, wagpb.Node{Events: []wagpb.Event{event}}))
 }
 
-// truncateWAGNodes repeatedly truncates applied WAG nodes until none remain
-// applied.
-func (e *testEngines) truncateWAGNodes(t *testing.T) {
+// writeRaftLogEntry writes a dummy raft log entry to the log engine.
+func (e *testEngines) writeRaftLogEntry(
+	t *testing.T, rangeID roachpb.RangeID, index kvpb.RaftIndex,
+) {
 	t.Helper()
-	ctx := context.Background()
+	sl := MakeStateLoader(rangeID)
+	require.NoError(t, e.LogEngine().PutUnversioned(sl.RaftLogKey(index), []byte("entry")))
+}
+
+// raftLogIndices returns the indices of all raft log entries for the given range
+// in the log engine.
+func (e *testEngines) raftLogIndices(t *testing.T, rangeID roachpb.RangeID) []kvpb.RaftIndex {
+	t.Helper()
+	sl := MakeStateLoader(rangeID)
+	prefix := sl.RaftLogPrefix()
+	var indices []kvpb.RaftIndex
+	iter, err := e.LogEngine().NewMVCCIterator(
+		context.Background(), storage.MVCCKeyIterKind, storage.IterOptions{
+			LowerBound:   prefix,
+			UpperBound:   prefix.PrefixEnd(),
+			ReadCategory: fs.ReplicationReadCategory,
+		})
+	require.NoError(t, err)
+	defer iter.Close()
+	iter.SeekGE(storage.MakeMVCCMetadataKey(prefix))
 	for {
-		b := e.LogEngine().NewWriteBatch()
-		truncated, err := truncateAppliedNode(
-			ctx, Raft{RO: e.LogEngine(), WO: b}, e.StateEngine(),
-		)
-		if err == nil && truncated {
-			err = b.Commit(false /* sync */)
-		}
-		b.Close()
-		require.NoError(t, err)
-		if !truncated {
-			return
+		ok, iterErr := iter.Valid()
+		require.NoError(t, iterErr)
+		if !ok {
+			break
 		}
+		suffix := iter.UnsafeKey().Key[len(prefix):]
+		index, decErr := keys.DecodeRaftLogKeyFromSuffix(suffix)
+		require.NoError(t, decErr)
+		indices = append(indices, index)
+		iter.Next()
 	}
+	return indices
 }
 
 // listWAGNodes returns the indices of all WAG nodes in the log engine.
@@ -78,6 +104,39 @@ func (e *testEngines) listWAGNodes(t *testing.T) []uint64 {
 	return indices
 }
 
+// makeSideloadClearer returns a SideloadClearer that truncates sideloaded
+// files for a given range using disk-based sideload storage.
+func (e *testEngines) makeSideloadClearer(
+	st *cluster.Settings, baseDir string, limiter *rate.Limiter, env *fs.Env,
+) SideloadClearer {
+	return func(ctx context.Context, rangeID roachpb.RangeID, lastIndex kvpb.RaftIndex) error {
+		ss := logstore.NewDiskSideloadStorage(st, rangeID, baseDir, limiter, env)
+		if err := ss.TruncateTo(ctx, lastIndex); err != nil {
+			return err
+		}
+		return ss.Sync()
+	}
+}
+
+// truncateWAGNodes repeatedly calls truncateAppliedWAGNodeAndClearRaftState
+// until no more nodes can be deleted.
+func (e *testEngines) truncateWAGNodes(t *testing.T, clearSideloaded SideloadClearer) {
+	t.Helper()
+	require.NoError(t, e.StateEngine().Flush())
+	stateReader := e.StateEngine().NewReader(storage.GuaranteedDurability)
+	defer stateReader.Close()
+	for {
+		b := e.LogEngine().NewBatch()
+		ok, err := truncateAppliedWAGNodeAndClearRaftState(context.Background(),
+			Raft{RO: e.LogEngine(), WO: b}, stateReader, clearSideloaded)
+		require.NoError(t, err)
+		require.NoError(t, b.Commit(false /* sync */))
+		if !ok {
+			return
+		}
+	}
+}
+
 func TestTruncateApplied(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -215,8 +274,88 @@ func TestTruncateApplied(t *testing.T) {
 			e := makeTestEngines()
 			defer e.Close()
 			tc.setup(t, &e)
-			e.truncateWAGNodes(t)
+			e.truncateWAGNodes(t, nil /* clearSideloaded */)
 			require.Equal(t, tc.wantWAGIndices, e.listWAGNodes(t))
 		})
 	}
 }
+
+// TestTruncateAndClearRaftState verifies that
+// truncateAppliedWAGNodeAndClearRaftState only clears raft log entries and
+// sideloaded files up to the destroyed/subsumed replica's last index. Entries
+// and files beyond that index may belong to a newer replica and must be
+// preserved.
+func TestTruncateAndClearRaftState(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1}
+	r2 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 2}
+	sl := MakeStateLoader(1 /* rangeID */)
+
+	for _, eventType := range []wagpb.EventType{wagpb.EventDestroy, wagpb.EventSubsume} {
+		t.Run(eventType.String(), func(t *testing.T) {
+			e := makeTestEngines()
+			defer e.Close()
+			ctx := context.Background()
+
+			// Write WAG nodes: init then destroy/subsume at index 20.
+			e.writeWAGNode(t, wagpb.Event{
+				Addr: wagpb.MakeAddr(r1, 10), Type: wagpb.EventInit,
+			})
+			e.writeWAGNode(t, wagpb.Event{
+				Addr: wagpb.MakeAddr(r1, 20), Type: eventType,
+			})
+
+			// Create a WAG node for a newer replica for the same range.
+			e.writeWAGNode(t, wagpb.Event{
+				Addr: wagpb.MakeAddr(r2, 0), Type: wagpb.EventCreate,
+			})
+
+			// Tombstone confirms destruction/subsumption.
+			require.NoError(t, sl.SetRangeTombstone(ctx, e.StateEngine(),
+				kvserverpb.RangeTombstone{NextReplicaID: 2}))
+
+			// Raft log entries at indices 10-25. (Indices at or below lastIndex=20)
+			// belong to the old replica. Entries at 21-25 (above lastIndex=20)
+			// belong to a newer replica reusing the same RangeID.
+			for idx := 10; idx <= 25; idx++ {
+				e.writeRaftLogEntry(t, 1 /* rangeID */, kvpb.RaftIndex(idx))
+			}
+
+			// Create sideloaded files using the log engine's VFS.
+			st := cluster.MakeTestingClusterSettings()
+			baseDir := e.LogEngine().GetAuxiliaryDir()
+			limiter := rate.NewLimiter(rate.Inf, math.MaxInt64)
+			env := e.LogEngine().Env()
+			ss := logstore.NewDiskSideloadStorage(st, 1, baseDir, limiter, env)
+			for _, idx := range []kvpb.RaftIndex{10, 15, 20, 21, 25} {
+				require.NoError(t, ss.Put(ctx, idx, 1 /* term */, []byte("sst-data")))
+			}
+			clearer := e.makeSideloadClearer(st, baseDir, limiter, env)
+
+			e.truncateWAGNodes(t, clearer)
+			// Raft log entries 10-20 (at or below lastIndex=20) must be deleted. The
+			// rest shouldn't be deleted by the WAG truncator.
+			require.Equal(t,
+				[]kvpb.RaftIndex{21, 22, 23, 24, 25},
+				e.raftLogIndices(t, 1 /* rangeID */),
+			)
+			// Sideloaded files at indices 10, 15, and 20 (at or below lastIndex=20)
+			// must be deleted.
+			for _, idx := range []kvpb.RaftIndex{10, 15, 20} {
+				_, err := ss.Get(ctx, idx, 1)
+				require.Errorf(t, err, "sideloaded file at index %d should have been deleted", idx)
+			}
+			// Sideloaded files at indices 21 and 25 (above lastIndex=20) must be
+			// preserved.
+			for _, idx := range []kvpb.RaftIndex{21, 25} {
+				data, err := ss.Get(ctx, idx, 1)
+				require.NoErrorf(t, err, "sideloaded file at index %d should be preserved", idx)
+				require.Equal(t, []byte("sst-data"), data)
+			}
+			// Only the WAG node for the newer replica should be left.
+			require.Equal(t, []uint64{3}, e.listWAGNodes(t))
+		})
+	}
+}