kvstorage: add datadriven harness

See cockroachdb#93310. This is
also the beginning of
cockroachdb#93247.

Epic: CRDB-220
Release note: None

pkg/kv/kvserver/kvstorage/BUILD.bazel

@@ -27,16 +27,29 @@ go_library(
 
 go_test(
     name = "kvstorage_test",
-    srcs = ["cluster_version_test.go"],
+    srcs = [
+        "cluster_version_test.go",
+        "datadriven_test.go",
+    ],
     args = ["-test.timeout=295s"],
+    data = glob(["testdata/**"]),
     embed = [":kvstorage"],
     deps = [
         "//pkg/clusterversion",
+        "//pkg/keys",
+        "//pkg/kv/kvserver/logstore",
         "//pkg/roachpb",
         "//pkg/storage",
         "//pkg/testutils",
+        "//pkg/testutils/datapathutils",
+        "//pkg/util/hlc",
         "//pkg/util/leaktest",
         "//pkg/util/log",
+        "//pkg/util/tracing",
+        "//pkg/util/uuid",
+        "@com_github_cockroachdb_datadriven//:datadriven",
+        "@com_github_stretchr_testify//require",
+        "@io_etcd_go_raft_v3//raftpb",
     ],
 )
 

pkg/kv/kvserver/kvstorage/datadriven_test.go

@@ -0,0 +1,208 @@
+// Copyright 2023 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package kvstorage
+
+import (
+	"context"
+	"fmt"
+	"regexp"
+	"sort"
+	"strings"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/clusterversion"
+	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/tracing"
+	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/cockroachdb/datadriven"
+	"github.com/stretchr/testify/require"
+	"go.etcd.io/raft/v3/raftpb"
+)
+
+type env struct {
+	eng storage.Engine
+	tr  *tracing.Tracer
+}
+
+func newEnv(t *testing.T) *env {
+	ctx := context.Background()
+	eng := storage.NewDefaultInMemForTesting()
+	// TODO(tbg): ideally this would do full bootstrap, which requires
+	// moving a lot more code from kvserver. But then we could unit test
+	// all of it with the datadriven harness!
+	require.NoError(t, WriteClusterVersion(ctx, eng, clusterversion.TestingClusterVersion))
+	require.NoError(t, InitEngine(ctx, eng, roachpb.StoreIdent{
+		ClusterID: uuid.FastMakeV4(),
+		NodeID:    1,
+		StoreID:   1,
+	}))
+	tr := tracing.NewTracer()
+	tr.SetRedactable(true)
+	return &env{
+		eng: eng,
+		tr:  tr,
+	}
+}
+
+func (e *env) close() {
+	e.eng.Close()
+	e.tr.Close()
+}
+
+func (e *env) handleNewReplica(
+	t *testing.T,
+	ctx context.Context,
+	id storage.FullReplicaID,
+	skipRaftReplicaID bool,
+	k, ek roachpb.RKey,
+) *roachpb.RangeDescriptor {
+	sl := logstore.NewStateLoader(id.RangeID)
+	require.NoError(t, sl.SetHardState(ctx, e.eng, raftpb.HardState{}))
+	if !skipRaftReplicaID && id.ReplicaID != 0 {
+		require.NoError(t, sl.SetRaftReplicaID(ctx, e.eng, id.ReplicaID))
+	}
+	if len(ek) == 0 {
+		return nil
+	}
+	desc := &roachpb.RangeDescriptor{
+		RangeID:  id.RangeID,
+		StartKey: keys.MustAddr(roachpb.Key(k)),
+		EndKey:   keys.MustAddr(roachpb.Key(ek)),
+		InternalReplicas: []roachpb.ReplicaDescriptor{{
+			NodeID:    1,
+			StoreID:   1,
+			ReplicaID: id.ReplicaID,
+		}},
+		NextReplicaID: id.ReplicaID + 1,
+	}
+	var v roachpb.Value
+	require.NoError(t, v.SetProto(desc))
+	ts := hlc.Timestamp{WallTime: 123}
+	require.NoError(t, e.eng.PutMVCC(storage.MVCCKey{
+		Key:       keys.RangeDescriptorKey(desc.StartKey),
+		Timestamp: ts,
+	}, storage.MVCCValue{Value: v}))
+	return desc
+}
+
+func TestDataDriven(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	reStripFileLinePrefix := regexp.MustCompile(`^[^ ]+ `)
+
+	datadriven.Walk(t, datapathutils.TestDataPath(t), func(t *testing.T, path string) {
+		e := newEnv(t)
+		defer e.close()
+		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) (output string) {
+			ctx, finishAndGet := tracing.ContextWithRecordingSpan(context.Background(), e.tr, path)
+			// This method prints all output to `buf`.
+			var buf strings.Builder
+			var printTrace bool // if true, trace printed to buf on return
+			if d.HasArg("trace") {
+				d.ScanArgs(t, "trace", &printTrace)
+			}
+
+			defer func() {
+				if r := recover(); r != nil {
+					fmt.Fprintln(&buf, r)
+				}
+				rec := finishAndGet()[0]
+				for _, l := range rec.Logs {
+					if !printTrace || !strings.Contains(string(l.Message), "kvstorage") {
+						continue
+					}
+
+					fmt.Fprintln(&buf, reStripFileLinePrefix.ReplaceAllString(string(l.Message), ``))
+				}
+				if buf.Len() == 0 {
+					fmt.Fprintln(&buf, "ok")
+				}
+				output = buf.String()
+			}()
+
+			switch d.Cmd {
+			case "new-replica":
+				var rangeID int
+				d.ScanArgs(t, "range-id", &rangeID)
+				var replicaID int
+				if d.HasArg("replica-id") { // optional to allow making incomplete state
+					d.ScanArgs(t, "replica-id", &replicaID)
+				}
+				var k string
+				if d.HasArg("k") {
+					d.ScanArgs(t, "k", &k)
+				}
+				var ek string
+				if d.HasArg("ek") {
+					d.ScanArgs(t, "ek", &ek)
+				}
+				var skipRaftReplicaID bool
+				if d.HasArg("skip-raft-replica-id") {
+					d.ScanArgs(t, "skip-raft-replica-id", &skipRaftReplicaID)
+				}
+				if desc := e.handleNewReplica(t, ctx,
+					storage.FullReplicaID{RangeID: roachpb.RangeID(rangeID), ReplicaID: roachpb.ReplicaID(replicaID)},
+					skipRaftReplicaID, keys.MustAddr(roachpb.Key(k)), keys.MustAddr(roachpb.Key(ek)),
+				); desc != nil {
+					fmt.Fprintln(&buf, desc)
+				}
+			case "list-range-ids":
+				m, err := loadFullReplicaIDsFromDisk(ctx, e.eng)
+				require.NoError(t, err)
+				var sl []storage.FullReplicaID
+				for id := range m {
+					sl = append(sl, id)
+				}
+				sort.Slice(sl, func(i, j int) bool {
+					return sl[i].RangeID < sl[j].RangeID
+				})
+				for _, id := range sl {
+					fmt.Fprintf(&buf, "r%d/%d\n", id.RangeID, id.ReplicaID)
+				}
+			case "load-and-reconcile":
+				rs, err := LoadAndReconcileReplicas(ctx, e.eng)
+				if err != nil {
+					fmt.Fprintln(&buf, err)
+					break
+				}
+				var merged []storage.FullReplicaID
+				for id := range rs.Initialized {
+					merged = append(merged, id)
+				}
+				for id := range rs.Uninitialized {
+					merged = append(merged, id)
+				}
+				sort.Slice(merged, func(i, j int) bool {
+					return merged[i].RangeID < merged[j].RangeID
+				})
+				for _, id := range merged {
+					fmt.Fprintf(&buf, "r%d/%d: ", id.RangeID, id.ReplicaID)
+					if desc := rs.Initialized[id]; desc != nil {
+						fmt.Fprint(&buf, desc)
+					} else {
+						fmt.Fprintf(&buf, "uninitialized")
+					}
+					fmt.Fprintln(&buf)
+				}
+			default:
+				t.Fatalf("unknown command %s", d.Cmd)
+			}
+			return "" // defer will do it
+		})
+	})
+
+}

