Merge #150173

150173: kvserver: use FullReplicaID in creation stack r=arulajmani a=pav-kv

This PR moves `FullReplicaID` to be next to `roachpb.{RangeID,ReplicaID}`, and uses it (for type safety) in a few replica creation places where `RangeID` and `ReplicaID` are passed in together.

Epic: CRDB-49111

Co-authored-by: Pavel Kalinnikov <pavel@cockroachlabs.com>

Author: craig[bot]

pkg/kv/kvserver/client_raft_log_queue_test.go

@@ -26,7 +26,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
-	"github.com/cockroachdb/cockroach/pkg/storage"
 	"github.com/cockroachdb/cockroach/pkg/storage/fs"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/listenerutil"
@@ -263,7 +262,7 @@ func TestCrashWhileTruncatingSideloadedEntries(t *testing.T) {
 		propFilter := newAtomicFunc(func(kvserverbase.ProposalFilterArgs) *kvpb.Error {
 			return nil
 		})
-		applyThrottle := newAtomicFunc(func(storage.FullReplicaID) {})
+		applyThrottle := newAtomicFunc(func(roachpb.FullReplicaID) {})
 		postSideEffects := newAtomicFunc(func(args kvserverbase.ApplyFilterArgs) (int, *kvpb.Error) {
 			return 0, nil
 		})
@@ -293,7 +292,7 @@ func TestCrashWhileTruncatingSideloadedEntries(t *testing.T) {
 					Store: &kvserver.StoreTestingKnobs{
 						DisableRaftLogQueue:     true, // we send a log truncation manually
 						DisableSyncLogWriteToss: true, // always use async log writes
-						TestingAfterRaftLogSync: func(id storage.FullReplicaID) { applyThrottle.get()(id) },
+						TestingAfterRaftLogSync: func(id roachpb.FullReplicaID) { applyThrottle.get()(id) },
 						TestingProposalFilter: func(args kvserverbase.ProposalFilterArgs) *kvpb.Error {
 							return propFilter.get()(args)
 						},
@@ -347,7 +346,7 @@ func TestCrashWhileTruncatingSideloadedEntries(t *testing.T) {
 		// Before writing more commands, block the raft commands application flow on
 		// the follower replica.
 		unblockApply := make(chan struct{})
-		applyThrottle.set(func(id storage.FullReplicaID) {
+		applyThrottle.set(func(id roachpb.FullReplicaID) {
 			if id == follower.ID() {
 				applyThrottle.reset()
 				<-unblockApply

pkg/kv/kvserver/kvstorage/datadriven_test.go

@@ -60,7 +60,7 @@ func (e *env) close() {
 func (e *env) handleNewReplica(
 	t *testing.T,
 	ctx context.Context,
-	id storage.FullReplicaID,
+	id roachpb.FullReplicaID,
 	skipRaftReplicaID bool,
 	k, ek roachpb.RKey,
 ) *roachpb.RangeDescriptor {
@@ -154,7 +154,7 @@ func TestDataDriven(t *testing.T) {
 					d.ScanArgs(t, "skip-raft-replica-id", &skipRaftReplicaID)
 				}
 				if desc := e.handleNewReplica(t, ctx,
-					storage.FullReplicaID{RangeID: roachpb.RangeID(rangeID), ReplicaID: roachpb.ReplicaID(replicaID)},
+					roachpb.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)

pkg/kv/kvserver/kvstorage/init.go

@@ -388,8 +388,8 @@ type Replica struct {
 }
 
 // ID returns the FullReplicaID.
-func (r Replica) ID() storage.FullReplicaID {
-	return storage.FullReplicaID{
+func (r Replica) ID() roachpb.FullReplicaID {
+	return roachpb.FullReplicaID{
 		RangeID:   r.RangeID,
 		ReplicaID: r.ReplicaID,
 	}

pkg/kv/kvserver/kvstorage/replica_state.go

@@ -71,6 +71,10 @@ func LoadReplicaState(
 	return ls, nil
 }
 
+func (r LoadedReplicaState) FullReplicaID() roachpb.FullReplicaID {
+	return roachpb.FullReplicaID{RangeID: r.ReplState.Desc.RangeID, ReplicaID: r.ReplicaID}
+}
+
 // check makes sure that the replica invariants hold for the loaded state.
 func (r LoadedReplicaState) check(storeID roachpb.StoreID) error {
 	desc := r.ReplState.Desc
@@ -116,20 +120,16 @@ const CreateUninitReplicaTODO = 0
 // Returns kvpb.RaftGroupDeletedError if this replica can not be created
 // because it has been deleted.
 func CreateUninitializedReplica(
-	ctx context.Context,
-	eng storage.Engine,
-	storeID roachpb.StoreID,
-	rangeID roachpb.RangeID,
-	replicaID roachpb.ReplicaID,
+	ctx context.Context, eng storage.Engine, storeID roachpb.StoreID, id roachpb.FullReplicaID,
 ) error {
-	sl := stateloader.Make(rangeID)
+	sl := stateloader.Make(id.RangeID)
 	// Before creating the replica, see if there is a tombstone which would
 	// indicate that this replica has been removed.
 	// TODO(pav-kv): should also check that there is no existing replica, i.e.
 	// ReplicaID load should find nothing.
 	if ts, err := sl.LoadRangeTombstone(ctx, eng); err != nil {
 		return err
-	} else if replicaID < ts.NextReplicaID {
+	} else if id.ReplicaID < ts.NextReplicaID {
 		return &kvpb.RaftGroupDeletedError{}
 	}
 
@@ -140,12 +140,12 @@ func CreateUninitializedReplica(
 	// non-existent. The only RangeID-specific key that can be present is the
 	// RangeTombstone inspected above.
 	_ = CreateUninitReplicaTODO
-	if err := sl.SetRaftReplicaID(ctx, eng, replicaID); err != nil {
+	if err := sl.SetRaftReplicaID(ctx, eng, id.ReplicaID); err != nil {
 		return err
 	}
 
 	// Make sure that storage invariants for this uninitialized replica hold.
-	uninitDesc := roachpb.RangeDescriptor{RangeID: rangeID}
-	_, err := LoadReplicaState(ctx, eng, storeID, &uninitDesc, replicaID)
+	uninitDesc := roachpb.RangeDescriptor{RangeID: id.RangeID}
+	_, err := LoadReplicaState(ctx, eng, storeID, &uninitDesc, id.ReplicaID)
 	return err
 }

pkg/kv/kvserver/replica.go

@@ -1096,8 +1096,8 @@ func (r *Replica) ReplicaID() roachpb.ReplicaID {
 }
 
 // ID returns the FullReplicaID for the Replica.
-func (r *Replica) ID() storage.FullReplicaID {
-	return storage.FullReplicaID{RangeID: r.RangeID, ReplicaID: r.replicaID}
+func (r *Replica) ID() roachpb.FullReplicaID {
+	return roachpb.FullReplicaID{RangeID: r.RangeID, ReplicaID: r.replicaID}
 }
 
 // LogStorageRaftMuLocked returns the Replica's log storage.

pkg/kv/kvserver/replica_init.go

@@ -79,7 +79,7 @@ func loadInitializedReplicaForTesting(
 func newInitializedReplica(
 	store *Store, loaded kvstorage.LoadedReplicaState, waitForPrevLeaseToExpire bool,
 ) (*Replica, error) {
-	r := newUninitializedReplicaWithoutRaftGroup(store, loaded.ReplState.Desc.RangeID, loaded.ReplicaID)
+	r := newUninitializedReplicaWithoutRaftGroup(store, loaded.FullReplicaID())
 	r.raftMu.Lock()
 	defer r.raftMu.Unlock()
 	r.mu.Lock()
@@ -99,10 +99,8 @@ func newInitializedReplica(
 //
 // TODO(#94912): we actually have another initialization path which should be
 // refactored: Replica.initFromSnapshotLockedRaftMuLocked().
-func newUninitializedReplica(
-	store *Store, rangeID roachpb.RangeID, replicaID roachpb.ReplicaID,
-) (*Replica, error) {
-	r := newUninitializedReplicaWithoutRaftGroup(store, rangeID, replicaID)
+func newUninitializedReplica(store *Store, id roachpb.FullReplicaID) (*Replica, error) {
+	r := newUninitializedReplicaWithoutRaftGroup(store, id)
 	r.raftMu.Lock()
 	defer r.raftMu.Unlock()
 	r.mu.Lock()
@@ -118,17 +116,15 @@ func newUninitializedReplica(
 // newUninitializedReplica() instead. This only exists for
 // newInitializedReplica() to avoid creating the Raft group twice (once when
 // creating the uninitialized replica, and once when initializing it).
-func newUninitializedReplicaWithoutRaftGroup(
-	store *Store, rangeID roachpb.RangeID, replicaID roachpb.ReplicaID,
-) *Replica {
-	uninitState := stateloader.UninitializedReplicaState(rangeID)
+func newUninitializedReplicaWithoutRaftGroup(store *Store, id roachpb.FullReplicaID) *Replica {
+	uninitState := stateloader.UninitializedReplicaState(id.RangeID)
 	r := &Replica{
 		AmbientContext: store.cfg.AmbientCtx,
-		RangeID:        rangeID,
-		replicaID:      replicaID,
+		RangeID:        id.RangeID,
+		replicaID:      id.ReplicaID,
 		creationTime:   timeutil.Now(),
 		store:          store,
-		abortSpan:      abortspan.New(rangeID),
+		abortSpan:      abortspan.New(id.RangeID),
 		concMgr: concurrency.NewManager(concurrency.Config{
 			NodeDesc:           store.nodeDesc,
 			RangeDesc:          uninitState.Desc,
@@ -145,7 +141,7 @@ func newUninitializedReplicaWithoutRaftGroup(
 		}),
 		allocatorToken: &plan.AllocatorToken{},
 	}
-	r.sideTransportClosedTimestamp.init(store.cfg.ClosedTimestampReceiver, rangeID)
+	r.sideTransportClosedTimestamp.init(store.cfg.ClosedTimestampReceiver, id.RangeID)
 	r.cachedClosedTimestampPolicy.Store(new(ctpb.RangeClosedTimestampPolicy))
 
 	r.mu.pendingLeaseRequest = makePendingLeaseRequest(r)
@@ -165,9 +161,9 @@ func newUninitializedReplicaWithoutRaftGroup(
 			// Expose proposal data for external test packages.
 			return store.cfg.TestingKnobs.TestingProposalSubmitFilter(kvserverbase.ProposalFilterArgs{
 				Ctx:        p.Context(),
-				RangeID:    rangeID,
+				RangeID:    id.RangeID,
 				StoreID:    store.StoreID(),
-				ReplicaID:  replicaID,
+				ReplicaID:  id.ReplicaID,
 				Cmd:        p.command,
 				QuotaAlloc: p.quotaAlloc,
 				CmdID:      p.idKey,
@@ -196,7 +192,7 @@ func newUninitializedReplicaWithoutRaftGroup(
 	// NB: state will be loaded when the replica gets initialized.
 	r.shMu.state = uninitState
 
-	r.rangeStr.store(replicaID, uninitState.Desc)
+	r.rangeStr.store(id.ReplicaID, uninitState.Desc)
 	// Add replica log tag - the value is rangeStr.String().
 	r.AmbientContext.AddLogTag("r", &r.rangeStr)
 	r.raftCtx = logtags.AddTag(r.AnnotateCtx(context.Background()), "raft", nil /* value */)
@@ -205,14 +201,14 @@ func newUninitializedReplicaWithoutRaftGroup(
 	// r.AmbientContext.AddLogTag("@", fmt.Sprintf("%x", unsafe.Pointer(r)))
 
 	r.raftMu.rangefeedCTLagObserver = newRangeFeedCTLagObserver()
-	r.raftMu.stateLoader = stateloader.Make(rangeID)
+	r.raftMu.stateLoader = stateloader.Make(id.RangeID)
 
 	// Initialize all the components of the log storage. The state of the log
 	// storage, such as RaftTruncatedState and the last entry ID, will be loaded
 	// when the replica is initialized.
 	sideloaded := logstore.NewDiskSideloadStorage(
 		store.cfg.Settings,
-		rangeID,
+		id.RangeID,
 		// NB: sideloaded log entries are persisted in the state engine so that they
 		// can be ingested to the state machine locally, when being applied.
 		store.StateEngine().GetAuxiliaryDir(),
@@ -229,13 +225,13 @@ func newUninitializedReplicaWithoutRaftGroup(
 	r.logStorage.mu.RWMutex = (*syncutil.RWMutex)(&r.mu.ReplicaMutex)
 	r.logStorage.raftMu.Mutex = &r.raftMu.Mutex
 	r.logStorage.ls = &logstore.LogStore{
-		RangeID:     rangeID,
+		RangeID:     id.RangeID,
 		Engine:      store.LogEngine(),
 		Sideload:    sideloaded,
 		StateLoader: r.raftMu.stateLoader.StateLoader,
 		// NOTE: use the same SyncWaiter loop for all raft log writes performed by a
 		// given range ID, to ensure that callbacks are processed in order.
-		SyncWaiter: store.syncWaiters[int(rangeID)%len(store.syncWaiters)],
+		SyncWaiter: store.syncWaiters[int(id.RangeID)%len(store.syncWaiters)],
 		Settings:   store.cfg.Settings,
 		DisableSyncLogWriteToss: buildutil.CrdbTestBuild &&
 			store.TestingKnobs().DisableSyncLogWriteToss,

pkg/kv/kvserver/replica_raft.go

@@ -2817,9 +2817,10 @@ func (r *Replica) acquireSplitLock(
 	ctx context.Context, split *roachpb.SplitTrigger,
 ) (func(), error) {
 	rightReplDesc, _ := split.RightDesc.GetReplicaDescriptor(r.StoreID())
-	rightRepl, _, err := r.store.getOrCreateReplica(
-		ctx, split.RightDesc.RangeID, rightReplDesc.ReplicaID, nil, /* creatingReplica */
-	)
+	rightRepl, _, err := r.store.getOrCreateReplica(ctx, roachpb.FullReplicaID{
+		RangeID:   split.RightDesc.RangeID,
+		ReplicaID: rightReplDesc.ReplicaID,
+	}, nil /* creatingReplica */)
 	// If getOrCreateReplica returns RaftGroupDeletedError we know that the RHS
 	// has already been removed. This case is handled properly in splitPostApply.
 	if errors.HasType(err, (*kvpb.RaftGroupDeletedError)(nil)) {
@@ -2852,9 +2853,10 @@ func (r *Replica) acquireMergeLock(
 	// complete, after which the replica will realize it has been destroyed and
 	// reject the snapshot.
 	rightReplDesc, _ := merge.RightDesc.GetReplicaDescriptor(r.StoreID())
-	rightRepl, _, err := r.store.getOrCreateReplica(
-		ctx, merge.RightDesc.RangeID, rightReplDesc.ReplicaID, nil, /* creatingReplica */
-	)
+	rightRepl, _, err := r.store.getOrCreateReplica(ctx, roachpb.FullReplicaID{
+		RangeID:   merge.RightDesc.RangeID,
+		ReplicaID: rightReplDesc.ReplicaID,
+	}, nil /* creatingReplica */)
 	if err != nil {
 		return nil, err
 	}

pkg/kv/kvserver/snapshot_apply_prepare.go

@@ -23,7 +23,7 @@ import (
 // snapWriteBuilder contains the data needed to prepare the on-disk state for a
 // snapshot.
 type snapWriteBuilder struct {
-	id storage.FullReplicaID
+	id roachpb.FullReplicaID
 
 	todoEng  storage.Engine
 	sl       stateloader.StateLoader

pkg/kv/kvserver/snapshot_apply_prepare_test.go

@@ -69,7 +69,7 @@ func TestPrepareSnapApply(t *testing.T) {
 		}
 	}
 
-	id := storage.FullReplicaID{RangeID: 123, ReplicaID: 4}
+	id := roachpb.FullReplicaID{RangeID: 123, ReplicaID: 4}
 	descA := desc(101, "a", "b")
 	descB := desc(102, "b", "z")
 	createRangeData(t, eng, *descA)