Merge #155192 #155284 #155412

155192: backfill: split merge memory accounting mutex acquisitions r=mw5h a=mw5h

Previously, constructMergeBatch() would get the mutex on memory accounting once per batch to minimize the number of atomics needed to perform memory accounting. With the addition of vector indexing, we now need to perform KV operations in order to construct batches, which can lead to deadlocks between KV and the memory accounting mutex.

This patch splits the mutex acquisition into two acquisitions, releasing the mutex in between. Memory used is accumulated as merged entries are built and then accounting is notified for the entire batch. This has the effect of making the accounting lag behind allocation a bit, but that seems preferable than paying the cost of a mutex acquisition for every row merged.

Fixes: #155190
Release note (bug fix): A potential deadlock during vector index creation has been corrected.

155284: backfill: add a separate control for number of vectors merged per batch r=mw5h a=mw5h

Add the 'bulkio.index_backfill.vector_merge_batch_size', which controls the number of rows to merge into vector indexes per transaction while the index is being created. This is analogous to the 'bulkio.index_backfill.merge_batch_size' setting, but it only applies when the target index is a vector index. By default, 3 vectors will be merged per transaction to reduce contention with fixup tasks.

Fixes: #155283
Release note (sql change): Added the
bulkio.index_backfill.vector_merge_batch_size cluster setting to control how many vectors to merge into a vector index per transaction during create operations. By default, this defaults to 3.

155412: kvserver: introduce test harness for replica lifecycle events  r=pav-kv a=arulajmani

First two commits from #155408.

This patch introduces the scaffolding for a harness to test various
replica lifecycle events in a data-driven manner. The harness works
from th context of a single node/store -- n1,s1. The intention is to
print all engine interactions of the replica on this store on various
replica lifecycle events. To kick things off, we add the ability to
create a replica -- both initialized and uninitialized.

Epic: none

Release note: None

Co-authored-by: Matt White <[email protected]>
Co-authored-by: Arul Ajmani <[email protected]>

Author: 3 people

pkg/kv/kvserver/BUILD.bazel

