Merge #150462

150462: kvserver: pass ReplicationTarget for TransferLease r=tbg a=wenyihu6

**asim: add node id for simulator replica**

Previously, the simulator replica interface did not include a NodeID() method.
Future changes to replica_rankings.go will require this as part of the candidate
replica interface. This commit updates the simulator replica to implement
NodeID() method.

Epic: none
Release note: none

---

**kvserver: pass ReplicationTarget for TransferLease**

Previously, the replicate queue’s TransferLease used store IDs to represent the
source and target. This will change in upcoming mma work, which requires full
replication targets as arguments. This change also brings it in line with
RangeRebalancer.RelocateRange for consistency.

Epic: none
Release note: none


Co-authored-by: wenyihu6 <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/asim/state/impl.go

@@ -628,6 +628,7 @@ func (s *state) addReplica(
 
 	desc := rng.desc.AddReplica(roachpb.NodeID(nodeID), roachpb.StoreID(storeID), rtype)
 	replica := &replica{
+		nodeID:    nodeID,
 		replicaID: ReplicaID(desc.ReplicaID),
 		storeID:   storeID,
 		rangeID:   rangeID,
@@ -1607,6 +1608,7 @@ func (r *rng) Size() int64 {
 // replica is an implementation of the Replica interface.
 type replica struct {
 	replicaID  ReplicaID
+	nodeID     NodeID
 	storeID    StoreID
 	rangeID    RangeID
 	desc       roachpb.ReplicaDescriptor
@@ -1618,6 +1620,11 @@ func (r *replica) ReplicaID() ReplicaID {
 	return r.replicaID
 }
 
+// NodeID returns the ID of the node this replica is on.
+func (r *replica) NodeID() NodeID {
+	return r.nodeID
+}
+
 // StoreID returns the ID of the store this replica is on.
 func (r *replica) StoreID() StoreID {
 	return r.storeID

pkg/kv/kvserver/asim/state/state.go

@@ -270,6 +270,8 @@ type Range interface {
 type Replica interface {
 	// ReplicaID returns the ID of this replica.
 	ReplicaID() ReplicaID
+	// NodeID returns the ID of the node this replica is on.
+	NodeID() NodeID
 	// StoreID returns the ID of the store this replica is on.
 	StoreID() StoreID
 	// Descriptor returns the descriptor for this replica.

pkg/kv/kvserver/asim/storerebalancer/candidate_replica.go

@@ -49,6 +49,11 @@ func (sr *simulatorReplica) OwnsValidLease(context.Context, hlc.ClockTimestamp)
 	return sr.repl.HoldsLease()
 }
 
+// NodeID returns the Replica's NodeID.
+func (sr *simulatorReplica) NodeID() roachpb.NodeID {
+	return roachpb.NodeID(sr.repl.NodeID())
+}
+
 // StoreID returns the Replica's StoreID.
 func (sr *simulatorReplica) StoreID() roachpb.StoreID {
 	return roachpb.StoreID(sr.repl.StoreID())

pkg/kv/kvserver/replica_rankings.go

@@ -31,6 +31,8 @@ type CandidateReplica interface {
 	// OwnsValidLease returns whether this replica is the current valid
 	// leaseholder.
 	OwnsValidLease(context.Context, hlc.ClockTimestamp) bool
+	// NodeID returns the Replica's NodeID.
+	NodeID() roachpb.NodeID
 	// StoreID returns the Replica's StoreID.
 	StoreID() roachpb.StoreID
 	// GetRangeID returns the Range ID.

pkg/kv/kvserver/replicate_queue.go

@@ -809,7 +809,7 @@ func (rq *replicateQueue) applyChange(
 	case plan.AllocationFinalizeAtomicReplicationOp:
 		err = rq.finalizeAtomicReplication(ctx, replica)
 	case plan.AllocationTransferLeaseOp:
-		err = rq.TransferLease(ctx, replica, op.Source.StoreID, op.Target.StoreID, op.Usage)
+		err = rq.TransferLease(ctx, replica, op.Source, op.Target, op.Usage)
 	case plan.AllocationChangeReplicasOp:
 		err = rq.changeReplicas(
 			ctx,
@@ -949,7 +949,7 @@ func (rq *replicateQueue) shedLease(
 	rangeUsageInfo := repl.RangeUsageInfo()
 	// Learner replicas aren't allowed to become the leaseholder or raft leader,
 	// so only consider the `VoterDescriptors` replicas.
-	target := rq.allocator.TransferLeaseTarget(
+	targetDesc := rq.allocator.TransferLeaseTarget(
 		ctx,
 		rq.storePool,
 		desc,
@@ -960,11 +960,18 @@ func (rq *replicateQueue) shedLease(
 		false, /* forceDecisionWithoutStats */
 		opts,
 	)
-	if target == (roachpb.ReplicaDescriptor{}) {
+	if targetDesc == (roachpb.ReplicaDescriptor{}) {
 		return allocator.NoSuitableTarget, nil
 	}
-
-	if err := rq.TransferLease(ctx, repl, repl.store.StoreID(), target.StoreID, rangeUsageInfo); err != nil {
+	source := roachpb.ReplicationTarget{
+		NodeID:  repl.NodeID(),
+		StoreID: repl.StoreID(),
+	}
+	target := roachpb.ReplicationTarget{
+		NodeID:  targetDesc.NodeID,
+		StoreID: targetDesc.StoreID,
+	}
+	if err := rq.TransferLease(ctx, repl, source, target, rangeUsageInfo); err != nil {
 		return allocator.TransferErr, err
 	}
 	return allocator.TransferOK, nil
@@ -990,7 +997,7 @@ type RangeRebalancer interface {
 	TransferLease(
 		ctx context.Context,
 		rlm ReplicaLeaseMover,
-		source, target roachpb.StoreID,
+		source, target roachpb.ReplicationTarget,
 		rangeUsageInfo allocator.RangeUsageInfo,
 	) error
 
@@ -1019,16 +1026,16 @@ func (rq *replicateQueue) finalizeAtomicReplication(ctx context.Context, repl *R
 func (rq *replicateQueue) TransferLease(
 	ctx context.Context,
 	rlm ReplicaLeaseMover,
-	source, target roachpb.StoreID,
+	source, target roachpb.ReplicationTarget,
 	rangeUsageInfo allocator.RangeUsageInfo,
 ) error {
 	rq.metrics.TransferLeaseCount.Inc(1)
 	log.KvDistribution.Infof(ctx, "transferring lease to s%d", target)
-	if err := rlm.AdminTransferLease(ctx, target, false /* bypassSafetyChecks */); err != nil {
+	if err := rlm.AdminTransferLease(ctx, target.StoreID, false /* bypassSafetyChecks */); err != nil {
 		return errors.Wrapf(err, "%s: unable to transfer lease to s%d", rlm, target)
 	}
 
-	rq.storePool.UpdateLocalStoresAfterLeaseTransfer(source, target, rangeUsageInfo)
+	rq.storePool.UpdateLocalStoresAfterLeaseTransfer(source.StoreID, target.StoreID, rangeUsageInfo)
 	return nil
 }
 

pkg/kv/kvserver/store_rebalancer.go

@@ -581,8 +581,14 @@ func (sr *StoreRebalancer) applyLeaseRebalance(
 		return sr.rr.TransferLease(
 			ctx,
 			candidateReplica,
-			candidateReplica.StoreID(),
-			target.StoreID,
+			roachpb.ReplicationTarget{
+				NodeID:  candidateReplica.NodeID(),
+				StoreID: candidateReplica.StoreID(),
+			},
+			roachpb.ReplicationTarget{
+				NodeID:  target.NodeID,
+				StoreID: target.StoreID,
+			},
 			candidateReplica.RangeUsageInfo(),
 		)
 	}); err != nil {