Merge #151351

151351: kvserver: remove `ApplyImpact` from `AllocationOp` r=sumeerbhola a=wenyihu6

Epic: CRDB-25222
Release note: none

---

**kvserver: update local store pool directly with AllocationChangeReplicasOp**

Previously, `change.Op.ApplyImpact` was used to update the replicateQueue's
local store pool state after a successful change replica call to reflect the
result of applying the operation. But in practice only
`AllocationChangeReplicasOp` used it - other ops update the store pool directly.
This commit updates `AllocationChangeReplicasOp` to also update the local store
pool directly after a successful `repl.changeReplicasImpl`.

In a follow-up, we'll remove `ApplyImpact` from the `AllocationOp` interface
entirely, since it adds an unnecessary layer of indirection without providing
much value.

---

**kvserver: remove `ApplyImpact` from `AllocationOp`**

This commit removes `ApplyImpact` from the `AllocationOp` interface entirely,
since it adds an unnecessary layer of indirection without providing much value.


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

Author: craig[bot]

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

@@ -8,7 +8,6 @@ package plan
 import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
-	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 )
@@ -18,9 +17,6 @@ import (
 //
 // TODO(kvoli): Add AllocationRelocateRangeOp.
 type AllocationOp interface {
-	// ApplyImpact updates the given storepool to reflect the result of
-	// applying this operation.
-	ApplyImpact(storepool storepool.AllocatorStorePool)
 	// LHBeingRemoved returns true when the leaseholder is will be removed if
 	// this operation succeeds, otherwise false.
 	LHBeingRemoved() bool
@@ -43,12 +39,6 @@ func (o AllocationTransferLeaseOp) LHBeingRemoved() bool {
 	return true
 }
 
-func (o AllocationTransferLeaseOp) ApplyImpact(storepool storepool.AllocatorStorePool) {
-	// TODO(kvoli): Currently the local storepool is updated directly in the
-	// lease transfer call, rather than in this function. Move the storepool
-	// tracking from rq.TransferLease to this function once #89771 is merged.
-}
-
 // AllocationChangeReplicasOp represents an operation to execute a change
 // replicas txn.
 type AllocationChangeReplicasOp struct {
@@ -73,14 +63,6 @@ func (o AllocationChangeReplicasOp) LHBeingRemoved() bool {
 	return false
 }
 
-// applyEstimatedImpact updates the given storepool to reflect the result
-// of applying this operation.
-func (o AllocationChangeReplicasOp) ApplyImpact(storepool storepool.AllocatorStorePool) {
-	for _, chg := range o.Chgs {
-		storepool.UpdateLocalStoreAfterRebalance(chg.Target.StoreID, o.Usage, chg.ChangeType)
-	}
-}
-
 // AllocationFinalizeAtomicReplicationOp represents an operation to finalize an
 // atomic change replicas operation and remove any remaining learners.
 type AllocationFinalizeAtomicReplicationOp struct{}
@@ -89,13 +71,11 @@ var _ AllocationOp = &AllocationFinalizeAtomicReplicationOp{}
 
 // TODO(kvoli): This always returns false, however it is possible that the LH
 // may have been removed here.
-func (o AllocationFinalizeAtomicReplicationOp) LHBeingRemoved() bool                               { return false }
-func (o AllocationFinalizeAtomicReplicationOp) ApplyImpact(storepool storepool.AllocatorStorePool) {}
+func (o AllocationFinalizeAtomicReplicationOp) LHBeingRemoved() bool { return false }
 
 // AllocationNoop represents no operation.
 type AllocationNoop struct{}
 
 var _ AllocationOp = &AllocationNoop{}
 
-func (o AllocationNoop) LHBeingRemoved() bool                               { return false }
-func (o AllocationNoop) ApplyImpact(storepool storepool.AllocatorStorePool) {}
+func (o AllocationNoop) LHBeingRemoved() bool { return false }

pkg/kv/kvserver/lease_queue.go

@@ -140,12 +140,6 @@ func (lq *leaseQueue) process(
 			return false, errors.Wrapf(err, "%s: unable to transfer lease to s%d", repl, transferOp.Target)
 		}
 
-		// TODO(wenyihu6): Initially, change.Op.ApplyImpact was used here. This was
-		// a problem since AllocationTransferLeaseOp.ApplyImpact was left
-		// unimplemented. We should either implement
-		// AllocationTransferLeaseOp.ApplyImpact correctly or remove the use of
-		// ApplyImpact entirely. The replicate queue does not have this issue since
-		// it uses rq.TransferLease, which updates the local store pool directly.
 		lq.storePool.UpdateLocalStoresAfterLeaseTransfer(
 			transferOp.Source.StoreID, transferOp.Target.StoreID, transferOp.Usage)
 	}

pkg/kv/kvserver/replicate_queue.go

@@ -816,6 +816,7 @@ func (rq *replicateQueue) applyChange(
 			replica,
 			op.Chgs,
 			replica.Desc(),
+			op.Usage,
 			op.AllocatorPriority,
 			op.Reason,
 			op.Details,
@@ -914,10 +915,6 @@ func (rq *replicateQueue) processOneChange(
 		return false, maybeAnnotateDecommissionErr(err, change.Action)
 	}
 
-	// Update the local storepool state to reflect the successful application
-	// of the change.
-	change.Op.ApplyImpact(rq.storePool)
-
 	// Requeue the replica if it meets the criteria in ShouldRequeue.
 	return ShouldRequeue(ctx, change, conf), nil
 }
@@ -1060,18 +1057,26 @@ func (rq *replicateQueue) changeReplicas(
 	repl *Replica,
 	chgs kvpb.ReplicationChanges,
 	desc *roachpb.RangeDescriptor,
+	rangeUsageInfo allocator.RangeUsageInfo,
 	allocatorPriority float64,
 	reason kvserverpb.RangeLogEventReason,
 	details string,
 ) error {
 	// NB: this calls the impl rather than ChangeReplicas because
 	// the latter traps tests that try to call it while the replication
 	// queue is active.
-	_, err := repl.changeReplicasImpl(
+	if _, err := repl.changeReplicasImpl(
 		ctx, desc, kvserverpb.SnapshotRequest_REPLICATE_QUEUE, allocatorPriority, reason,
 		details, chgs,
-	)
-	return err
+	); err != nil {
+		return err
+	}
+	// On success, update local store pool to reflect the result of applying the
+	// operations.
+	for _, chg := range chgs {
+		rq.storePool.UpdateLocalStoreAfterRebalance(chg.Target.StoreID, rangeUsageInfo, chg.ChangeType)
+	}
+	return nil
 }
 
 func (*replicateQueue) postProcessScheduled(