feat: add max_entries parameter (#244)

## Description

This PR adds `max_entries` restaking module parameter to limit the
number of unbonding entries per (delegator, delegation target) pair.

Closes: MILK-204

<!-- Add a description of the changes that this PR introduces and the
files that
are the most critical to review. -->

---

### Author Checklist

*All items are required. Please add a note to the item if the item is
not applicable and
please add links to any relevant follow up issues.*

I have...

- [ ] included the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] added `!` to the type prefix if API or client breaking change
- [ ] targeted the correct branch (see [PR
Targeting](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#pr-targeting))
- [ ] provided a link to the relevant issue or specification
- [ ] followed the guidelines for [building
modules](https://docs.cosmos.network/v0.44/building-modules/intro.html)
- [ ] included the necessary unit and integration
[tests](https://github.com/milkyway-labs/milkyway/blob/master/CONTRIBUTING.md#testing)
- [ ] added a changelog entry to `CHANGELOG.md`
- [ ] included comments for [documenting Go
code](https://blog.golang.org/godoc)
- [ ] updated the relevant documentation or specification
- [ ] reviewed "Files changed" and left comments if necessary
- [ ] confirmed all CI checks have passed

### Reviewers Checklist

*All items are required. Please add a note if the item is not applicable
and please add
your handle next to the items reviewed if you only reviewed selected
items.*

I have...

- [ ] confirmed the correct [type
prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json)
in the PR title
- [ ] confirmed `!` in the type prefix if API or client breaking change
- [ ] confirmed all author checklist items have been addressed
- [ ] reviewed state machine logic
- [ ] reviewed API design and naming
- [ ] reviewed documentation is accurate
- [ ] reviewed tests and test coverage
- [ ] manually tested (if applicable)

app/app.go

@@ -72,6 +72,7 @@ import (
 	"github.com/milkyway-labs/milkyway/v7/app/keepers"
 	"github.com/milkyway-labs/milkyway/v7/app/upgrades"
 	v6 "github.com/milkyway-labs/milkyway/v7/app/upgrades/v6"
+	v9 "github.com/milkyway-labs/milkyway/v7/app/upgrades/v9"
 	_ "github.com/milkyway-labs/milkyway/v7/client/docs/statik"
 	liquidvestingtypes "github.com/milkyway-labs/milkyway/v7/x/liquidvesting/types"
 )
@@ -87,6 +88,7 @@ var (
 
 	Upgrades = []upgrades.Upgrade{
 		v6.Upgrade,
+		v9.Upgrade,
 	}
 )
 

app/upgrades/v9/constants.go

@@ -0,0 +1,18 @@
+package v9
+
+import (
+	storetypes "cosmossdk.io/store/types"
+
+	"github.com/milkyway-labs/milkyway/v7/app/upgrades"
+)
+
+const UpgradeName = "v9"
+
+var Upgrade = upgrades.Upgrade{
+	UpgradeName:          UpgradeName,
+	CreateUpgradeHandler: CreateUpgradeHandler,
+	StoreUpgrades: storetypes.StoreUpgrades{
+		Added:   []string{},
+		Deleted: []string{},
+	},
+}

app/upgrades/v9/upgrades.go

@@ -0,0 +1,37 @@
+package v9
+
+import (
+	"context"
+
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+
+	"github.com/milkyway-labs/milkyway/v7/app/keepers"
+	"github.com/milkyway-labs/milkyway/v7/x/restaking/types"
+)
+
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	configuration module.Configurator,
+	keepers *keepers.AppKeepers,
+) upgradetypes.UpgradeHandler {
+	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		vm, err := mm.RunMigrations(ctx, configuration, fromVM)
+		if err != nil {
+			return nil, err
+		}
+
+		// Set the default max entries parameter
+		params, err := keepers.RestakingKeeper.GetParams(ctx)
+		if err != nil {
+			return nil, err
+		}
+		params.MaxEntries = types.DefaultMaxEntries
+		err = keepers.RestakingKeeper.SetParams(ctx, params)
+		if err != nil {
+			return nil, err
+		}
+
+		return vm, nil
+	}
+}

proto/milkyway/restaking/v1/params.proto

@@ -26,4 +26,8 @@ message Params {
     (gogoproto.customtype) = "cosmossdk.io/math.LegacyDec",
     (gogoproto.nullable) = false
   ];
+
+  // MaxEntries represents the maximum number of entries for unbonding
+  // delegation.
+  uint32 max_entries = 4;
 }

x/liquidvesting/keeper/end_blocker_test.go

@@ -38,6 +38,7 @@ func (suite *KeeperTestSuite) TestKeeper_EndBlocker() {
 					7*24*time.Hour,
 					nil,
 					restakingtypes.DefaultRestakingCap,
+					restakingtypes.DefaultMaxEntries,
 				))
 				suite.Require().NoError(err)
 
@@ -125,6 +126,7 @@ func (suite *KeeperTestSuite) TestKeeper_EndBlocker() {
 					7*24*time.Hour,
 					nil,
 					restakingtypes.DefaultRestakingCap,
+					restakingtypes.DefaultMaxEntries,
 				))
 				suite.Require().NoError(err)
 

