Skip to content

Commit 67d9183

Browse files
pav-kvpav-kv
committed
wag: add test for WAG writing
Epic: none Release note: none
1 parent 2536417 commit 67d9183

File tree

5 files changed

+136
-1
lines changed

5 files changed

+136
-1
lines changed

‎pkg/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,7 @@ ALL_TESTS = [
264264
"//pkg/kv/kvserver/kvflowcontrol/rac2:rac2_test",
265265
"//pkg/kv/kvserver/kvflowcontrol/replica_rac2:replica_rac2_test",
266266
"//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv_test",
267+
"//pkg/kv/kvserver/kvstorage/wag:wag_test",
267268
"//pkg/kv/kvserver/kvstorage:kvstorage_test",
268269
"//pkg/kv/kvserver/leases:leases_test",
269270
"//pkg/kv/kvserver/liveness:liveness_test",
@@ -1562,6 +1563,7 @@ GO_TARGETS = [
15621563
"//pkg/kv/kvserver/kvstorage/snaprecv:snaprecv_test",
15631564
"//pkg/kv/kvserver/kvstorage/wag/wagpb:wagpb",
15641565
"//pkg/kv/kvserver/kvstorage/wag:wag",
1566+
"//pkg/kv/kvserver/kvstorage/wag:wag_test",
15651567
"//pkg/kv/kvserver/kvstorage:kvstorage",
15661568
"//pkg/kv/kvserver/kvstorage:kvstorage_test",
15671569
"//pkg/kv/kvserver/leases:leases",

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
load("@io_bazel_rules_go//go:def.bzl", "go_library")
1+
load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
22

33
go_library(
44
name = "wag",
@@ -11,3 +11,21 @@ go_library(
1111
"//pkg/storage",
1212
],
1313
)
14+
15+
go_test(
16+
name = "wag_test",
17+
srcs = ["store_test.go"],
18+
data = glob(["testdata/**"]),
19+
embed = [":wag"],
20+
deps = [
21+
"//pkg/kv/kvpb",
22+
"//pkg/kv/kvserver/kvstorage/wag/wagpb",
23+
"//pkg/kv/kvserver/print",
24+
"//pkg/roachpb",
25+
"//pkg/storage",
26+
"//pkg/testutils/echotest",
27+
"//pkg/util/leaktest",
28+
"//pkg/util/log",
29+
"@com_github_stretchr_testify//require",
30+
],
31+
)

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

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
// Copyright 2025 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the CockroachDB Software License
4+
// included in the /LICENSE file.
5+
6+
package wag
7+
8+
import (
9+
"fmt"
10+
"path/filepath"
11+
"testing"
12+
13+
"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
14+
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage/wag/wagpb"
15+
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/print"
16+
"github.com/cockroachdb/cockroach/pkg/roachpb"
17+
"github.com/cockroachdb/cockroach/pkg/storage"
18+
"github.com/cockroachdb/cockroach/pkg/testutils/echotest"
19+
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
20+
"github.com/cockroachdb/cockroach/pkg/util/log"
21+
"github.com/stretchr/testify/require"
22+
)
23+
24+
func TestWrite(t *testing.T) {
25+
defer leaktest.AfterTest(t)()
26+
defer log.Scope(t).Close(t)
27+
28+
eng := storage.NewDefaultInMemForTesting()
29+
defer eng.Close()
30+
s := store{eng: eng}
31+
32+
var out string
33+
write := func(name string, f func(w storage.Writer) error) {
34+
b := eng.NewWriteBatch()
35+
defer b.Close()
36+
require.NoError(t, f(b))
37+
38+
str, err := print.DecodeWriteBatch(b.Repr())
39+
require.NoError(t, err)
40+
out += fmt.Sprintf(">> %s\n%s", name, str)
41+
}
42+
43+
id := roachpb.FullReplicaID{RangeID: 123, ReplicaID: 4}
44+
write("create", func(w storage.Writer) error { return createReplica(&s, w, id) })
45+
write("init", func(w storage.Writer) error { return initReplica(&s, w, id, 10) })
46+
write("split", func(w storage.Writer) error { return splitReplica(&s, w, id, 200) })
47+
48+
echotest.Require(t, out, filepath.Join("testdata", t.Name()+".txt"))
49+
}
50+
51+
type store struct {
52+
eng storage.Engine
53+
seq Seq
54+
}
55+
56+
func createReplica(s *store, w storage.Writer, id roachpb.FullReplicaID) error {
57+
b := s.eng.NewWriteBatch()
58+
defer b.Close()
59+
if err := writeStateMachine(b, "state-machine-key", "state"); err != nil {
60+
return err
61+
}
62+
return Write(w, s.seq.Next(1), wagpb.Node{
63+
Addr: wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: 0},
64+
Type: wagpb.NodeType_NodeCreate,
65+
Mutation: wagpb.Mutation{Batch: b.Repr()},
66+
})
67+
}
68+
69+
func initReplica(s *store, w storage.Writer, id roachpb.FullReplicaID, index uint64) error {
70+
return Write(w, s.seq.Next(1), wagpb.Node{
71+
Addr: wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: kvpb.RaftIndex(index)},
72+
Type: wagpb.NodeType_NodeSnap,
73+
Mutation: wagpb.Mutation{Ingestion: &wagpb.Ingestion{
74+
SSTs: []string{"tmp/1.sst", "tmp/2.sst"},
75+
}},
76+
})
77+
}
78+
79+
func splitReplica(s *store, w storage.Writer, id roachpb.FullReplicaID, index uint64) error {
80+
b := s.eng.NewWriteBatch()
81+
defer b.Close()
82+
if err := writeStateMachine(b, "lhs-key", "lhs-state"); err != nil {
83+
return err
84+
} else if err := writeStateMachine(b, "rhs-key", "rhs-state"); err != nil {
85+
return err
86+
}
87+
88+
seq := s.seq.Next(2)
89+
if err := Write(w, seq, wagpb.Node{
90+
Addr: wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: kvpb.RaftIndex(index - 1)},
91+
Type: wagpb.NodeType_NodeApply,
92+
}); err != nil {
93+
return err
94+
}
95+
return Write(w, seq+1, wagpb.Node{
96+
Addr: wagpb.Addr{RangeID: id.RangeID, ReplicaID: id.ReplicaID, Index: kvpb.RaftIndex(index)},
97+
Type: wagpb.NodeType_NodeSplit,
98+
Mutation: wagpb.Mutation{Batch: b.Repr()},
99+
Create: 567, // the RHS range ID
100+
})
101+
}
102+
103+
func writeStateMachine(w storage.Writer, k, v string) error {
104+
return w.PutUnversioned(roachpb.Key(k), []byte(v))
105+
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
echo
2+
----
3+
>> create
4+
Put: 0,0 /Local/Store"wagn\x00\x00\x00\x00\x00\x00\x00\x01" (0x01737761676e000000000000000100): "\n\x04\b{\x10\x04\x10\x01\x1a(\n&\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x01\x12state-machine-key\x00\x05state"
5+
>> init
6+
Put: 0,0 /Local/Store"wagn\x00\x00\x00\x00\x00\x00\x00\x02" (0x01737761676e000000000000000200): "\n\x06\b{\x10\x04 \n\x10\x02\x1a\x18\x12\x16\n\ttmp/1.sst\n\ttmp/2.sst"
7+
>> split
8+
Put: 0,0 /Local/Store"wagn\x00\x00\x00\x00\x00\x00\x00\x03" (0x01737761676e000000000000000300): "\n\a\b{\x10\x04 \xc7\x01\x10\x03\x1a\x00"
9+
Put: 0,0 /Local/Store"wagn\x00\x00\x00\x00\x00\x00\x00\x04" (0x01737761676e000000000000000400): "\n\a\b{\x10\x04 \xc8\x01\x10\x04\x1a6\n4\x00\x00\x00\x00\x00\x00\x00\x00\x02\x00\x00\x00\x01\blhs-key\x00\tlhs-state\x01\brhs-key\x00\trhs-state \xb7\x04"

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ go_proto_library(
2121
visibility = ["//visibility:public"],
2222
deps = [
2323
"//pkg/kv/kvserver/kvserverpb",
24+
"//pkg/roachpb", # keep
2425
"@com_github_gogo_protobuf//gogoproto",
2526
],
2627
)

0 commit comments

Comments
 (0)