fix: fix the cancel unbond

CHANGELOG.md

@@ -82,6 +82,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/staking) [#12409](https://github.com/cosmos/cosmos-sdk/pull/12409) Migrate `x/staking` to self-managed parameters and deprecate it's usage of `x/params`.
 * (x/bank) [#11859](https://github.com/cosmos/cosmos-sdk/pull/11859) Move the SendEnabled information out of the Params and into the state store directly.
 * (x/gov) [#12771](https://github.com/cosmos/cosmos-sdk/pull/12771) Initial deposit requirement for proposals at submission time.
+* (x/staking) [#12967](https://github.com/cosmos/cosmos-sdk/pull/12967) `unbond` now creates only one unbonding delegation entry when multiple unbondings exist at a single height (e.g. through multiple messages in a transaction).
 
 ### API Breaking Changes
 

x/staking/keeper/delegation_test.go

@@ -271,8 +271,9 @@ func TestUnbondingDelegationsMaxEntries(t *testing.T) {
 
 	// should all pass
 	var completionTime time.Time
-	for i := uint32(0); i < maxEntries; i++ {
+	for i := int64(0); i < int64(maxEntries); i++ {
 		var err error
+		ctx = ctx.WithBlockHeight(i)
 		completionTime, err = app.StakingKeeper.Undelegate(ctx, addrDels[0], addrVals[0], math.LegacyNewDec(1))
 		require.NoError(t, err)
 	}

x/staking/migrations/v4/migrations_test.go

@@ -2,16 +2,19 @@ package v4_test
 
 import (
 	"testing"
+	"time"
 
-	"github.com/stretchr/testify/require"
-
+	"github.com/cosmos/cosmos-sdk/codec"
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/testutil"
+	"github.com/cosmos/cosmos-sdk/testutil/sims"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	moduletestutil "github.com/cosmos/cosmos-sdk/types/module/testutil"
-	"github.com/cosmos/cosmos-sdk/x/distribution"
 	paramtypes "github.com/cosmos/cosmos-sdk/x/params/types"
-	"github.com/cosmos/cosmos-sdk/x/staking/migrations/v4"
+	"github.com/cosmos/cosmos-sdk/x/staking"
+	v4 "github.com/cosmos/cosmos-sdk/x/staking/migrations/v4"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
+	"github.com/stretchr/testify/require"
 )
 
 type mockSubspace struct {
@@ -27,19 +30,112 @@ func (ms mockSubspace) GetParamSet(ctx sdk.Context, ps paramtypes.ParamSet) {
 }
 
 func TestMigrate(t *testing.T) {
-	encCfg := moduletestutil.MakeTestEncodingConfig(distribution.AppModuleBasic{})
+	encCfg := moduletestutil.MakeTestEncodingConfig(staking.AppModuleBasic{})
 	cdc := encCfg.Codec
 
 	storeKey := sdk.NewKVStoreKey(v4.ModuleName)
 	tKey := sdk.NewTransientStoreKey("transient_test")
 	ctx := testutil.DefaultContext(storeKey, tKey)
 	store := ctx.KVStore(storeKey)
+	duplicateCreationHeight := int64(1)
+
+	accAddrs := sims.CreateIncrementalAccounts(1)
+	accAddr := accAddrs[0]
+
+	valAddrs := sims.ConvertAddrsToValAddrs(accAddrs)
+	valAddr := valAddrs[0]
+
+	// creating 10 ubdEntries with same height and 10 ubdEntries with different creation height
+	err := createOldStateUnbonding(t, duplicateCreationHeight, valAddr, accAddr, cdc, store)
+	require.NoError(t, err)
 
 	legacySubspace := newMockSubspace(types.DefaultParams())
-	require.NoError(t, v4.MigrateStore(ctx, storeKey, cdc, legacySubspace))
 
-	var res types.Params
-	bz := store.Get(v4.ParamsKey)
-	require.NoError(t, cdc.Unmarshal(bz, &res))
-	require.Equal(t, legacySubspace.ps, res)
+	testCases := []struct {
+		name        string
+		doMigration bool
+	}{
+		{
+			name:        "without state migration",
+			doMigration: false,
+		},
+		{
+			name:        "with state migration",
+			doMigration: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			if tc.doMigration {
+				require.NoError(t, v4.MigrateStore(ctx, storeKey, cdc, legacySubspace))
+			}
+
+			ubd := getUBD(t, accAddr, valAddr, store, cdc)
+			if tc.doMigration {
+				var res types.Params
+				bz := store.Get(v4.ParamsKey)
+				require.NoError(t, cdc.Unmarshal(bz, &res))
+				require.Equal(t, legacySubspace.ps, res)
+
+				// checking the updated balance for duplicateCreationHeight
+				for _, ubdEntry := range ubd.Entries {
+					if ubdEntry.CreationHeight == duplicateCreationHeight {
+						require.Equal(t, sdk.NewInt(100*10), ubdEntry.Balance)
+						break
+					}
+				}
+				require.Equal(t, 11, len(ubd.Entries))
+			} else {
+				require.Equal(t, true, true)
+				require.Equal(t, 20, len(ubd.Entries))
+			}
+		})
+	}
+}
+
+// createOldStateUnbonding will create the ubd entries with duplicate heights
+// 10 duplicate heights and 10 unique ubd with creation height
+func createOldStateUnbonding(t *testing.T, creationHeight int64, valAddr sdk.ValAddress, accAddr sdk.AccAddress, cdc codec.BinaryCodec, store storetypes.KVStore) error {
+	unbondBalance := sdk.NewInt(100)
+	completionTime := time.Now()
+	ubdEntries := make([]types.UnbondingDelegationEntry, 0, 10)
+
+	for i := int64(0); i < 10; i++ {
+		ubdEntry := types.UnbondingDelegationEntry{
+			CreationHeight: creationHeight,
+			Balance:        unbondBalance,
+			InitialBalance: unbondBalance,
+			CompletionTime: completionTime,
+		}
+		ubdEntries = append(ubdEntries, ubdEntry)
+		// creating more entries for testing the creation_heights
+		ubdEntry.CreationHeight = i + 2
+		ubdEntry.CompletionTime = completionTime.Add(time.Minute * 10)
+		ubdEntries = append(ubdEntries, ubdEntry)
+	}
+
+	ubd := types.UnbondingDelegation{
+		ValidatorAddress: valAddr.String(),
+		DelegatorAddress: accAddr.String(),
+		Entries:          ubdEntries,
+	}
+
+	// set the unbond delegation with validator and delegator
+	bz := types.MustMarshalUBD(cdc, ubd)
+	key := getUBDKey(accAddr, valAddr)
+	store.Set(key, bz)
+	return nil
+}
+
+func getUBD(t *testing.T, accAddr sdk.AccAddress, valAddr sdk.ValAddress, store storetypes.KVStore, cdc codec.BinaryCodec) types.UnbondingDelegation {
+	// get the unbonding delegations
+	var ubdRes types.UnbondingDelegation
+	ubdbz := store.Get(getUBDKey(accAddr, valAddr))
+	require.NoError(t, cdc.Unmarshal(ubdbz, &ubdRes))
+	return ubdRes
+}
+
+func getUBDKey(accAddr sdk.AccAddress, valAddr sdk.ValAddress) []byte {
+	return types.GetUBDKey(accAddr, valAddr)
 }

x/staking/migrations/v4/store.go

@@ -1,6 +1,8 @@
 package v4
 
 import (
+	"sort"
+
 	"github.com/cosmos/cosmos-sdk/codec"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -11,6 +13,22 @@ import (
 // MigrateStore performs in-place store migrations from v3 to v4.
 func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, legacySubspace exported.Subspace) error {
 	store := ctx.KVStore(storeKey)
+
+	// migrate params
+	if err := migrateParams(ctx, store, cdc, legacySubspace); err != nil {
+		return err
+	}
+
+	// migrate unbonding delegations
+	if err := migrateUBDEntries(ctx, store, cdc, legacySubspace); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// migrateParams will set the params to store from legacySubspace
+func migrateParams(ctx sdk.Context, store storetypes.KVStore, cdc codec.BinaryCodec, legacySubspace exported.Subspace) error {
 	var legacyParams types.Params
 	legacySubspace.GetParamSet(ctx, &legacyParams)
 
@@ -20,6 +38,63 @@ func MigrateStore(ctx sdk.Context, storeKey storetypes.StoreKey, cdc codec.Binar
 
 	bz := cdc.MustMarshal(&legacyParams)
 	store.Set(types.ParamsKey, bz)
+	return nil
+}
+
+// migrateUBDEntries will remove the ubdEntries with same creation_height
+// and create a new ubdEntry with updated balance and initial_balance
+func migrateUBDEntries(ctx sdk.Context, store storetypes.KVStore, cdc codec.BinaryCodec, legacySubspace exported.Subspace) error {
+	iterator := sdk.KVStorePrefixIterator(store, types.UnbondingDelegationKey)
+	defer iterator.Close()
 
+	for ; iterator.Valid(); iterator.Next() {
+		ubd := types.MustUnmarshalUBD(cdc, iterator.Value())
+
+		entriesAtSameCreationHeight := make(map[int64][]types.UnbondingDelegationEntry)
+		for _, ubdEntry := range ubd.Entries {
+			entriesAtSameCreationHeight[ubdEntry.CreationHeight] = append(entriesAtSameCreationHeight[ubdEntry.CreationHeight], ubdEntry)
+		}
+
+		creationHeights := make([]int64, 0, len(entriesAtSameCreationHeight))
+		for k := range entriesAtSameCreationHeight {
+			creationHeights = append(creationHeights, k)
+		}
+
+		sort.Slice(creationHeights, func(i, j int) bool { return creationHeights[i] < creationHeights[j] })
+
+		ubd.Entries = make([]types.UnbondingDelegationEntry, 0, len(creationHeights))
+
+		for _, h := range creationHeights {
+			ubdEntry := types.UnbondingDelegationEntry{
+				Balance:        sdk.ZeroInt(),
+				InitialBalance: sdk.ZeroInt(),
+			}
+			for _, entry := range entriesAtSameCreationHeight[h] {
+				ubdEntry.Balance = ubdEntry.Balance.Add(entry.Balance)
+				ubdEntry.InitialBalance = ubdEntry.InitialBalance.Add(entry.InitialBalance)
+				ubdEntry.CreationHeight = entry.CreationHeight
+				ubdEntry.CompletionTime = entry.CompletionTime
+			}
+			ubd.Entries = append(ubd.Entries, ubdEntry)
+		}
+
+		// set the new ubd to the store
+		setUBDToStore(ctx, store, cdc, ubd)
+	}
 	return nil
 }
+
+func setUBDToStore(ctx sdk.Context, store storetypes.KVStore, cdc codec.BinaryCodec, ubd types.UnbondingDelegation) {
+	delegatorAddress := sdk.MustAccAddressFromBech32(ubd.DelegatorAddress)
+
+	bz := types.MustMarshalUBD(cdc, ubd)
+
+	addr, err := sdk.ValAddressFromBech32(ubd.ValidatorAddress)
+	if err != nil {
+		panic(err)
+	}
+
+	key := types.GetUBDKey(delegatorAddress, addr)
+
+	store.Set(key, bz)
+}

x/staking/types/delegation.go

@@ -130,8 +130,27 @@ func NewUnbondingDelegation(
 
 // AddEntry - append entry to the unbonding delegation
 func (ubd *UnbondingDelegation) AddEntry(creationHeight int64, minTime time.Time, balance math.Int) {
-	entry := NewUnbondingDelegationEntry(creationHeight, minTime, balance)
-	ubd.Entries = append(ubd.Entries, entry)
+	// Check the entries exists with creation_height and complete_time
+	entryIndex := -1
+	for index, ubdEntry := range ubd.Entries {
+		if ubdEntry.CreationHeight == creationHeight && ubdEntry.CompletionTime.Equal(minTime) {
+			entryIndex = index
+			break
+		}
+	}
+	// entryIndex exists
+	if entryIndex != -1 {
+		ubdEntry := ubd.Entries[entryIndex]
+		ubdEntry.Balance = ubdEntry.Balance.Add(balance)
+		ubdEntry.InitialBalance = ubdEntry.InitialBalance.Add(balance)
+
+		// update the entry
+		ubd.Entries[entryIndex] = ubdEntry
+	} else {
+		// append the new unbond delegation entry
+		entry := NewUnbondingDelegationEntry(creationHeight, minTime, balance)
+		ubd.Entries = append(ubd.Entries, entry)
+	}
 }
 
 // RemoveEntry - remove entry at index i to the unbonding delegation