x/liquidvesting/keeper/genesis_test.go

@@ -70,6 +70,7 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 					7*24*time.Hour,
 					nil,
 					restakingtypes.DefaultRestakingCap,
+					restakingtypes.DefaultMaxEntries,
 				))
 				suite.Require().NoError(err)
 

x/restaking/keeper/alias_functions.go

@@ -650,6 +650,30 @@ func (k *Keeper) getUnbondingDelegationKeyBuilder(ud types.UnbondingDelegation)
 	}
 }
 
+// HasMaxUnbondingDelegationEntries checks if unbonding delegation has maximum number of entries.
+func (k Keeper) HasMaxUnbondingDelegationEntries(ctx context.Context, delegator string, target types.DelegationTarget) (bool, error) {
+	delType, err := types.GetDelegationTypeFromTarget(target)
+	if err != nil {
+		return false, err
+	}
+
+	ubd, found, err := k.GetUnbondingDelegation(ctx, delegator, delType, target.GetID())
+	if err != nil {
+		return false, err
+	}
+	if !found {
+		// If there's no unbonding delegation, then we know it hasn't reached the max
+		// entries
+		return false, nil
+	}
+
+	maxEntries, err := k.MaxEntries(ctx)
+	if err != nil {
+		return false, err
+	}
+	return len(ubd.Entries) >= int(maxEntries), nil
+}
+
 // SetUnbondingDelegation stores the given unbonding delegation in the store
 func (k *Keeper) SetUnbondingDelegation(ctx context.Context, ud types.UnbondingDelegation) ([]byte, error) {
 	// Get the key to be used to store the unbonding delegation
@@ -720,10 +744,14 @@ func (k *Keeper) RemoveUnbondingDelegation(ctx context.Context, ubd types.Unbond
 // an unbonding object and inserting it into the unbonding queue which will be
 // processed during the staking EndBlocker.
 func (k *Keeper) PerformUndelegation(ctx context.Context, data types.UndelegationData) (time.Time, error) {
-	// TODO: Probably we should implement this as well
-	// if k.HasMaxUnbondingDelegationEntries(ctx, delAddr, valAddr) {
-	//	 return time.Time{}, types.ErrMaxUnbondingDelegationEntries
-	// }
+	// Check if the unbonding delegation has reached the maximum number of entries
+	haxMaxEntries, err := k.HasMaxUnbondingDelegationEntries(ctx, data.Delegator, data.Target)
+	if err != nil {
+		return time.Time{}, err
+	}
+	if haxMaxEntries {
+		return time.Time{}, types.ErrMaxUnbondingDelegationEntries
+	}
 
 	// Unbond the tokens
 	returnAmount, err := k.Unbond(ctx, data)

x/restaking/keeper/genesis_test.go

@@ -336,6 +336,7 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 					30*24*time.Hour,
 					nil,
 					sdkmath.LegacyNewDec(100000),
+					100,
 				))
 				suite.Require().NoError(err)
 			},
@@ -344,6 +345,7 @@ func (suite *KeeperTestSuite) TestKeeper_ExportGenesis() {
 					30*24*time.Hour,
 					nil,
 					sdkmath.LegacyNewDec(100000),
+					100,
 				),
 			},
 		},
