From 39d4b4d347a982fb3200824c7f6d19d6dc4ef58a Mon Sep 17 00:00:00 2001 From: iskettaneh <173953022+iskettaneh@users.noreply.github.com> Date: Fri, 3 Apr 2026 19:41:28 -0400 Subject: [PATCH] kvstorage/wag: prepare for WAG truncation This commit sets up the premitives for WAG truncations by: 1) Adding a `Delete` function for removing WAG nodes by index. 2) Changing the WAG `Iterator.Iter` method to return `iter.Seq2[uint64, wagpb.Node]` pair. This will be useful when we iterate over the WAG and truncate the nodes that have been applied and synced. Epic: none Release note: None Co-Authored-By: roachdev-claude --- pkg/kv/kvserver/kvstorage/wag/store.go | 25 ++++++-- pkg/kv/kvserver/kvstorage/wag/store_test.go | 61 ++++++++++++++++--- .../kvstorage/wag/testdata/TestWrite.txt | 4 +- 3 files changed, 75 insertions(+), 15 deletions(-) diff --git a/pkg/kv/kvserver/kvstorage/wag/store.go b/pkg/kv/kvserver/kvstorage/wag/store.go index 0d56765d9ba..cc8b1e27c01 100644 --- a/pkg/kv/kvserver/kvstorage/wag/store.go +++ b/pkg/kv/kvserver/kvstorage/wag/store.go @@ -76,11 +76,18 @@ func Write(w storage.Writer, index uint64, node wagpb.Node) error { return w.PutUnversioned(keys.StoreWAGNodeKey(index), data) } +// Delete removes the WAG node at the given sequence number. +// TODO(ibrahim): Consider SingleClearEngineKey if we can guarantee that the +// index is written only once (even after restarts). +func Delete(w storage.Writer, index uint64) error { + return w.ClearUnversioned(keys.StoreWAGNodeKey(index), storage.ClearOptions{}) +} + // Iterator helps to scan the WAG sequence. // // var iter wag.Iterator -// for node := range iter.Iter(ctx, reader) { -// // process node +// for index, node := range iter.Iter(ctx, reader) { +// // process index, node // } // if err := iter.Error(); err != nil { // return err @@ -92,8 +99,9 @@ type Iterator struct { err error } -// Iter returns an iterator that scans the WAG sequence. -func (it *Iterator) Iter(ctx context.Context, r storage.Reader) iter.Seq[wagpb.Node] { +// Iter returns an iterator that scans the WAG sequence. The iterator yields a +// pair containing the WAG node index and the WAG node itself. +func (it *Iterator) Iter(ctx context.Context, r storage.Reader) iter.Seq2[uint64, wagpb.Node] { prefix := keys.StoreWAGPrefix() mi, err := r.NewMVCCIterator(ctx, storage.MVCCKeyIterKind, storage.IterOptions{ UpperBound: prefix.PrefixEnd(), @@ -104,13 +112,18 @@ func (it *Iterator) Iter(ctx context.Context, r storage.Reader) iter.Seq[wagpb.N } mi.SeekGE(storage.MakeMVCCMetadataKey(prefix)) - return func(yield func(wagpb.Node) bool) { + return func(yield func(uint64, wagpb.Node) bool) { defer mi.Close() for ; ; mi.Next() { if ok, err := mi.Valid(); err != nil || !ok { it.err = err return } + index, err := keys.DecodeWAGNodeKey(mi.UnsafeKey().Key) + if err != nil { + it.err = err + return + } v, err := mi.UnsafeValue() if err != nil { it.err = err @@ -120,7 +133,7 @@ func (it *Iterator) Iter(ctx context.Context, r storage.Reader) iter.Seq[wagpb.N if it.err = node.Unmarshal(v); it.err != nil { // nolint:protounmarshal return } - if !yield(node) { + if !yield(index, node) { return } } diff --git a/pkg/kv/kvserver/kvstorage/wag/store_test.go b/pkg/kv/kvserver/kvstorage/wag/store_test.go index 044f41aea1a..d6006820955 100644 --- a/pkg/kv/kvserver/kvstorage/wag/store_test.go +++ b/pkg/kv/kvserver/kvstorage/wag/store_test.go @@ -47,6 +47,9 @@ func TestWrite(t *testing.T) { id := roachpb.FullReplicaID{RangeID: 123, ReplicaID: 4} rhsID := roachpb.FullReplicaID{RangeID: 567, ReplicaID: 1} write("create", func(w storage.Writer) error { return createReplica(&s, w, id) }) + // Intentionally introduce a gap in the sequence. We will later make sure that + // the iterator correctly skips over this gap. + s.seq.Next(1) write("init", func(w storage.Writer) error { return initReplica(&s, w, id, 10) }) write("split", func(w storage.Writer) error { return splitReplica(&s, w, id, rhsID, 200) }) @@ -55,16 +58,60 @@ func TestWrite(t *testing.T) { out = strings.ReplaceAll(out, "\n\n", "\n") echotest.Require(t, out, filepath.Join("testdata", t.Name()+".txt")) - // Smoke check that the iterator works. - var iter Iterator - count := 0 - for range iter.Iter(context.Background(), s.eng) { - count++ + // readIndices returns the WAG node indices by scanning the engine. + readIndices := func() []uint64 { + var it Iterator + var res []uint64 + for index := range it.Iter(context.Background(), s.eng) { + res = append(res, index) + } + require.NoError(t, it.Error()) + return res } - require.NoError(t, iter.Error()) + + // Verify that the iterator returns nodes with the correct indices. // 3 WAG nodes: create, init, split. The split is a single node with two // events (Split + Init) rather than two separate nodes (dep + event). - require.Equal(t, 3, count) + require.Equal(t, []uint64{1, 3, 4}, readIndices()) +} + +func TestDelete(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + eng := storage.NewDefaultInMemForTesting() + defer eng.Close() + + id := roachpb.FullReplicaID{RangeID: 1, ReplicaID: 1} + node := wagpb.Node{ + Events: []wagpb.Event{ + {Addr: wagpb.MakeAddr(id, 10), Type: wagpb.EventApply}, + }, + } + + // Write 5 WAG nodes with indices 1 through 5. + for i := uint64(1); i <= 5; i++ { + require.NoError(t, Write(eng, i, node)) + } + + // Read back all indices. + readIndices := func() []uint64 { + var it Iterator + var res []uint64 + for index := range it.Iter(context.Background(), eng) { + res = append(res, index) + } + require.NoError(t, it.Error()) + return res + } + require.Equal(t, []uint64{1, 2, 3, 4, 5}, readIndices()) + + // Delete nodes 2 and 4. + require.NoError(t, Delete(eng, 2)) + require.NoError(t, Delete(eng, 4)) + + // Verify that only nodes 1, 3, 5 remain. + require.Equal(t, []uint64{1, 3, 5}, readIndices()) } type store struct { diff --git a/pkg/kv/kvserver/kvstorage/wag/testdata/TestWrite.txt b/pkg/kv/kvserver/kvstorage/wag/testdata/TestWrite.txt index 6e10700ecdf..8a071d848ac 100644 --- a/pkg/kv/kvserver/kvstorage/wag/testdata/TestWrite.txt +++ b/pkg/kv/kvserver/kvstorage/wag/testdata/TestWrite.txt @@ -4,9 +4,9 @@ echo Put: 0,0 /Local/Store/wag/1 (0x01737761676e000000000000000100): (r123/4:0,EventCreate) > Put: 0,0 "state-machine-key" (0x73746174652d6d616368696e652d6b657900): "state" >> init -Put: 0,0 /Local/Store/wag/2 (0x01737761676e000000000000000200): (r123/4:10,EventInit) +Put: 0,0 /Local/Store/wag/3 (0x01737761676e000000000000000300): (r123/4:10,EventInit) ingestion: SSTs:"tmp/1.sst" SSTs:"tmp/2.sst" >> split -Put: 0,0 /Local/Store/wag/3 (0x01737761676e000000000000000300): (r567/1:10,EventInit) (r123/4:200,EventSplit) +Put: 0,0 /Local/Store/wag/4 (0x01737761676e000000000000000400): (r567/1:10,EventInit) (r123/4:200,EventSplit) > Put: 0,0 "lhs-key" (0x6c68732d6b657900): "lhs-state" > Put: 0,0 "rhs-key" (0x7268732d6b657900): "rhs-state"