pkg/kv/kvserver/kvstorage/testdata/assert_duplicate_replica

@@ -0,0 +1,14 @@
+new-replica range-id=1 replica-id=10 k=a ek=c
+----
+r1:{a-c} [(n1,s1):10, next=11, gen=0]
+
+new-replica range-id=1 replica-id=20 k=c ek=d
+----
+r1:{c-d} [(n1,s1):20, next=21, gen=0]
+
+# When we created replica-id=20, it clobbered replica-id=10
+# and so we see the descriptor but cannot match it up to a
+# raft state.
+load-and-reconcile
+----
+r1:{a-c} [(n1,s1):10, next=11, gen=0] has no persisted replicaID

pkg/kv/kvserver/kvstorage/testdata/assert_duplicate_replica_2

@@ -0,0 +1,14 @@
+# Variant of the original test but there's overlap between an
+# initialized and uninitialized copy. The result is exactly
+# the same.
+new-replica range-id=1 replica-id=10 k=a ek=c
+----
+r1:{a-c} [(n1,s1):10, next=11, gen=0]
+
+new-replica range-id=1 replica-id=20
+----
+ok
+
+load-and-reconcile
+----
+r1:{a-c} [(n1,s1):10, next=11, gen=0] has no persisted replicaID

pkg/kv/kvserver/kvstorage/testdata/assert_overlapping_replica

