diff --git a/pkg/kv/kvserver/kvstorage/BUILD.bazel b/pkg/kv/kvserver/kvstorage/BUILD.bazel index b2a76561a9e0..44fc5b579bf4 100644 --- a/pkg/kv/kvserver/kvstorage/BUILD.bazel +++ b/pkg/kv/kvserver/kvstorage/BUILD.bazel @@ -12,6 +12,7 @@ go_library( "stateloader.go", "storage.go", "wag_replay.go", + "wag_truncator.go", ], importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage", visibility = ["//visibility:public"], @@ -51,6 +52,7 @@ go_test( "stateloader_test.go", "storage_test.go", "wag_replay_test.go", + "wag_truncator_test.go", ], data = glob(["testdata/**"]), embed = [":kvstorage"], @@ -69,6 +71,7 @@ go_test( "//pkg/roachpb", "//pkg/settings/cluster", "//pkg/storage", + "//pkg/storage/fs", "//pkg/testutils", "//pkg/testutils/datapathutils", "//pkg/testutils/dd", @@ -82,5 +85,6 @@ go_test( "//pkg/util/uuid", "@com_github_cockroachdb_datadriven//:datadriven", "@com_github_stretchr_testify//require", + "@org_golang_x_time//rate", ], ) diff --git a/pkg/kv/kvserver/kvstorage/wag_truncator.go b/pkg/kv/kvserver/kvstorage/wag_truncator.go new file mode 100644 index 000000000000..03103a7a3232 --- /dev/null +++ b/pkg/kv/kvserver/kvstorage/wag_truncator.go @@ -0,0 +1,100 @@ +// Copyright 2026 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +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" +) + +// 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 +// corresponding raft log prefix from the engine and the sideloaded entries +// storage. +// +// 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. +// +// Returns a boolean indicating whether a node was successfully truncated or +// not. If the return value is false, it means that either there are no WAG +// nodes left, or that the WAG node has not been applied to the state engine. +// Also, an error is returned if the WAG node could not be fetched or deleted. +// 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) { + // TODO(ibrahim): Right now, the canApplyWAGNode function returns a list of + // raftCatchUpTargets that are not needed for the purposes of truncation, + // consider refactoring the function to return only the needed info. + replayAction, err := canApplyWAGNode(ctx, node, stateRO) + if err != nil { + return false, err + } + if replayAction.apply { + // If an event needs to be applied, the WAG node cannot be deleted yet. + return false, nil + } + if err := wag.Delete(raft.WO, index); err != nil { + return false, err + } + + // 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 + } + return false, iter.Error() +} diff --git a/pkg/kv/kvserver/kvstorage/wag_truncator_test.go b/pkg/kv/kvserver/kvstorage/wag_truncator_test.go new file mode 100644 index 000000000000..38e5d8120604 --- /dev/null +++ b/pkg/kv/kvserver/kvstorage/wag_truncator_test.go @@ -0,0 +1,360 @@ +// Copyright 2026 The Cockroach Authors. +// +// Use of this software is governed by the CockroachDB Software License +// included in the /LICENSE file. + +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 +// tests. +type testEngines struct { + Engines + seq wag.Seq +} + +// makeTestEngines creates separated log and state engines. The caller must +// defer Close() to ensure engines are closed before leak checks. +func makeTestEngines() testEngines { + return testEngines{ + Engines: MakeSeparatedEnginesForTesting( + storage.NewDefaultInMemForTesting(), + storage.NewDefaultInMemForTesting(), + ), + } +} + +// writeWAGNode writes a WAG node containing the given event to the log engine, +// auto-incrementing the WAG sequence number. +func (e *testEngines) writeWAGNode(t *testing.T, event wagpb.Event) { + t.Helper() + index := e.seq.Next() + require.NoError(t, wag.Write(e.LogEngine(), index, wagpb.Node{Events: []wagpb.Event{event}})) +} + +// 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() + require.NoError(t, e.LogEngine().PutUnversioned( + keys.RaftLogKeyFromPrefix(keys.RaftLogPrefix(rangeID), 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 ; ; iter.Next() { + ok, err := iter.Valid() + require.NoError(t, err) + if !ok { + break + } + suffix := iter.UnsafeKey().Key[len(prefix):] + index, err := keys.DecodeRaftLogKeyFromSuffix(suffix) + require.NoError(t, err) + indices = append(indices, index) + } + return indices +} + +// listWAGNodes returns the indices of all WAG nodes in the log engine. +func (e *testEngines) listWAGNodes(t *testing.T) []uint64 { + t.Helper() + var indices []uint64 + var iter wag.Iterator + for index := range iter.Iter(context.Background(), e.LogEngine()) { + indices = append(indices, index) + } + require.NoError(t, iter.Error()) + 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, env *fs.Env, +) SideloadClearer { + return func(ctx context.Context, rangeID roachpb.RangeID, lastIndex kvpb.RaftIndex) error { + limiter := rate.NewLimiter(rate.Inf, math.MaxInt64) + 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 */)) + b.Close() + if !ok { + return + } + } +} + +func TestTruncateApplied(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + ctx := context.Background() + r1 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1} + r2 := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 2} + sl := MakeStateLoader(1 /* rangeID */) + + for _, tc := range []struct { + setup func(t *testing.T, e *testEngines) + wantWAGIndices []uint64 + }{ + { + setup: func(t *testing.T, e *testEngines) { + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID)) + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 10})) + }, + // There are no WAG nodes, so the truncation is a no-op. + wantWAGIndices: nil, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r2, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r2, 15), Type: wagpb.EventInit}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r2, 20), Type: wagpb.EventApply}) + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID)) + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 9})) + }, + // No WAG node has been applied yet because the current + // replicaID is < the event replicaID. + wantWAGIndices: []uint64{1, 2, 3}, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 15), Type: wagpb.EventInit}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventApply}) + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID)) + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 15})) + }, + // Raft applied index is 15, only the first two WAG nodes should be + // truncated. + wantWAGIndices: []uint64{3}, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 15), Type: wagpb.EventInit}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventApply}) + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID)) + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 20})) + }, + // Raft applied index is 20, all WAG nodes should be truncated. + wantWAGIndices: nil, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventDestroy}) + require.NoError(t, sl.SetRangeTombstone(ctx, e.StateEngine(), + kvserverpb.RangeTombstone{NextReplicaID: r1.ReplicaID + 1})) + }, + // A range tombstone confirms destruction, WAG nodes to that replica + // should be truncated. + wantWAGIndices: nil, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventSubsume}) + require.NoError(t, sl.SetRangeTombstone(ctx, e.StateEngine(), + kvserverpb.RangeTombstone{NextReplicaID: r1.ReplicaID + 1})) + }, + // A range tombstone confirms subsumption, WAG nodes to that replica + // should be truncated. + wantWAGIndices: nil, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventDestroy}) + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r2.ReplicaID)) + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 20})) + }, + // A newer replica supersedes the old one, WAG nodes for the old + // replica should be truncated. + wantWAGIndices: nil, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventSubsume}) + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r2.ReplicaID)) + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 20})) + }, + // A newer replica supersedes the subsumed one, WAG nodes for the old + // replica should be truncated. + wantWAGIndices: nil, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventDestroy}) + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID)) + // EventDestroy won't be deleted if the replicaID still matches the + // WAG node's replica ID because it means that the event hasn't applied + // yet. In this case, we won't even read the applied index. + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 20})) + }, + wantWAGIndices: []uint64{2}, + }, + { + setup: func(t *testing.T, e *testEngines) { + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 0), Type: wagpb.EventCreate}) + e.writeWAGNode(t, wagpb.Event{Addr: wagpb.MakeAddr(r1, 20), Type: wagpb.EventSubsume}) + require.NoError(t, sl.SetRaftReplicaID(ctx, e.StateEngine(), r1.ReplicaID)) + // EventSubsume won't be deleted if the replicaID still matches the + // WAG node's replica ID because it means that the event hasn't applied + // yet. In this case, we won't even read the applied index. + require.NoError(t, sl.SetRangeAppliedState(ctx, e.StateEngine(), + &kvserverpb.RangeAppliedState{RaftAppliedIndex: 20})) + }, + wantWAGIndices: []uint64{2}, + }, + } { + t.Run("", func(t *testing.T) { + e := makeTestEngines() + defer e.Close() + tc.setup(t, &e) + 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})) + + // Write raft log entries 10-25. Entries <= 20 belong to the old replica. + // Entries >= 21 belong to the new replica using 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() + env := e.LogEngine().Env() + limiter := rate.NewLimiter(rate.Inf, math.MaxInt64) + 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, env) + + e.truncateWAGNodes(t, clearer) + // Raft entries <= 20 belong to the old replica and 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 entries 10, 15, and 20 belong to the old replica + // and must be deleted. The rest shouldn't be deleted by the WAG + // truncator. + 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) + } + 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)) + }) + } +}