Merge #152665

152665: wag: introduce WAG storage r=arulajmani a=pav-kv

This PR introduces the WAG sequence writer and the corresponding Store-local schema. The writes are tested with a basic `echotest` imitating the future use of the WAG.

Resolves #149604

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

Author: craig[bot]

pkg/BUILD.bazel

@@ -265,6 +265,7 @@ ALL_TESTS = [
     "//pkg/kv/kvserver/kvflowcontrol/rac2:rac2_test",
     "//pkg/kv/kvserver/kvflowcontrol/replica_rac2:replica_rac2_test",
     "//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv_test",
+    "//pkg/kv/kvserver/kvstorage/wag:wag_test",
     "//pkg/kv/kvserver/kvstorage:kvstorage_test",
     "//pkg/kv/kvserver/leases:leases_test",
     "//pkg/kv/kvserver/liveness:liveness_test",
@@ -1564,6 +1565,8 @@ GO_TARGETS = [
     "//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv",
     "//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv_test",
     "//pkg/kv/kvserver/kvstorage/wag/wagpb:wagpb",
+    "//pkg/kv/kvserver/kvstorage/wag:wag",
+    "//pkg/kv/kvserver/kvstorage/wag:wag_test",
     "//pkg/kv/kvserver/kvstorage:kvstorage",
     "//pkg/kv/kvserver/kvstorage:kvstorage_test",
     "//pkg/kv/kvserver/leases:leases",

pkg/keys/constants.go

@@ -226,6 +226,8 @@ var (
 	// SupporterMeta stores the highest timestamp at which support has been
 	// withdrawn.
 	localStoreLivenessSupporterMeta = []byte("slsm")
+	// localStoreWAGNodeSuffix is the suffix for WAG nodes.
+	localStoreWAGNodeSuffix = []byte("wagn")
 	// localStoreLivenessSupportFor stores the Store Liveness support by the local
 	// store for a store in the cluster. It includes the epoch and expiration of
 	// support.

pkg/keys/doc.go

@@ -219,6 +219,7 @@ var _ = [...]interface{}{
 	StoreLivenessSupporterMetaKey,    // "slsm"
 	StoreCachedSettingsKey,           // "stng"
 	StoreLastUpKey,                   // "uptm"
+	StoreWAGNodeKey,                  // "wagn"
 
 	//   5. Range lock keys for all replicated locks. All range locks share
 	//   LocalRangeLockTablePrefix. Locks can be acquired on global keys and on

pkg/keys/keys.go

@@ -140,6 +140,32 @@ func DecodeStoreLivenessSupportForKey(key roachpb.Key) (roachpb.NodeID, roachpb.
 	return nodeID, storeID, nil
 }
 
+// StoreWAGPrefix returns the key prefix for WAG nodes.
+func StoreWAGPrefix() roachpb.Key {
+	return MakeStoreKey(localStoreWAGNodeSuffix, nil)
+}
+
+// StoreWAGNodeKey returns the key for a WAG node at the given index.
+func StoreWAGNodeKey(index uint64) roachpb.Key {
+	return MakeStoreKey(localStoreWAGNodeSuffix, encoding.EncodeUint64Ascending(nil, index))
+}
+
+// DecodeWAGNodeKey returns the index of the WAG node from its key.
+func DecodeWAGNodeKey(key roachpb.Key) (uint64, error) {
+	suffix, detail, err := DecodeStoreKey(key)
+	if err != nil {
+		return 0, err
+	}
+	if !suffix.Equal(localStoreWAGNodeSuffix) {
+		return 0, errors.Errorf("key with suffix %q != %q", suffix, localStoreWAGNodeSuffix)
+	}
+	detail, index, err := encoding.DecodeUint64Ascending(detail)
+	if len(detail) != 0 {
+		return 0, errors.Errorf("invalid key has trailing garbage: %q", detail)
+	}
+	return index, err
+}
+
 // StoreCachedSettingsKey returns a store-local key for store's cached settings.
 func StoreCachedSettingsKey(settingKey roachpb.Key) roachpb.Key {
 	return MakeStoreKey(localStoreCachedSettingsSuffix, encoding.EncodeBytesAscending(nil, settingKey))

pkg/keys/printer.go

@@ -146,6 +146,7 @@ var constSubKeyDict = []struct {
 	{"/clusterVersion", localStoreClusterVersionSuffix},
 	{"/nodeTombstone", localStoreNodeTombstoneSuffix},
 	{"/cachedSettings", localStoreCachedSettingsSuffix},
+	{"/wag", localStoreWAGNodeSuffix},
 	{"/lossOfQuorumRecovery/applied", localStoreUnsafeReplicaRecoverySuffix},
 	{"/lossOfQuorumRecovery/status", localStoreLossOfQuorumRecoveryStatusSuffix},
 	{"/lossOfQuorumRecovery/cleanup", localStoreLossOfQuorumRecoveryCleanupActionsSuffix},
@@ -167,10 +168,19 @@ func cachedSettingsKeyPrint(buf *redact.StringBuilder, key roachpb.Key) {
 	buf.Print(settingKey.String())
 }
 
+func wagNodeKeyPrint(buf *redact.StringBuilder, key roachpb.Key) {
+	index, err := DecodeWAGNodeKey(key)
+	if err != nil {
+		buf.Printf("<invalid: %s>", err)
+	}
+	buf.Printf("%d", index)
+}
+
 func localStoreKeyPrint(buf *redact.StringBuilder, _ []encoding.Direction, key roachpb.Key) {
 	for _, v := range constSubKeyDict {
 		if bytes.HasPrefix(key, v.key) {
 			buf.Print(v.name)
+			// TODO(pav-kv): make this switch more efficient with a lookup.
 			if v.key.Equal(localStoreNodeTombstoneSuffix) {
 				buf.SafeRune('/')
 				nodeTombstoneKeyPrint(
@@ -186,6 +196,9 @@ func localStoreKeyPrint(buf *redact.StringBuilder, _ []encoding.Direction, key r
 				lossOfQuorumRecoveryEntryKeyPrint(
 					buf, append(roachpb.Key(nil), append(LocalStorePrefix, key...)...),
 				)
+			} else if v.key.Equal(localStoreWAGNodeSuffix) {
+				buf.SafeRune('/')
+				wagNodeKeyPrint(buf, append(append(roachpb.Key(nil), LocalStorePrefix...), key...))
 			}
 			return
 		}

pkg/kv/kvserver/kvstorage/wag/BUILD.bazel

@@ -0,0 +1,31 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "wag",
+    srcs = ["store.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/keys",
+        "//pkg/kv/kvserver/kvstorage/wag/wagpb",
+        "//pkg/storage",
+    ],
+)
+
+go_test(
+    name = "wag_test",
+    srcs = ["store_test.go"],
+    data = glob(["testdata/**"]),
+    embed = [":wag"],
+    deps = [
+        "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver/kvstorage/wag/wagpb",
+        "//pkg/kv/kvserver/print",
+        "//pkg/roachpb",
+        "//pkg/storage",
+        "//pkg/testutils/echotest",
+        "//pkg/util/leaktest",
+        "//pkg/util/log",
+        "@com_github_stretchr_testify//require",
+    ],
+)

pkg/kv/kvserver/kvstorage/wag/store.go

@@ -0,0 +1,105 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package wag
+
+import (
+	"context"
+	"iter"
+	"sync/atomic"
+
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag/wagpb"
+	"github.com/cockroachdb/cockroach/pkg/storage"
+)
+
+// Seq is the WAG sequencer. It allocates consecutive unique sequence numbers to
+// assist the writes of the WAG nodes in a topologically sorted order.
+type Seq struct {
+	// index is the last allocated index into the WAG nodes sequence.
+	// TODO(pav-kv): initialize it on store restarts.
+	index atomic.Uint64
+}
+
+// Next allocates the given number of consecutive WAG nodes in the sequence, and
+// returns the index of the first allocated node. The caller can subsequently
+// use indices [Next, Next+count) for writing WAG nodes, and must not use
+// indices outside this span.
+//
+// The caller must make sure that conflicting writers call Next and do the
+// corresponding writes according to the topological ordering of these
+// mutations. For example, replica lifecycle events of a single range must be
+// ordered. Independent / concurrent writers can call Next and perform the
+// corresponding writes in any order (e.g. different ranges can write their
+// events concurrently).
+func (s *Seq) Next(count uint64) uint64 {
+	return s.index.Add(count) - count + 1
+}
+
+// Write puts the WAG node under the specific sequence number into the given
+// writer. The index must have been allocated to the caller by the sequencer.
+func Write(w storage.Writer, index uint64, node wagpb.Node) error {
+	data, err := node.Marshal() // nolint:protomarshal
+	if err != nil {
+		return err
+	}
+	return w.PutUnversioned(keys.StoreWAGNodeKey(index), data)
+}
+
+// Iterator helps to scan the WAG sequence.
+//
+//	var iter wag.Iterator
+//	for node := range iter.Iter(ctx, reader) {
+//		// process node
+//	}
+//	if err := iter.Error(); err != nil {
+//		return err
+//	}
+//
+// TODO(pav-kv): make it more flexible, e.g. iterate from a particular index.
+type Iterator struct {
+	// err is the last error encountered during the WAG iteration.
+	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] {
+	prefix := keys.StoreWAGPrefix()
+	mi, err := r.NewMVCCIterator(ctx, storage.MVCCKeyIterKind, storage.IterOptions{
+		UpperBound: prefix.PrefixEnd(),
+	})
+	if err != nil {
+		it.err = err
+		return nil
+	}
+	mi.SeekGE(storage.MakeMVCCMetadataKey(prefix))
+
+	return func(yield func(wagpb.Node) bool) {
+		defer mi.Close()
+		for ; ; mi.Next() {
+			if ok, err := mi.Valid(); err != nil || !ok {
+				it.err = err
+				return
+			}
+			v, err := mi.UnsafeValue()
+			if err != nil {
+				it.err = err
+				return
+			}
+			var node wagpb.Node
+			if it.err = node.Unmarshal(v); it.err != nil { // nolint:protounmarshal
+				return
+			}
+			if !yield(node) {
+				return
+			}
+		}
+	}
+}
+
+// Error returns the last encountered error during the iteration.
+func (it *Iterator) Error() error {
+	return it.err
+}

pkg/kv/kvserver/kvstorage/wag/store_test.go

@@ -0,0 +1,121 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package wag
+
+import (
+	"context"
+	"fmt"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag/wagpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/print"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/testutils/echotest"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWrite(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	eng := storage.NewDefaultInMemForTesting()
+	defer eng.Close()
+	s := store{eng: eng}
+
+	var out string
+	write := func(name string, f func(w storage.Writer) error) {
+		b := eng.NewWriteBatch()
+		defer b.Close()
+		require.NoError(t, f(b))
+
+		str, err := print.DecodeWriteBatch(b.Repr())
+		require.NoError(t, err)
+		out += fmt.Sprintf(">> %s\n%s", name, str)
+
+		require.NoError(t, b.Commit(false /* sync */))
+	}
+
+	id := roachpb.FullReplicaID{RangeID: 123, ReplicaID: 4}
+	write("create", func(w storage.Writer) error { return createReplica(&s, w, id) })
+	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, 200) })
+
+	// TODO(pav-kv): the trailing \n in DecodeWriteBatch is duplicated with
+	// recursion. Remove it, and let the caller handle new lines.
+	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++
+	}
+	require.NoError(t, iter.Error())
+	require.Equal(t, 4, count)
+}
+
+type store struct {
+	eng storage.Engine
+	seq Seq
+}
+
+func createReplica(s *store, w storage.Writer, id roachpb.FullReplicaID) error {
+	b := s.eng.NewWriteBatch()
+	defer b.Close()
+	if err := writeStateMachine(b, "state-machine-key", "state"); err != nil {
+		return err
+	}
+	return Write(w, s.seq.Next(1), wagpb.Node{
+		Addr:     wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: 0},
+		Type:     wagpb.NodeType_NodeCreate,
+		Mutation: wagpb.Mutation{Batch: b.Repr()},
+	})
+}
+
+func initReplica(s *store, w storage.Writer, id roachpb.FullReplicaID, index uint64) error {
+	return Write(w, s.seq.Next(1), wagpb.Node{
+		Addr: wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: kvpb.RaftIndex(index)},
+		Type: wagpb.NodeType_NodeSnap,
+		Mutation: wagpb.Mutation{Ingestion: &wagpb.Ingestion{
+			SSTs: []string{"tmp/1.sst", "tmp/2.sst"},
+		}},
+	})
+}
+
+func splitReplica(s *store, w storage.Writer, id roachpb.FullReplicaID, index uint64) error {
+	b := s.eng.NewWriteBatch()
+	defer b.Close()
+	if err := writeStateMachine(b, "lhs-key", "lhs-state"); err != nil {
+		return err
+	} else if err := writeStateMachine(b, "rhs-key", "rhs-state"); err != nil {
+		return err
+	}
+
+	seq := s.seq.Next(2)
+	if err := Write(w, seq, wagpb.Node{
+		Addr: wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: kvpb.RaftIndex(index - 1)},
+		Type: wagpb.NodeType_NodeApply,
+	}); err != nil {
+		return err
+	}
+	return Write(w, seq+1, wagpb.Node{
+		Addr:     wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: kvpb.RaftIndex(index)},
+		Type:     wagpb.NodeType_NodeSplit,
+		Mutation: wagpb.Mutation{Batch: b.Repr()},
+		Create:   567, // the RHS range ID
+	})
+}
+
+func writeStateMachine(w storage.Writer, k, v string) error {
+	return w.PutUnversioned(roachpb.Key(k), []byte(v))
+}

pkg/kv/kvserver/kvstorage/wag/testdata/TestWrite.txt

@@ -0,0 +1,13 @@
+echo
+----
+>> create
+Put: 0,0 /Local/Store/wag/1 (0x01737761676e000000000000000100): NodeCreate r123/4:0
+> Put: 0,0 "state-machine-key" (0x73746174652d6d616368696e652d6b657900): "state"
+>> init
+Put: 0,0 /Local/Store/wag/2 (0x01737761676e000000000000000200): NodeSnap r123/4:10
+ingestion: SSTs:"tmp/1.sst" SSTs:"tmp/2.sst" 
+>> split
+Put: 0,0 /Local/Store/wag/3 (0x01737761676e000000000000000300): NodeApply r123/4:199
+Put: 0,0 /Local/Store/wag/4 (0x01737761676e000000000000000400): NodeSplit r123/4:200 create:567
+> Put: 0,0 "lhs-key" (0x6c68732d6b657900): "lhs-state"
+> Put: 0,0 "rhs-key" (0x7268732d6b657900): "rhs-state"

pkg/kv/kvserver/kvstorage/wag/wagpb/BUILD.bazel

@@ -21,6 +21,7 @@ go_proto_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/kv/kvserver/kvserverpb",
+        "//pkg/roachpb",  # keep
         "@com_github_gogo_protobuf//gogoproto",
     ],
 )
@@ -31,5 +32,8 @@ go_library(
     embed = [":wagpb_go_proto"],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag/wagpb",
     visibility = ["//visibility:public"],
-    deps = ["//pkg/kv/kvpb"],
+    deps = [
+        "//pkg/kv/kvpb",
+        "@com_github_cockroachdb_redact//:redact",
+    ],
 )