kvserver: remove maybeMarkReplicaInitializedLockedReplLocked

[tbg]

Related to #94911.

We seem to have three paths that initialize a Replica. The first one is

cockroach/pkg/kv/kvserver/replica_init.go

Lines 157 to 238 in 3d6fe9e

	// loadRaftMuLockedReplicaMuLocked will load the state of the replica from disk.
	// If desc is initialized, the Replica will be initialized when this method
	// returns. An initialized Replica may not be reloaded. If this method is called
	// with an uninitialized desc it may be called again later with an initialized
	// desc.
	//
	// This method is called in three places:
	//
	// 1. newReplica - used when the store is initializing and during testing
	// 2. tryGetOrCreateReplica - see newUnloadedReplica
	// 3. splitPostApply - this call initializes a previously uninitialized Replica.
	func (r *Replica) loadRaftMuLockedReplicaMuLocked(desc *roachpb.RangeDescriptor) error {
	ctx := r.AnnotateCtx(context.TODO())
	if r.mu.state.Desc != nil && r.IsInitialized() {
	log.Fatalf(ctx, "r%d: cannot reinitialize an initialized replica", desc.RangeID)
	} else if r.replicaID == 0 {
	// NB: This is just a defensive check as r.mu.replicaID should never be 0.
	log.Fatalf(ctx, "r%d: cannot initialize replica without a replicaID", desc.RangeID)
	}
	if desc.IsInitialized() {
	r.setStartKeyLocked(desc.StartKey)
	}
	// Clear the internal raft group in case we're being reset. Since we're
	// reloading the raft state below, it isn't safe to use the existing raft
	// group.
	r.mu.internalRaftGroup = nil
	var err error
	if r.mu.state, err = r.mu.stateLoader.Load(ctx, r.Engine(), desc); err != nil {
	return err
	}
	r.mu.lastIndex, err = r.mu.stateLoader.LoadLastIndex(ctx, r.Engine())
	if err != nil {
	return err
	}
	r.mu.lastTerm = invalidLastTerm
	// Ensure that we're not trying to load a replica with a different ID than
	// was used to construct this Replica.
	replicaID := r.replicaID
	if replicaDesc, found := r.mu.state.Desc.GetReplicaDescriptor(r.StoreID()); found {
	replicaID = replicaDesc.ReplicaID
	} else if desc.IsInitialized() {
	log.Fatalf(ctx, "r%d: cannot initialize replica which is not in descriptor %v", desc.RangeID, desc)
	}
	if r.replicaID != replicaID {
	log.Fatalf(ctx, "attempting to initialize a replica which has ID %d with ID %d",
	r.replicaID, replicaID)
	}
	r.setDescLockedRaftMuLocked(ctx, desc)
	// Only do this if there was a previous lease. This shouldn't be important
	// to do but consider that the first lease which is obtained is back-dated
	// to a zero start timestamp (and this de-flakes some tests). If we set the
	// min proposed TS here, this lease could not be renewed (by the semantics
	// of minLeaseProposedTS); and since minLeaseProposedTS is copied on splits,
	// this problem would multiply to a number of replicas at cluster bootstrap.
	// Instead, we make the first lease special (which is OK) and the problem
	// disappears.
	if r.mu.state.Lease.Sequence > 0 {
	r.mu.minLeaseProposedTS = r.Clock().NowAsClockTimestamp()
	}
	ssBase := r.Engine().GetAuxiliaryDir()
	if r.raftMu.sideloaded, err = newDiskSideloadStorage(
	r.store.cfg.Settings,
	desc.RangeID,
	replicaID,
	ssBase,
	r.store.limiters.BulkIOWriteRate,
	r.store.engine,
	); err != nil {
	return errors.Wrap(err, "while initializing sideloaded storage")
	}
	r.assertStateRaftMuLockedReplicaMuRLocked(ctx, r.store.Engine())
	r.sideTransportClosedTimestamp.init(r.store.cfg.ClosedTimestampReceiver, desc.RangeID)
	return nil
	}

and as explain in its comment, this is what transitions most uninitialized replicas to initialized.

However, then there's also

cockroach/pkg/kv/kvserver/store_create_replica.go