@@ -658,6 +660,7 @@ func (suite *KeeperTestSuite) TestKeeper_InitGenesis() {
 					30*24*time.Hour,
 					nil,
 					sdkmath.LegacyNewDec(100000),
+					100,
 				),
 			},
 			check: func(ctx sdk.Context) {
@@ -667,6 +670,7 @@ func (suite *KeeperTestSuite) TestKeeper_InitGenesis() {
 					30*24*time.Hour,
 					nil,
 					sdkmath.LegacyNewDec(100000),
+					100,
 				), params)
 			},
 		},

x/restaking/keeper/grpc_query_test.go

@@ -3500,13 +3500,13 @@ func (suite *KeeperTestSuite) TestQuerier_Params() {
 		{
 			name: "params are returned properly",
 			store: func(ctx sdk.Context) {
-				params := types.NewParams(30*24*time.Hour, []string{"uinit", "umilk"}, sdkmath.LegacyNewDec(100000))
+				params := types.NewParams(30*24*time.Hour, []string{"uinit", "umilk"}, sdkmath.LegacyNewDec(100000), 5)
 				err := suite.k.SetParams(ctx, params)
 				suite.Require().NoError(err)
 			},
 			request:   types.NewQueryParamsRequest(),
 			shouldErr: false,
-			expParams: types.NewParams(30*24*time.Hour, []string{"uinit", "umilk"}, sdkmath.LegacyNewDec(100000)),
+			expParams: types.NewParams(30*24*time.Hour, []string{"uinit", "umilk"}, sdkmath.LegacyNewDec(100000), 5),
 		},
 	}
 

x/restaking/keeper/msg_server_test.go

