Improve CancelUnlocking to search for unbonding delegation entry by c…

…reationHeight
sourcenetwork · Dec 9, 2024 · cee7583 · cee7583
1 parent 5dd1b16
commit cee7583
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 11 deletions.
diff --git a/x/tier/keeper/keeper.go b/x/tier/keeper/keeper.go
@@ -285,7 +285,7 @@ func (k Keeper) Redelegate(ctx context.Context, delAddr sdk.AccAddress, srcValAd
 }
 
 // CancelUnlocking effectively cancels the pending unlocking lockup.
-func (k Keeper) CancelUnlocking(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress, amt math.Int) error {
+func (k Keeper) CancelUnlocking(ctx context.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress, creationHeight int64, amt math.Int) error {
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 
 	// use the module account address when interacting with unbonding delegations
@@ -301,36 +301,46 @@ func (k Keeper) CancelUnlocking(ctx context.Context, delAddr sdk.AccAddress, val
 		return errorsmod.Wrapf(err, "unbonding delegation not found for delegator %s and validator %s", modAddr, valAddr)
 	}
 
-	// search for the unbonding delegation entry that matches the amount and ensure it's valid
+	// find unbonding delegation entry by CreationHeight
 	var (
 		unbondEntryIndex int64 = -1
 		unbondEntry      stakingtypes.UnbondingDelegationEntry
 	)
 
 	for i, entry := range ubd.Entries {
-		if entry.Balance.GTE(amt) && entry.CompletionTime.After(sdkCtx.BlockTime()) {
+		if entry.CreationHeight == creationHeight && entry.CompletionTime.After(sdkCtx.BlockTime()) {
 			unbondEntryIndex = int64(i)
 			unbondEntry = entry
 			break
 		}
 	}
 
 	if unbondEntryIndex == -1 {
-		return errorsmod.Wrapf(stakingtypes.ErrNoUnbondingDelegation, "no valid unbonding entry found for amount %s", amt)
+		return errorsmod.Wrapf(
+			stakingtypes.ErrNoUnbondingDelegation,
+			"no valid unbonding entry found for creation height %d",
+			creationHeight,
+		)
+	}
+
+	// if amt is zero, revert entire UnbondingDelegationEntry. Otherwise, revert the specified amt.
+	restoreAmount := amt
+	if restoreAmount.IsZero() {
+		restoreAmount = unbondEntry.Balance
 	}
 
-	_, err = k.stakingKeeper.Delegate(ctx, modAddr, amt, stakingtypes.Unbonding, validator, false)
+	_, err = k.stakingKeeper.Delegate(ctx, modAddr, restoreAmount, stakingtypes.Unbonding, validator, false)
 	if err != nil {
 		return errorsmod.Wrap(err, "failed to delegate tokens back to validator")
 	}
 
 	// update or remove the unbonding delegation entry
-	remainingBalance := unbondEntry.Balance.Sub(amt)
+	remainingBalance := unbondEntry.Balance.Sub(restoreAmount)
 	if remainingBalance.IsZero() {
 		ubd.RemoveEntry(unbondEntryIndex)
 	} else {
 		unbondEntry.Balance = remainingBalance
-		unbondEntry.InitialBalance = unbondEntry.InitialBalance.Sub(amt)
+		unbondEntry.InitialBalance = unbondEntry.InitialBalance.Sub(restoreAmount)
 		ubd.Entries[unbondEntryIndex] = unbondEntry
 	}
 
@@ -344,14 +354,15 @@ func (k Keeper) CancelUnlocking(ctx context.Context, delAddr sdk.AccAddress, val
 		return errorsmod.Wrap(err, "failed to update unbonding delegation")
 	}
 
-	k.AddLockup(ctx, delAddr, valAddr, amt)
+	k.AddLockup(ctx, delAddr, valAddr, restoreAmount)
 
 	sdkCtx.EventManager().EmitEvent(
 		sdk.NewEvent(
 			types.EventTypeCancelUnlocking,
 			sdk.NewAttribute(stakingtypes.AttributeKeyDelegator, delAddr.String()),
 			sdk.NewAttribute(stakingtypes.AttributeKeyValidator, valAddr.String()),
-			sdk.NewAttribute(sdk.AttributeKeyAmount, amt.String()),
+			sdk.NewAttribute(sdk.AttributeKeyAmount, restoreAmount.String()),
+			sdk.NewAttribute(types.AttributeKeyCreationHeight, fmt.Sprintf("%d", creationHeight)),
 		),
 	)
 

diff --git a/x/tier/keeper/keeper_test.go b/x/tier/keeper/keeper_test.go
@@ -255,7 +255,7 @@ func TestCancelUnlocking(t *testing.T) {
 	require.Equal(t, unlockTime, unlTime)
 
 	// cancel (remove) the unlocking lockup
-	err = k.CancelUnlocking(ctx, delAddr, valAddr, adjustedUnlockAmount)
+	err = k.CancelUnlocking(ctx, delAddr, valAddr, creationHeight, adjustedUnlockAmount)
 	require.NoError(t, err)
 
 	// verify that lockup was updated

diff --git a/x/tier/keeper/msg_server.go b/x/tier/keeper/msg_server.go
@@ -68,7 +68,7 @@ func (m msgServer) CancelUnlocking(ctx context.Context, msg *types.MsgCancelUnlo
 	delAddr := sdk.MustAccAddressFromBech32(msg.DelegatorAddress)
 	valAddr := types.MustValAddressFromBech32(msg.ValidatorAddress)
 
-	err := m.Keeper.CancelUnlocking(ctx, delAddr, valAddr, msg.Stake.Amount)
+	err := m.Keeper.CancelUnlocking(ctx, delAddr, valAddr, msg.CreationHeight, msg.Stake.Amount)
 	if err != nil {
 		return nil, errorsmod.Wrap(err, "cancel unlocking")
 	}