@@ -0,0 +1,12 @@
+new-replica range-id=1 replica-id=10 k=a ek=c
+----
+r1:{a-c} [(n1,s1):10, next=11, gen=0]
+
+new-replica range-id=2 replica-id=20 k=b ek=d
+----
+r2:{b-d} [(n1,s1):20, next=21, gen=0]
+
+load-and-reconcile
+----
+r1/10: r1:{a-c} [(n1,s1):10, next=11, gen=0]
+r2/20: r2:{b-d} [(n1,s1):20, next=21, gen=0]

pkg/kv/kvserver/kvstorage/testdata/init

@@ -0,0 +1,39 @@
+# Uninitialized deprecated replica: doesn't have a RaftReplicaID, it's just a
+# HardState. We expect this to be removed by load-and-reconcile.
+new-replica range-id=5
+----
+ok
+
+# Uninitialized replica. Expect this to stay around and be returned as such.
+new-replica range-id=6 replica-id=60
+----
+ok
+
+# Initialized deprecated replica: no RaftReplicaID. Expect ID to be backfilled.
+new-replica range-id=7 replica-id=70 k=a ek=c skip-raft-replica-id=true
+----
+r7:{a-c} [(n1,s1):70, next=71, gen=0]
+
+# Initialized replica without a need for a backfill.
+new-replica range-id=8 replica-id=80 k=c ek=f
+----
+r8:{c-f} [(n1,s1):80, next=81, gen=0]
+
+list-range-ids
+----
+r6/60
+r8/80
+
+# Loading the replicas returns only the ones that
+# had a ReplicaID.
+load-and-reconcile trace=true
+----
+r7:{a-c} [(n1,s1):70, next=71, gen=0] has no persisted replicaID
+beginning range descriptor iteration
+iterated over 2 keys to find 2 range descriptors (by suffix: map[‹rdsc›:2])
+
+# r5 was removed, and r7 had its RaftReplicaID backfilled.
+list-range-ids
+----
+r6/60
+r8/80