@@ -351,6 +351,7 @@ go_test(
         "replica_init_test.go",
         "replica_learner_test.go",
         "replica_lease_renewal_test.go",
+        "replica_lifecycle_datadriven_test.go",
         "replica_metrics_test.go",
         "replica_probe_test.go",
         "replica_proposal_bench_test.go",

pkg/kv/kvserver/replica_lifecycle_datadriven_test.go

@@ -0,0 +1,295 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+// TODO(arul): As this test suite evolves, see if it can be moved into the
+// kvstorage package instead.
+
+package kvserver
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+
+	"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/kv/kvserver/stateloader"
+	"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/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/datadriven"
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/exp/maps"
+	"golang.org/x/exp/slices"
+)
+
+// TestReplicaLifecycleDataDriven is intended to test the behaviour of various
+// replica lifecycle events, such as splits, merges, replica destruction, etc.
+// The test has a single storage engine that corresponds to n1/s1, and all batch
+// operations to storage are printed out. It uses the following format:
+//
+// create-descriptor start=<key> end=<key> replicas=[<int>,<int>,...]
+// ----
+//
+//	Creates a range descriptor with the specified start and end keys and
+//	optional replica list. The range ID is auto-assigned. If provided,
+//	replicas specify NodeIDs for replicas of the range. Note that ReplicaIDs
+//	are assigned incrementally starting from 1.
+//
+// create-replica range-id=<int> [initialized]
+// ----
+//
+// Creates a replica on n1/s1 for the specified range ID. The created replica
+// may be initialized or uninitialized.
+//
+// print-range-state
+// ----
+//
+// Prints the current range state in the test context.
+func TestReplicaLifecycleDataDriven(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	datadriven.Walk(t, "testdata/replica_lifecycle", func(t *testing.T, path string) {
+		tc := newTestCtx()
+		defer tc.close()
+		ctx := context.Background()
+		datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
+			switch d.Cmd {
+			case "create-descriptor":
+				var startKey, endKey string
+				d.ScanArgs(t, "start", &startKey)
+				d.ScanArgs(t, "end", &endKey)
+				var replicasStr string
+				d.ScanArgs(t, "replicas", &replicasStr)
+				replicaNodeIDs := parseReplicas(t, replicasStr)
+
+				rangeID := tc.nextRangeID
+				tc.nextRangeID++
+				var internalReplicas []roachpb.ReplicaDescriptor
+				for i, id := range replicaNodeIDs {
+					internalReplicas = append(internalReplicas, roachpb.ReplicaDescriptor{
+						ReplicaID: roachpb.ReplicaID(i + 1),
+						NodeID:    id,
+						StoreID:   roachpb.StoreID(id),
+						Type:      roachpb.VOTER_FULL,
+					})
+				}
+				desc := roachpb.RangeDescriptor{
+					RangeID:          rangeID,
+					StartKey:         roachpb.RKey(startKey),
+					EndKey:           roachpb.RKey(endKey),
+					InternalReplicas: internalReplicas,
+					NextReplicaID:    roachpb.ReplicaID(len(internalReplicas) + 1),
+				}
+				require.True(t, desc.StartKey.Compare(desc.EndKey) < 0)
+
+				// Ranges are expected to be non-overlapping. Before creating a
+				// new one, sanity check that we're not violating this property
+				// in the test context.
+				for existingRangeID, existingRS := range tc.ranges {
+					existingDesc := existingRS.desc
+					require.False(t, desc.StartKey.Compare(existingDesc.EndKey) < 0 &&
+						existingDesc.StartKey.Compare(desc.EndKey) < 0,
+						"descriptor overlaps with existing range %d [%s,%s)",
+						existingRangeID, existingDesc.StartKey, existingDesc.EndKey)
+				}
+
+				rs := newRangeState(desc)
+				tc.ranges[rangeID] = rs
+				return fmt.Sprintf("created descriptor: %v", desc)
+
+			case "create-replica":
+				var rangeID int
+				d.ScanArgs(t, "range-id", &rangeID)
+				rs := tc.mustGetRangeState(t, roachpb.RangeID(rangeID))
+				if rs.replica != nil {
+					return errors.New("initialized replica already exists on n1/s1").Error()
+				}
+				repl := rs.getReplicaDescriptor(t)
+
+				initialized := d.HasArg("initialized")
+
+				batch := tc.storage.NewBatch()
+				defer batch.Close()
+
+				if initialized {
+					err := stateloader.WriteInitialRangeState(ctx, batch, rs.desc, repl.ReplicaID, rs.version)
+					require.NoError(t, err)
+				} else {
+					err := kvstorage.CreateUninitializedReplica(
+						ctx, batch, batch, 1, /* StoreID */
+						roachpb.FullReplicaID{RangeID: rs.desc.RangeID, ReplicaID: repl.ReplicaID},
+					)
+					require.NoError(t, err)
+				}
+				tc.updatePostReplicaCreateState(t, ctx, rs, batch)
+
+				// Print the descriptor and batch output.
+				var sb strings.Builder
+				output, err := print.DecodeWriteBatch(batch.Repr())
+				require.NoError(t, err, "error decoding batch")
+				sb.WriteString(fmt.Sprintf("created replica: %v", repl))
+				if output != "" {
+					sb.WriteString("\n")
+					sb.WriteString(output)
+				}
+				// Commit the batch.
+				err = batch.Commit(true)
+				require.NoError(t, err, "error committing batch")
+				return sb.String()
+
+			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.
+				rangeIDs := maps.Keys(tc.ranges)
+				slices.Sort(rangeIDs)
+
+				for _, rangeID := range rangeIDs {
+					rs := tc.ranges[rangeID]
+					sb.WriteString(fmt.Sprintf("%s\n", rs))
+				}
+				return sb.String()
+
+			default:
+				return fmt.Sprintf("unknown command: %s", d.Cmd)
+			}
+		})
+	})
+}
+
+// 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.
+}
+
+// replicaInfo contains the basic info about a replica, used for managing its
+// engine (both raft log and state machine) state.
+type replicaInfo struct {
+	roachpb.FullReplicaID
+	hs raftpb.HardState
+	ts kvserverpb.RaftTruncatedState
+}
+
+// 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
+	nextRangeID roachpb.RangeID // monotonically-increasing rangeID
+	st          *cluster.Settings
+	// 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()
+	return &testCtx{
+		ranges:      make(map[roachpb.RangeID]*rangeState),
+		nextRangeID: 1,
+		st:          st,
+		storage:     storage.NewDefaultInMemForTesting(),
+	}
+}
+
+// close closes the test context's storage engine.
+func (tc *testCtx) close() {
+	tc.storage.Close()
+}
+
+// newRangeState constructs a new rangeState for the supplied descriptor.
+func newRangeState(desc roachpb.RangeDescriptor) *rangeState {
+	return &rangeState{
+		desc:    desc,
+		version: roachpb.Version{Major: 10, Minor: 8, Internal: 7}, // dummy version to avoid churn
+	}
+}
+
+// mustGetRangeState returns the range state for the given range ID.
+func (tc *testCtx) mustGetRangeState(t *testing.T, rangeID roachpb.RangeID) *rangeState {
+	rs, ok := tc.ranges[rangeID]
+	require.True(t, ok, "range-id %d not found", rangeID)
+	return rs
+}
+
+func (tc *testCtx) updatePostReplicaCreateState(
+	t *testing.T, ctx context.Context, rs *rangeState, batch storage.Batch,
+) {
+	// Sanity check that we're not overwriting an existing replica.
+	require.Nil(t, rs.replica)
+	sl := stateloader.Make(rs.desc.RangeID)
+	hs, err := sl.LoadHardState(ctx, batch)
+	require.NoError(t, err)
+	ts, err := sl.LoadRaftTruncatedState(ctx, batch)
+	require.NoError(t, err)
+	replID, err := sl.LoadRaftReplicaID(ctx, batch)
+	require.NoError(t, err)
+	rs.replica = &replicaInfo{
+		FullReplicaID: roachpb.FullReplicaID{
+			RangeID:   rs.desc.RangeID,
+			ReplicaID: replID.ReplicaID,
+		},
+		hs: hs,
+		ts: ts,
+	}
+}
+
+func (rs *rangeState) getReplicaDescriptor(t *testing.T) *roachpb.ReplicaDescriptor {
+	for i, repl := range rs.desc.InternalReplicas {
+		if repl.NodeID == roachpb.NodeID(1) {
+			return &rs.desc.InternalReplicas[i]
+		}
+	}
+	t.Fatal("replica not found")
+	return nil // unreachable
+}
+
+func (rs *rangeState) String() string {
+	var sb strings.Builder
+	sb.WriteString(fmt.Sprintf("range desc: %s", rs.desc))
+	if rs.replica != nil {
+		sb.WriteString(fmt.Sprintf("\n		replica (n1/s1): %s", rs.replica))
+	}
+	return sb.String()
+}
+
+func (r *replicaInfo) String() string {
+	var sb strings.Builder
+	sb.WriteString(fmt.Sprintf("id=%s ", r.FullReplicaID.ReplicaID))
+	if r.hs == (raftpb.HardState{}) {
+		sb.WriteString("uninitialized")
+	} else {
+		sb.WriteString(fmt.Sprintf("HardState={Term:%d,Vote:%d,Commit:%d} ", r.hs.Term, r.hs.Vote, r.hs.Commit))
+		sb.WriteString(fmt.Sprintf("TruncatedState={Index:%d,Term:%d}", r.ts.Index, r.ts.Term))
+	}
+	return sb.String()
+}
+
+func parseReplicas(t *testing.T, val string) []roachpb.NodeID {
+	var replicaNodeIDs []roachpb.NodeID
+	require.True(t, len(val) >= 2 && val[0] == '[' && val[len(val)-1] == ']', "incorrect format")
+	val = val[1 : len(val)-1]
+	for _, s := range strings.Split(val, ",") {
+		var id int
+		_, err := fmt.Sscanf(strings.TrimSpace(s), "%d", &id)
+		require.NoError(t, err)
+		replicaNodeIDs = append(replicaNodeIDs, roachpb.NodeID(id))
+	}
+	// The test is written from the perspective of n1/s1, so not having n1 in
+	// this list should return an error.
+	require.True(t, slices.Contains(replicaNodeIDs, 1), "replica list must contain n1")
+	return replicaNodeIDs
+}