@@ -1492,7 +1492,7 @@ func (suite *KeeperTestSuite) TestMsgServer_UndelegatePool() {
 			},
 			store: func(ctx sdk.Context) {
 				// Set the unbonding time to 1 week
-				err := suite.k.SetParams(ctx, types.NewParams(7*24*time.Hour, nil, types.DefaultRestakingCap))
+				err := suite.k.SetParams(ctx, types.NewParams(7*24*time.Hour, nil, types.DefaultRestakingCap, types.DefaultMaxEntries))
 				suite.Require().NoError(err)
 
 				// Create the pool
@@ -1809,7 +1809,7 @@ func (suite *KeeperTestSuite) TestMsgServer_UndelegateOperator() {
 			},
 			store: func(ctx sdk.Context) {
 				// Set the unbonding time to 1 week
-				err := suite.k.SetParams(ctx, types.NewParams(7*24*time.Hour, nil, types.DefaultRestakingCap))
+				err := suite.k.SetParams(ctx, types.NewParams(7*24*time.Hour, nil, types.DefaultRestakingCap, types.DefaultMaxEntries))
 				suite.Require().NoError(err)
 
 				// Create the operator
@@ -2295,7 +2295,7 @@ func (suite *KeeperTestSuite) TestMsgServer_UndelegateService() {
 			},
 			store: func(ctx sdk.Context) {
 				// Set the unbonding time to 1 week
-				err := suite.k.SetParams(ctx, types.NewParams(7*24*time.Hour, nil, types.DefaultRestakingCap))
+				err := suite.k.SetParams(ctx, types.NewParams(7*24*time.Hour, nil, types.DefaultRestakingCap, types.DefaultMaxEntries))
 				suite.Require().NoError(err)
 
 				// Create the service

x/restaking/keeper/params.go

@@ -63,6 +63,15 @@ func (k *Keeper) RestakingCap(ctx context.Context) (math.LegacyDec, error) {
 	return params.RestakingCap, nil
 }
 
+// MaxEntries returns the maximum number of simultaneous unbonding delegations.
+func (k *Keeper) MaxEntries(ctx context.Context) (uint32, error) {
+	params, err := k.GetParams(ctx)
+	if err != nil {
+		return 0, err
+	}
+	return params.MaxEntries, nil
+}
+
 // SetParams sets module parameters
 func (k *Keeper) SetParams(ctx context.Context, params types.Params) error {
 	store := k.storeService.OpenKVStore(ctx)

x/restaking/keeper/restake_restriction_test.go

@@ -82,6 +82,7 @@ func (suite *KeeperTestSuite) TestKeeper_ValidateRestakeRestakingCap() {
 					types.DefaultUnbondingTime,
 					nil,
 					sdkmath.LegacyNewDec(5000),
+					types.DefaultMaxEntries,
 				))
 				suite.Require().NoError(err)
 			},
@@ -93,7 +94,7 @@ func (suite *KeeperTestSuite) TestKeeper_ValidateRestakeRestakingCap() {
 			store: func(ctx sdk.Context) {
 				// Set restaking cap
 				err := suite.k.SetParams(ctx, types.NewParams(
-					types.DefaultUnbondingTime, nil, sdkmath.LegacyNewDec(5000)),
+					types.DefaultUnbondingTime, nil, sdkmath.LegacyNewDec(5000), types.DefaultMaxEntries),
 				)
 				suite.Require().NoError(err)
 			},

x/restaking/keeper/unbond_test.go

@@ -0,0 +1,40 @@
+package keeper_test
+
+import (
+	"time"
+
+	"github.com/milkyway-labs/milkyway/v7/utils"
+	"github.com/milkyway-labs/milkyway/v7/x/restaking/keeper"
+	"github.com/milkyway-labs/milkyway/v7/x/restaking/types"
+)
+
+func (suite *KeeperTestSuite) TestKeeper_MaxUnbondingEntries() {
+	ctx, _ := suite.ctx.CacheContext()
+
+	params, err := suite.k.GetParams(ctx)
+	suite.Require().NoError(err)
+	params.MaxEntries = 2
+	err = suite.k.SetParams(ctx, params)
+	suite.Require().NoError(err)
+
+	delegator := "cosmos167x6ehhple8gwz5ezy9x0464jltvdpzl6qfdt4"
+	suite.fundAccount(ctx, delegator, utils.MustParseCoins("10000_000000umilk"))
+
+	// First delegate to pool
+	msgServer := keeper.NewMsgServer(suite.k)
+	_, err = msgServer.DelegatePool(ctx, types.NewMsgDelegatePool(utils.MustParseCoin("10000_000000umilk"), delegator))
+	suite.Require().NoError(err)
+
+	// Unbonding from pool for the first two times should be successful
+	_, err = msgServer.UndelegatePool(ctx, types.NewMsgUndelegatePool(utils.MustParseCoin("1000_000000umilk"), delegator))
+	suite.Require().NoError(err)
+	// Increase the block height by 1 so that a separate unbonding entry is created
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1).WithBlockTime(ctx.BlockTime().Add(5 * time.Second))
+	_, err = msgServer.UndelegatePool(ctx, types.NewMsgUndelegatePool(utils.MustParseCoin("1000_000000umilk"), delegator))
+	suite.Require().NoError(err)
+
+	// But it should fail for the third time, since it exceeds the max entries
+	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1).WithBlockTime(ctx.BlockTime().Add(5 * time.Second))
+	_, err = msgServer.UndelegatePool(ctx, types.NewMsgUndelegatePool(utils.MustParseCoin("1000_000000umilk"), delegator))
+	suite.Require().ErrorIs(err, types.ErrMaxUnbondingDelegationEntries)
+}

x/restaking/simulation/utils.go

