kvstorage/wag: prepare for WAG truncation

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
 			}
 		}

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 {

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"