pkg/kv/kvserver/testdata/replica_lifecycle/create_replica.txt

@@ -0,0 +1,32 @@
+create-descriptor start=a end=d replicas=[1,2,3]
+----
+created descriptor: r1:{a-d} [(n1,s1):1, (n2,s2):2, (n3,s3):3, next=4, gen=0]
+
+create-replica range-id=1
+----
+created replica: (n1,s1):1
+Put: 0,0 /Local/RangeID/1/u/RaftReplicaID (0x016989757266747200): replica_id:1 
+
+# Create an initialized replica this time around.
+create-descriptor start=d end=k replicas=[1,2,3]
+----
+created descriptor: r2:{d-k} [(n1,s1):1, (n2,s2):2, (n3,s3):3, next=4, gen=0]
+
+create-replica range-id=2 initialized
+----
+created replica: (n1,s1):1
+Put: 0,0 /Local/RangeID/2/r/RangeLease (0x01698a72726c6c2d00): <empty>
+Put: 0,0 /Local/RangeID/2/r/RangeGCThreshold (0x01698a726c67632d00): 0,0
+Put: 0,0 /Local/RangeID/2/r/RangeGCHint (0x01698a727267636800): latest_range_delete_timestamp:<> gc_timestamp:<> gc_timestamp_next:<> 
+Put: 0,0 /Local/RangeID/2/r/RangeVersion (0x01698a727276657200): 10.8-upgrading-step-007
+Put: 0,0 /Local/RangeID/2/r/RangeAppliedState (0x01698a727261736b00): raft_applied_index:10 lease_applied_index:10 range_stats:<sys_bytes:142 sys_count:4 > raft_closed_timestamp:<> raft_applied_index_term:5 
+Put: 0,0 /Local/RangeID/2/u/RaftReplicaID (0x01698a757266747200): replica_id:1 
+Put: 0,0 /Local/RangeID/2/u/RaftHardState (0x01698a757266746800): term:5 vote:0 commit:10 lead:0 lead_epoch:0 
+Put: 0,0 /Local/RangeID/2/u/RaftTruncatedState (0x01698a757266747400): index:10 term:5 
+
+print-range-state
+----
+range desc: r1:{a-d} [(n1,s1):1, (n2,s2):2, (n3,s3):3, next=4, gen=0]
+		replica (n1/s1): id=1 uninitialized
+range desc: r2:{d-k} [(n1,s1):1, (n2,s2):2, (n3,s3):3, next=4, gen=0]
+		replica (n1/s1): id=1 HardState={Term:5,Vote:0,Commit:10} TruncatedState={Index:10,Term:5}