Lines 367 to 430 in e1c24ed

	// maybeMarkReplicaInitializedLocked should be called whenever a previously
	// uninitialized replica has become initialized so that the store can update its
	// internal bookkeeping. It requires that Store.mu and Replica.raftMu
	// are locked.
	func (s *Store) maybeMarkReplicaInitializedLockedReplLocked(
	ctx context.Context, lockedRepl *Replica,
	) error {
	desc := lockedRepl.descRLocked()
	if !desc.IsInitialized() {
	return errors.Errorf("attempted to process uninitialized range %s", desc)
	}
	rangeID := lockedRepl.RangeID
	if _, ok := s.mu.uninitReplicas[rangeID]; !ok {
	// Do nothing if the range has already been initialized.
	return nil
	}
	delete(s.mu.uninitReplicas, rangeID)
	if it := s.getOverlappingKeyRangeLocked(desc); it.item != nil {
	return errors.AssertionFailedf("%s: cannot initialize replica; %s has overlapping range %s",
	s, desc, it.Desc())
	}
	// Copy of the start key needs to be set before inserting into replicasByKey.
	lockedRepl.setStartKeyLocked(desc.StartKey)
	if it := s.mu.replicasByKey.ReplaceOrInsertReplica(ctx, lockedRepl); it.item != nil {
	return errors.AssertionFailedf("range for key %v already exists in replicasByKey btree: %+v",
	it.item.key(), it)
	}
	// Unquiesce the replica. We don't allow uninitialized replicas to unquiesce,
	// but now that the replica has been initialized, we unquiesce it as soon as
	// possible. This replica was initialized in response to the reception of a
	// snapshot from another replica. This means that the other replica is not
	// quiesced, so we don't need to campaign or wake the leader. We just want
	// to start ticking.
	//
	// NOTE: The fact that this replica is being initialized in response to the
	// receipt of a snapshot means that its r.mu.internalRaftGroup must not be
	// nil.
	//
	// NOTE: Unquiescing the replica here is not strictly necessary. As of the
	// time of writing, this function is only ever called below handleRaftReady,
	// which will always unquiesce any eligible replicas before completing. So in
	// marking this replica as initialized, we have made it eligible to unquiesce.
	// However, there is still a benefit to unquiecing here instead of letting
	// handleRaftReady do it for us. The benefit is that handleRaftReady cannot
	// make assumptions about the state of the other replicas in the range when it
	// unquieces a replica, so when it does so, it also instructs the replica to
	// campaign and to wake the leader (see maybeUnquiesceAndWakeLeaderLocked).
	// We have more information here (see "This means that the other replica ..."
	// above) and can make assumptions about the state of the other replicas in
	// the range, so we can unquiesce without campaigning or waking the leader.
	if !lockedRepl.maybeUnquiesceWithOptionsLocked(false /* campaignOnWake */) {
	return errors.AssertionFailedf("expected replica %s to unquiesce after initialization", desc)
	}
	// Add the range to metrics and maybe gossip on capacity change.
	s.metrics.ReplicaCount.Inc(1)
	s.maybeGossipOnCapacityChange(ctx, rangeAddEvent)
	return nil
	}

which has as its sole caller the applySnapshot path.

We should unify them by replacing the latter use with the former.

Jira issue: CRDB-23225
Epic: CRDB-220

[blathers-crl]

cc @cockroachdb/replication

[tbg]

Only the second one even removes the uninited replica from store.mu.uninitReplicas, I wondered why we weren't leaking entries in that map on splits (which don't call that method) but it turns out we sort of randomly remove them here (and move the right hand side replica over to the initialized map a few lines down):

cockroach/pkg/kv/kvserver/store_split.go

Lines 312 to 323 in 352d2ef

	s.mu.Lock()
	defer s.mu.Unlock()
	if exRng, ok := s.mu.uninitReplicas[rightDesc.RangeID]; rightReplOrNil != nil && ok {
	// If we have an uninitialized replica of the new range we require pointer
	// equivalence with rightRepl. See Store.splitTriggerPostApply().
	if exRng != rightReplOrNil {
	log.Fatalf(ctx, "found unexpected uninitialized replica: %s vs %s", exRng, rightReplOrNil)
	}
	// NB: We only remove from uninitReplicas here so that we don't leave open a
	// window where a replica is temporarily not present in Store.mu.replicas.
	delete(s.mu.uninitReplicas, rightDesc.RangeID)
	}

There are probably reasons for that. We have already called loadRaftMuLocked[...] on the right hand side at that point, we did that earlier in splitPostApply:

cockroach/pkg/kv/kvserver/store_split.go

Lines 253 to 257 in 352d2ef

	// Finish initialization of the RHS.
	err := rightRepl.loadRaftMuLockedReplicaMuLocked(&split.RightDesc)
	if err != nil {
	log.Fatalf(ctx, "%v", err)
	}

Also, it seems sketchy how we're updating the span of the LHS replica and that of the RHS replica in separate s.mu critical sections. For example, say we split r1 [a,c) into r1 [a,b) and r2 [b,c). Then in the code above this paragraph, we've already updated r2 to have span [b,c). But r1 still owns [a,c). We get away with it because at that point, the store still believes that rhsRepl is an uninitialized replicas and so it isn't inserted yet into the data structures that need the key -> replica map to be injective. But this seems really messy.

cockroach/pkg/kv/kvserver/store_split.go

Lines 312 to 325 in 352d2ef

	s.mu.Lock()
	defer s.mu.Unlock()
	if exRng, ok := s.mu.uninitReplicas[rightDesc.RangeID]; rightReplOrNil != nil && ok {
	// If we have an uninitialized replica of the new range we require pointer
	// equivalence with rightRepl. See Store.splitTriggerPostApply().
	if exRng != rightReplOrNil {
	log.Fatalf(ctx, "found unexpected uninitialized replica: %s vs %s", exRng, rightReplOrNil)
	}
	// NB: We only remove from uninitReplicas here so that we don't leave open a
	// window where a replica is temporarily not present in Store.mu.replicas.
	delete(s.mu.uninitReplicas, rightDesc.RangeID)
	}
	leftRepl.setDescRaftMuLocked(ctx, newLeftDesc)

[tbg]

cc @pavelkalinnikov

[pav-kv]

There are mainly two aspects of initializing a Replica:

1. Initializing a Replica :)
2. Updating the Store state correspondingly.

Also, the two should probably be done with some synchronization.

The two code paths linked in the issue description reside in Replica and Store correspondingly, so it makes some sense and we could take it as a basis. We could eventually arrive at one function if there is too much commonality in the orchestration / "glue" between these 2 parts across different workflows (split or snapshot application).

Seemingly, it is the "glue" between these two things that is currently ad-hoc / scattered all over the place. For example, the split trigger application path does (1) through a "good" function, but not so well on (2) - it's hidden in split trigger code. The snapshot application does the opposite: on (2) it uses the "maybe mark" func, but does (1) in applySnapshot instead of using a shared helper.