asim: use roachpb.ReplicationTarget for AllocationTransferLeaseOp

Previously, AllocationTransferLeaseOp used store ID to represent the source and
target. This commit updates it to use roachpb.ReplicationTarget, which includes
both store id and node id. We made this change because future commits will pass
fields from AllocationTransferLeaseOp directly into NonMMAPreTransferLease,
which requires ReplicationTarget. Note that production behavior remains
unchanged, as only the store ID is accessed for them.

Epic: none
Release note: none

Author: wenyihu6

pkg/kv/kvserver/allocator/plan/lease.go

@@ -172,8 +172,12 @@ func (lp LeasePlanner) PlanOneChange(
 	}
 
 	change.Op = AllocationTransferLeaseOp{
-		Source:             repl.StoreID(),
-		Target:             target.StoreID,
+		Source: roachpb.ReplicationTarget{
+			StoreID: repl.StoreID(), NodeID: repl.NodeID(),
+		},
+		Target: roachpb.ReplicationTarget{
+			StoreID: target.StoreID, NodeID: target.NodeID,
+		},
 		Usage:              usage,
 		bypassSafetyChecks: false,
 	}

pkg/kv/kvserver/allocator/plan/op.go

@@ -29,7 +29,7 @@ type AllocationOp interface {
 // AllocationTransferLeaseOp represents an operation to transfer a range lease to another
 // store, from the current one.
 type AllocationTransferLeaseOp struct {
-	Target, Source     roachpb.StoreID
+	Target, Source     roachpb.ReplicationTarget
 	Usage              allocator.RangeUsageInfo
 	bypassSafetyChecks bool
 }

pkg/kv/kvserver/allocator/plan/replicate.go

@@ -100,6 +100,7 @@ type AllocatorReplica interface {
 	GetCompactedIndex() kvpb.RaftIndex
 	LastReplicaAdded() (roachpb.ReplicaID, time.Time)
 	StoreID() roachpb.StoreID
+	NodeID() roachpb.NodeID
 	GetRangeID() roachpb.RangeID
 	SendStreamStats(*rac2.RangeSendStreamStats)
 }
@@ -908,8 +909,14 @@ func (rp ReplicaPlanner) maybeTransferLeaseAwayTarget(
 	log.KvDistribution.Infof(ctx, "transferring away lease to s%d", target.StoreID)
 
 	op = AllocationTransferLeaseOp{
-		Source:             repl.StoreID(),
-		Target:             target.StoreID,
+		Source: roachpb.ReplicationTarget{
+			NodeID:  repl.NodeID(),
+			StoreID: repl.StoreID(),
+		},
+		Target: roachpb.ReplicationTarget{
+			NodeID:  target.NodeID,
+			StoreID: target.StoreID,
+		},
 		Usage:              usageInfo,
 		bypassSafetyChecks: false,
 	}

pkg/kv/kvserver/asim/queue/allocator_replica.go

@@ -160,3 +160,8 @@ func (sr *SimulatorReplica) RangeUsageInfo() allocator.RangeUsageInfo {
 
 func (sr *SimulatorReplica) SendStreamStats(stats *rac2.RangeSendStreamStats) {
 }
+
+// NodeID returns the Replica's NodeID.
+func (sr *SimulatorReplica) NodeID() roachpb.NodeID {
+	return sr.repl.Descriptor().NodeID
+}

pkg/kv/kvserver/asim/queue/lease_queue.go

@@ -113,6 +113,8 @@ func (lq *leaseQueue) MaybeAdd(ctx context.Context, replica state.Replica, s sta
 func (lq *leaseQueue) Tick(ctx context.Context, tick time.Time, s state.State) {
 	lq.AddLogTag("tick", tick)
 	ctx = lq.ResetAndAnnotateCtx(ctx)
+	// TODO(wenyihu6): it is unclear why next tick is forwarded to last tick
+	// here (see #149904 for more details).
 	if lq.lastTick.After(lq.next) {
 		lq.next = lq.lastTick
 	}

pkg/kv/kvserver/asim/queue/replicate_queue.go

@@ -112,6 +112,8 @@ func (rq *replicateQueue) MaybeAdd(ctx context.Context, replica state.Replica, s
 func (rq *replicateQueue) Tick(ctx context.Context, tick time.Time, s state.State) {
 	rq.AddLogTag("tick", tick)
 	ctx = rq.ResetAndAnnotateCtx(ctx)
+	// TODO(wenyihu6): it is unclear why next tick is forwarded to last tick
+	// here (see #149904 for more details).
 	if rq.lastTick.After(rq.next) {
 		rq.next = rq.lastTick
 	}
@@ -175,8 +177,8 @@ func pushReplicateChange(
 	case plan.AllocationTransferLeaseOp:
 		stateChange = &state.LeaseTransferChange{
 			RangeID:        state.RangeID(change.Replica.GetRangeID()),
-			TransferTarget: state.StoreID(op.Target),
-			Author:         state.StoreID(op.Source),
+			TransferTarget: state.StoreID(op.Target.StoreID),
+			Author:         state.StoreID(op.Source.StoreID),
 			Wait:           delayFn(rng.Size(), true),
 		}
 	case plan.AllocationChangeReplicasOp:

pkg/kv/kvserver/asim/queue/split_queue.go

@@ -68,6 +68,8 @@ func (sq *splitQueue) MaybeAdd(ctx context.Context, replica state.Replica, state
 // FIFO order on ties. The tick currently only considers size based range
 // splitting.
 func (sq *splitQueue) Tick(ctx context.Context, tick time.Time, s state.State) {
+	// TODO(wenyihu6): it is unclear why next tick is forwarded to last tick
+	// here (see #149904 for more details).
 	if sq.lastTick.After(sq.next) {
 		sq.next = sq.lastTick
 	}

pkg/kv/kvserver/lease_queue.go

@@ -136,7 +136,7 @@ func (lq *leaseQueue) process(
 		lease, _ := repl.GetLease()
 		log.KvDistribution.Infof(ctx, "transferring lease to s%d usage=%v, lease=[%v type=%v]", transferOp.Target, transferOp.Usage, lease, lease.Type())
 		lq.lastLeaseTransfer.Store(timeutil.Now())
-		if err := repl.AdminTransferLease(ctx, transferOp.Target, false /* bypassSafetyChecks */); err != nil {
+		if err := repl.AdminTransferLease(ctx, transferOp.Target.StoreID, false /* bypassSafetyChecks */); err != nil {
 			return false, errors.Wrapf(err, "%s: unable to transfer lease to s%d", repl, transferOp.Target)
 		}
 		change.Op.ApplyImpact(lq.storePool)

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, op.Target, op.Usage)
+		err = rq.TransferLease(ctx, replica, op.Source.StoreID, op.Target.StoreID, op.Usage)
 	case plan.AllocationChangeReplicasOp:
 		err = rq.changeReplicas(
 			ctx,