kvserver: add split trigger support to TestReplicaLifecycleDataDriven

This patch adds a few directives to construct and evaluate split
triggers.

We then use these directives to demonstrate that leases are
correctly copied over from the LHS to the RHS. The interesting case
is around leader leases, where if the LHS has a leader lease, the RHS
gets an expiration based lease. For the other two lease types, the RHS's
lease type stays the same.

Epic: none

Release note: None

Author: arulajmani

pkg/kv/kvserver/batcheval/cmd_end_transaction.go

@@ -1407,6 +1407,9 @@ func splitTriggerHelper(
 	// This avoids running the consistency checker on the RHS immediately after
 	// the split.
 	lastTS := hlc.Timestamp{}
+	// TODO(arul): instead of fetching the consistency checker timestamp here
+	// like this, we should instead pass it using the SplitTriggerHelperInput to
+	// make it easier to test.
 	if _, err := storage.MVCCGetProto(ctx, batch,
 		keys.QueueLastProcessedKey(split.LeftDesc.StartKey, "consistencyChecker"),
 		hlc.Timestamp{}, &lastTS, storage.MVCCGetOptions{}); err != nil {

pkg/kv/kvserver/print/debug_print.go

@@ -397,6 +397,7 @@ func tryTxn(kv storage.MVCCKeyValue) (string, error) {
 	if err := maybeUnmarshalInline(kv.Value, &txn); err != nil {
 		return "", err
 	}
+	// TODO(arul): investigate why we need this new line here.
 	return txn.String() + "\n", nil
 }
 

pkg/kv/kvserver/replica_lifecycle_datadriven_test.go

@@ -9,21 +9,28 @@
 package kvserver
 
 import (
+	"cmp"
 	"context"
 	"fmt"
 	"strings"
 	"testing"
+	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/abortspan"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/print"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/storage"
+	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/testutils/dd"
+	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/datadriven"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
@@ -50,10 +57,36 @@ import (
 // Creates a replica on n1/s1 for the specified range ID. The created replica
 // may be initialized or uninitialized.
 //
-// print-range-state
+// create-split range-id=<int> split-key=<key>
 // ----
 //
-// Prints the current range state in the test context.
+//	Creates a split for the specified range at the given split key, which
+//	entails creating a SplitTrigger with both the LHS and RHS descriptors.
+//	Much like how things work in CRDB, the LHS descriptor is created by
+//	narrowing the original range and a new range descriptor is created for
+//	the RHS with the same replica set.
+//
+// set-lease range-id=<int> replica=<int> [lease-type=leader-lease|epoch|expiration]
+// ----
+//
+//	Sets the lease for the specified range to the supplied replica. Note that
+//	the replica parameter specifies NodeIDs, not to be confused with
+//	ReplicaIDs. By default, the lease is of the leader-lease variety, but
+//	this may be overriden to an epoch or expiration based lease by using the
+//	lease-type parameter. For now, we treat the associated lease metadata as
+//	uninteresting.
+//
+// run-split-trigger range-id=<int>
+// ----
+//
+//	Executes the split trigger for the specified range on n1.
+//
+// print-range-state [sort-keys=<bool>]
+// ----
+//
+// Prints the current range state in the test context. By default, ranges are
+// sorted by range ID. If sort-keys is set to true, ranges are sorted by their
+// descriptor's start key instead.
 func TestReplicaLifecycleDataDriven(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
@@ -116,7 +149,7 @@ func TestReplicaLifecycleDataDriven(t *testing.T) {
 				if rs.replica != nil {
 					return errors.New("initialized replica already exists on n1/s1").Error()
 				}
-				repl := rs.getReplicaDescriptor(t)
+				repl := rs.getReplicaDescriptor(t, roachpb.NodeID(1))
 
 				batch := tc.storage.NewBatch()
 				defer batch.Close()
@@ -148,14 +181,130 @@ func TestReplicaLifecycleDataDriven(t *testing.T) {
 				require.NoError(t, err, "error committing batch")
 				return sb.String()
 
+			case "create-split":
+				rangeID := dd.ScanArg[int](t, d, "range-id")
+				splitKey := dd.ScanArg[string](t, d, "split-key")
+				rs := tc.mustGetRangeState(t, roachpb.RangeID(rangeID))
+				desc := rs.desc
+				require.True(
+					t,
+					roachpb.RKey(splitKey).Compare(desc.StartKey) > 0 &&
+						roachpb.RKey(splitKey).Compare(desc.EndKey) < 0,
+					"split key not within range",
+				)
+				leftDesc := desc
+				leftDesc.EndKey = roachpb.RKey(splitKey)
+				rightDesc := desc
+				rightDesc.RangeID = tc.nextRangeID
+				tc.nextRangeID++
+				rightDesc.StartKey = roachpb.RKey(splitKey)
+				split := &roachpb.SplitTrigger{
+					LeftDesc:  leftDesc,
+					RightDesc: rightDesc,
+				}
+				tc.splits[roachpb.RangeID(rangeID)] = split
+				return "ok"
+
+			case "set-lease":
+				rangeID := dd.ScanArg[int](t, d, "range-id")
+				replicaNodeID := dd.ScanArg[int](t, d, "replica")
+				leaseType := "leader-lease" // default to a leader-lease
+				if d.HasArg("lease-type") {
+					leaseType = dd.ScanArg[string](t, d, "lease-type")
+				}
+				rs := tc.mustGetRangeState(t, roachpb.RangeID(rangeID))
+				targetReplica := rs.getReplicaDescriptor(t, roachpb.NodeID(replicaNodeID))
+				// NB: The details of the lease are not important to the test;
+				// only the type is.
+				var lease roachpb.Lease
+				switch leaseType {
+				case "leader-lease":
+					lease = roachpb.Lease{
+						Replica:       *targetReplica,
+						Term:          10,
+						MinExpiration: hlc.Timestamp{WallTime: 100},
+					}
+				case "epoch":
+					lease = roachpb.Lease{
+						Replica: *targetReplica,
+						Epoch:   20,
+					}
+				case "expiration":
+					lease = roachpb.Lease{
+						Replica:    *targetReplica,
+						Expiration: &hlc.Timestamp{WallTime: 300},
+					}
+				default:
+					t.Fatalf("unknown lease type: %s", leaseType)
+				}
+				rs.lease = lease
+				return "ok"
+
+			case "run-split-trigger":
+				rangeID := dd.ScanArg[int](t, d, "range-id")
+				split, ok := tc.splits[roachpb.RangeID(rangeID)]
+				require.True(t, ok, "split trigger not found for range-id %d", rangeID)
+				rs := tc.mustGetRangeState(t, roachpb.RangeID(rangeID))
+				desc := rs.desc
+				batch := tc.storage.NewBatch()
+				defer batch.Close()
+
+				rec := (&batcheval.MockEvalCtx{
+					ClusterSettings:        tc.st,
+					Desc:                   &desc,
+					Clock:                  tc.clock,
+					AbortSpan:              rs.abortspan,
+					LastReplicaGCTimestamp: rs.lastGCTimestamp,
+					RangeLeaseDuration:     tc.rangeLeaseDuration,
+				}).EvalContext()
+
+				in := batcheval.SplitTriggerHelperInput{
+					LeftLease:      rs.lease,
+					GCThreshold:    &rs.gcThreshold,
+					GCHint:         &rs.gcHint,
+					ReplicaVersion: rs.version,
+				}
+				// Actually run the split trigger.
+				_, _, err := batcheval.TestingSplitTrigger(
+					ctx, rec, batch /* bothDeltaMS */, enginepb.MVCCStats{}, split, in, hlc.Timestamp{},
+				)
+				require.NoError(t, err)
+
+				// Update the test context's notion of the range state after the
+				// split.
+				tc.updatePostSplitRangeState(t, ctx, batch, roachpb.RangeID(rangeID), split)
+				// Print the state of the batch (all keys/values written as part
+				// of the split trigger).
+				output, err := print.DecodeWriteBatch(batch.Repr())
+				require.NoError(t, err)
+				// Commit the batch.
+				err = batch.Commit(true)
+				require.NoError(t, err, "error committing batch")
+				// TODO(arul): There are double lines in the output (see tryTxn
+				// in debug_print.go) that we need to strip out for the benefit
+				// of the datadriven test driver. Until that TODO is addressed,
+				// we manually split things out here.
+				return strings.ReplaceAll(output, "\n\n", "\n")
+
 			case "print-range-state":
 				var sb strings.Builder
 				if len(tc.ranges) == 0 {
 					return "no ranges in test context"
 				}
-				// Sort by range IDs for consistent output.
+
+				sortByKeys := false
+				if d.HasArg("sort-keys") {
+					sortByKeys = dd.ScanArg[bool](t, d, "sort-keys")
+				}
+
 				rangeIDs := maps.Keys(tc.ranges)
-				slices.Sort(rangeIDs)
+				slices.SortFunc(rangeIDs, func(a, b roachpb.RangeID) int {
+					if sortByKeys {
+						return tc.ranges[a].desc.StartKey.Compare(tc.ranges[b].desc.StartKey)
+					}
+					// Else sort by range IDs for consistent output.
+					return cmp.Compare(a, b)
+				})
 
 				for _, rangeID := range rangeIDs {
 					rs := tc.ranges[rangeID]
@@ -172,9 +321,14 @@ func TestReplicaLifecycleDataDriven(t *testing.T) {
 
 // rangeState represents the state of a single range in the test context.
 type rangeState struct {
-	desc    roachpb.RangeDescriptor
-	version roachpb.Version
-	replica *replicaInfo // replica on n1/s1.
+	desc            roachpb.RangeDescriptor
+	version         roachpb.Version
+	lease           roachpb.Lease
+	gcThreshold     hlc.Timestamp
+	gcHint          roachpb.GCHint
+	abortspan       *abortspan.AbortSpan
+	lastGCTimestamp hlc.Timestamp
+	replica         *replicaInfo // replica on n1/s1.
 }
 
 // replicaInfo contains the basic info about a replica, used for managing its
@@ -188,20 +342,30 @@ type replicaInfo struct {
 // testCtx is a single test's context. It tracks the state of all ranges and any
 // intermediate steps when performing replica lifecycle events.
 type testCtx struct {
-	ranges      map[roachpb.RangeID]*rangeState
+	st                 *cluster.Settings
+	clock              *hlc.Clock
+	rangeLeaseDuration time.Duration
+
 	nextRangeID roachpb.RangeID // monotonically-increasing rangeID
-	st          *cluster.Settings
+	ranges      map[roachpb.RangeID]*rangeState
+	splits      map[roachpb.RangeID]*roachpb.SplitTrigger
 	// The storage engine corresponds to a single store, (n1, s1).
 	storage storage.Engine
 }
 
 // newTestCtx constructs and returns a new testCtx.
 func newTestCtx() *testCtx {
 	st := cluster.MakeTestingClusterSettings()
+	manual := timeutil.NewManualTime(timeutil.Unix(0, 10))
+	clock := hlc.NewClockForTesting(manual)
 	return &testCtx{
-		ranges:      make(map[roachpb.RangeID]*rangeState),
+		st:                 st,
+		clock:              clock,
+		rangeLeaseDuration: 99 * time.Nanosecond,
+
 		nextRangeID: 1,
-		st:          st,
+		ranges:      make(map[roachpb.RangeID]*rangeState),
+		splits:      make(map[roachpb.RangeID]*roachpb.SplitTrigger),
 		storage:     storage.NewDefaultInMemForTesting(),
 	}
 }
@@ -213,9 +377,16 @@ func (tc *testCtx) close() {
 
 // newRangeState constructs a new rangeState for the supplied descriptor.
 func newRangeState(desc roachpb.RangeDescriptor) *rangeState {
+	gcThreshold := hlc.Timestamp{WallTime: 4}
+	gcHint := roachpb.GCHint{GCTimestamp: gcThreshold}
+
 	return &rangeState{
-		desc:    desc,
-		version: roachpb.Version{Major: 10, Minor: 8, Internal: 7}, // dummy version to avoid churn
+		desc:            desc,
+		version:         roachpb.Version{Major: 10, Minor: 8, Internal: 7}, // dummy version to avoid churn
+		gcThreshold:     gcThreshold,
+		gcHint:          gcHint,
+		abortspan:       abortspan.New(desc.RangeID),
+		lastGCTimestamp: hlc.Timestamp{},
 	}
 }
 
@@ -248,13 +419,42 @@ func (tc *testCtx) updatePostReplicaCreateState(
 	}
 }
 
-func (rs *rangeState) getReplicaDescriptor(t *testing.T) *roachpb.ReplicaDescriptor {
+// updatePostSplitRangeState updates the range state after a split.
+func (tc *testCtx) updatePostSplitRangeState(
+	t *testing.T,
+	ctx context.Context,
+	reader storage.Reader,
+	lhsRangeID roachpb.RangeID,
+	split *roachpb.SplitTrigger,
+) {
+	originalRangeState := tc.mustGetRangeState(t, lhsRangeID)
+	// The range ID should not change for LHS since it's the same range.
+	require.Equal(t, lhsRangeID, split.LeftDesc.RangeID)
+	// Update LHS by just updating the descriptor.
+	originalRangeState.desc = split.LeftDesc
+	tc.ranges[lhsRangeID] = originalRangeState
+	rhsRangeState := newRangeState(split.RightDesc)
+	// Create RHS range state by reading from the reader.
+	rhsSl := kvstorage.MakeStateLoader(split.RightDesc.RangeID)
+	rhsState, err := rhsSl.Load(ctx, reader, &split.RightDesc)
+	require.NoError(t, err)
+	rhsRangeState.lease = *rhsState.Lease
+	rhsRangeState.gcThreshold = *rhsState.GCThreshold
+	rhsRangeState.gcHint = *rhsState.GCHint
+	rhsRangeState.version = *rhsState.Version
+
+	tc.ranges[split.RightDesc.RangeID] = rhsRangeState
+}
+
+func (rs *rangeState) getReplicaDescriptor(
+	t *testing.T, nodeID roachpb.NodeID,
+) *roachpb.ReplicaDescriptor {
 	for i, repl := range rs.desc.InternalReplicas {
-		if repl.NodeID == roachpb.NodeID(1) {
+		if repl.NodeID == nodeID {
 			return &rs.desc.InternalReplicas[i]
 		}
 	}
-	t.Fatal("replica not found")
+	t.Fatalf("replica with NodeID %d not found in range descriptor", nodeID)
 	return nil // unreachable
 }
 
@@ -264,6 +464,9 @@ func (rs *rangeState) String() string {
 	if rs.replica != nil {
 		sb.WriteString(fmt.Sprintf("\n		replica (n1/s1): %s", rs.replica))
 	}
+	if (rs.lease != roachpb.Lease{}) {
+		sb.WriteString(fmt.Sprintf("\n		lease: %s", rs.lease))
+	}
 	return sb.String()
 }