@@ -124,7 +124,7 @@ func RandomUserPreferencesEntries(
 
 func RandomParams(r *rand.Rand) types.Params {
 	unbondingDays := time.Duration(r.Intn(7) + 1)
-	return types.NewParams(time.Hour*24*unbondingDays, nil, simulation.RandomDecAmount(r, math.LegacyNewDec(10000)))
+	return types.NewParams(time.Hour*24*unbondingDays, nil, simulation.RandomDecAmount(r, math.LegacyNewDec(10000)), uint32(r.Intn(20)+1))
 }
 
 func RandomUserPreferences(r *rand.Rand, services []servicestypes.Service) types.UserPreferences {

x/restaking/types/errors.go

@@ -20,4 +20,5 @@ var (
 	ErrPoolNotSecuringService         = errors.Register(ModuleName, 13, "pool not securing the service")
 	ErrDenomNotRestakable             = errors.Register(ModuleName, 14, "denom not restakable")
 	ErrRestakingCapExceeded           = errors.Register(ModuleName, 15, "restaking cap exceeded")
+	ErrMaxUnbondingDelegationEntries  = errors.Register(ModuleName, 16, "too many unbonding delegation entries for (delegator, delegation target) tuple")
 )

x/restaking/types/genesis_test.go

@@ -227,7 +227,7 @@ func TestGenesis_Validate(t *testing.T) {
 				nil,
 				nil,
 				nil,
-				types.NewParams(0, nil, types.DefaultRestakingCap),
+				types.NewParams(0, nil, types.DefaultRestakingCap, types.DefaultMaxEntries),
 			),
 			shouldErr: true,
 		},
@@ -323,7 +323,7 @@ func TestGenesis_Validate(t *testing.T) {
 						}),
 					),
 				},
-				types.NewParams(5*24*time.Hour, nil, sdkmath.LegacyNewDec(100000)),
+				types.NewParams(5*24*time.Hour, nil, sdkmath.LegacyNewDec(100000), types.DefaultMaxEntries),
 			),
 			shouldErr: false,
 		},

x/restaking/types/messages_test.go

@@ -269,7 +269,7 @@ func TestMsgUpdateParams_ValidateBasic(t *testing.T) {
 		{
 			name: "invalid params return error",
 			msg: types.NewMsgUpdateParams(
-				types.NewParams(0, nil, types.DefaultRestakingCap),
+				types.NewParams(0, nil, types.DefaultRestakingCap, types.DefaultMaxEntries),
 				msgUpdateParams.Authority,
 			),
 			shouldErr: true,
@@ -302,7 +302,7 @@ func TestMsgUpdateParams_ValidateBasic(t *testing.T) {
 }
 
 func TestMsgUpdateParams_GetSignBytes(t *testing.T) {
-	expected := `{"type":"milkyway/restaking/MsgUpdateParams","value":{"authority":"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd","params":{"restaking_cap":"0.000000000000000000","unbonding_time":"259200000000000"}}}`
+	expected := `{"type":"milkyway/restaking/MsgUpdateParams","value":{"authority":"cosmos13t6y2nnugtshwuy0zkrq287a95lyy8vzleaxmd","params":{"max_entries":7,"restaking_cap":"0.000000000000000000","unbonding_time":"259200000000000"}}}`
 	require.Equal(t, expected, string(msgUpdateParams.GetSignBytes()))
 }
 

x/restaking/types/params.go

@@ -10,24 +10,31 @@ import (
 
 const (
 	DefaultUnbondingTime = 3 * 24 * time.Hour
+	DefaultMaxEntries    = 7
 )
 
 var (
 	DefaultRestakingCap = math.LegacyZeroDec() // no cap
 )
 
 // NewParams returns a new Params instance
-func NewParams(unbondingTime time.Duration, allowedDenoms []string, restakingCap math.LegacyDec) Params {
+func NewParams(
+	unbondingTime time.Duration,
+	allowedDenoms []string,
+	restakingCap math.LegacyDec,
+	maxEntries uint32,
+) Params {
 	return Params{
 		UnbondingTime: unbondingTime,
 		AllowedDenoms: allowedDenoms,
 		RestakingCap:  restakingCap,
+		MaxEntries:    maxEntries,
 	}
 }
 
 // DefaultParams return a Params instance with default values set
 func DefaultParams() Params {
-	return NewParams(DefaultUnbondingTime, nil, DefaultRestakingCap)
+	return NewParams(DefaultUnbondingTime, nil, DefaultRestakingCap, DefaultMaxEntries)
 }
 
 // Validate performs basic validation of params
@@ -47,5 +54,9 @@ func (p *Params) Validate() error {
 		return fmt.Errorf("restaking cap cannot be negative: %s", p.RestakingCap)
 	}
 
+	if p.MaxEntries == 0 {
+		return fmt.Errorf("max entries must be positive: %d", p.MaxEntries)
+	}
+
 	return nil
 }