feat: add EndBlocker

proto/milkyway/restaking/v1/models.proto

@@ -166,7 +166,7 @@ message DTData {
   option (gogoproto.equal) = false;
   option (gogoproto.goproto_getters) = false;
 
-  UnbondingDelegationType type = 1;
+  UnbondingDelegationType unbonding_delegation_type = 1;
   string delegator_address = 2
       [ (cosmos_proto.scalar) = "cosmos.AddressString" ];
   uint32 target_id = 3 [ (gogoproto.customname) = "TargetID" ];

x/restaking/abci.go

@@ -0,0 +1,18 @@
+package restaking
+
+import (
+	"time"
+
+	"github.com/cosmos/cosmos-sdk/telemetry"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/milkyway-labs/milkyway/x/restaking/keeper"
+	"github.com/milkyway-labs/milkyway/x/restaking/types"
+)
+
+// EndBlocker is called every block and is responsible for maturing unbonding delegations
+func EndBlocker(ctx sdk.Context, k *keeper.Keeper) error {
+	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyEndBlocker)
+
+	return k.CompleteMatureUnbondingDelegations(ctx)
+}

x/restaking/keeper/alias_functions.go

@@ -6,7 +6,6 @@ import (
 	storetypes "cosmossdk.io/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
-	"github.com/milkyway-labs/milkyway/utils"
 	operatorstypes "github.com/milkyway-labs/milkyway/x/operators/types"
 	poolstypes "github.com/milkyway-labs/milkyway/x/pools/types"
 	"github.com/milkyway-labs/milkyway/x/restaking/types"
@@ -251,6 +250,34 @@ func (k *Keeper) PerformDelegation(ctx sdk.Context, data types.DelegationData) (
 // --- Unbonding operations
 // --------------------------------------------------------------------------------------------------------------------
 
+func (k *Keeper) getUnbondingDelegationTarget(ctx sdk.Context, ubd types.UnbondingDelegation) (types.DelegationTarget, error) {
+	switch ubd.Type {
+	case types.UNBONDING_DELEGATION_TYPE_POOL:
+		pool, found := k.poolsKeeper.GetPool(ctx, ubd.TargetID)
+		if !found {
+			return nil, poolstypes.ErrPoolNotFound
+		}
+		return &pool, nil
+
+	case types.UNBONDING_DELEGATION_TYPE_OPERATOR:
+		operator, found := k.operatorsKeeper.GetOperator(ctx, ubd.TargetID)
+		if !found {
+			return nil, operatorstypes.ErrOperatorNotFound
+		}
+		return &operator, nil
+
+	case types.UNBONDING_DELEGATION_TYPE_SERVICE:
+		service, found := k.servicesKeeper.GetService(ctx, ubd.TargetID)
+		if !found {
+			return nil, servicestypes.ErrServiceNotFound
+		}
+		return &service, nil
+
+	default:
+		return nil, types.ErrInvalidDelegationType
+	}
+}
+
 // getUnbondingDelegationKeyBuilder returns the key builder for the given unbonding delegation
 func (k *Keeper) getUnbondingDelegationKeyBuilder(ud types.UnbondingDelegation) (types.UnbondingDelegationKeyBuilder, error) {
 	switch ud.Type {
@@ -269,42 +296,51 @@ func (k *Keeper) getUnbondingDelegationKeyBuilder(ud types.UnbondingDelegation)
 }
 
 // SetUnbondingDelegation stores the given unbonding delegation in the store
-func (k *Keeper) SetUnbondingDelegation(ctx sdk.Context, ud types.UnbondingDelegation, entryID uint64) error {
+func (k *Keeper) SetUnbondingDelegation(ctx sdk.Context, ud types.UnbondingDelegation) ([]byte, error) {
 	// Get the key to be used to store the unbonding delegation
 	getUnbondingDelegation, err := k.getUnbondingDelegationKeyBuilder(ud)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	unbondingDelegationKey := getUnbondingDelegation(ud.DelegatorAddress, ud.TargetID)
 
 	// Store the unbonding delegation
 	store := ctx.KVStore(k.storeKey)
 	store.Set(unbondingDelegationKey, types.MustMarshalUnbondingDelegation(k.cdc, ud))
 
-	// Set the index allowing to lookup the UnbondingDelegation by the unbondingID of an
-	// UnbondingDelegationEntry that it contains
-	store.Set(types.GetUnbondingIndexKey(entryID), unbondingDelegationKey)
-
-	// Set the type of the unbonding delegation so that we know how to deserialize id
-	store.Set(types.GetUnbondingTypeKey(entryID), utils.Uint32ToBytes(ud.TargetID))
-
-	return nil
+	return unbondingDelegationKey, nil
 }
 
 // GetUnbondingDelegation returns the unbonding delegation for the given delegator and target.
-func (k *Keeper) GetUnbondingDelegation(ctx sdk.Context, delegatorAddress string, target types.DelegationTarget) (types.UnbondingDelegation, bool) {
-	switch target.(type) {
-	case *poolstypes.Pool:
-		return k.GetPoolUnbondingDelegation(ctx, target.GetID(), delegatorAddress)
-	case *operatorstypes.Operator:
-		return k.GetOperatorUnbondingDelegation(ctx, delegatorAddress, target.GetID())
-	case *servicestypes.Service:
-		return k.GetServiceUnbondingDelegation(ctx, delegatorAddress, target.GetID())
+func (k *Keeper) GetUnbondingDelegation(
+	ctx sdk.Context, delegatorAddress string, ubdType types.UnbondingDelegationType, targetID uint32,
+) (types.UnbondingDelegation, bool) {
+	switch ubdType {
+	case types.UNBONDING_DELEGATION_TYPE_POOL:
+		return k.GetPoolUnbondingDelegation(ctx, targetID, delegatorAddress)
+	case types.UNBONDING_DELEGATION_TYPE_OPERATOR:
+		return k.GetOperatorUnbondingDelegation(ctx, delegatorAddress, targetID)
+	case types.UNBONDING_DELEGATION_TYPE_SERVICE:
+		return k.GetServiceUnbondingDelegation(ctx, delegatorAddress, targetID)
 	default:
 		return types.UnbondingDelegation{}, false
 	}
 }
 
+// RemoveUnbondingDelegation removes the unbonding delegation object and associated index.
+func (k *Keeper) RemoveUnbondingDelegation(ctx sdk.Context, ubd types.UnbondingDelegation) error {
+	// Get the key to be used to store the unbonding delegation
+	getUnbondingDelegation, err := k.getUnbondingDelegationKeyBuilder(ubd)
+	if err != nil {
+		return err
+	}
+	unbondingDelegationKey := getUnbondingDelegation(ubd.DelegatorAddress, ubd.TargetID)
+
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(unbondingDelegationKey)
+	return nil
+}
+
 // PerformUndelegation unbonds an amount of delegator shares from a given validator. It
 // will verify that the unbonding entries between the delegator and validator
 // are not exceeded and unbond the staked tokens (based on shares) by creating

x/restaking/keeper/end_blocker.go

@@ -0,0 +1,34 @@
+package keeper
+
+import (
+	"fmt"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/milkyway-labs/milkyway/x/restaking/types"
+)
+
+// CompleteMatureUnbondingDelegations runs the endblocker logic for delegations
+func (k *Keeper) CompleteMatureUnbondingDelegations(ctx sdk.Context) error {
+	// Remove all mature unbonding delegations from the ubd queue.
+	matureUnbonds := k.DequeueAllMatureUBDQueue(ctx, ctx.BlockHeader().Time)
+	for _, data := range matureUnbonds {
+
+		balances, err := k.CompleteUnbonding(ctx, data)
+		if err != nil {
+			return err
+		}
+
+		ctx.EventManager().EmitEvent(
+			sdk.NewEvent(
+				types.EventTypeCompleteUnbonding,
+				sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
+				sdk.NewAttribute(types.AttributeUnbondingDelegationType, data.UnbondingDelegationType.String()),
+				sdk.NewAttribute(types.AttributeTargetID, fmt.Sprintf("%d", data.TargetID)),
+				sdk.NewAttribute(types.AttributeKeyDelegator, data.DelegatorAddress),
+			),
+		)
+	}
+
+	return nil
+}

x/restaking/keeper/unbond.go

@@ -36,6 +36,12 @@ func (k *Keeper) IncrementUnbondingID(ctx sdk.Context) (unbondingID uint64) {
 	return unbondingID
 }
 
+// DeleteUnbondingIndex removes a mapping from UnbondingId to unbonding operation
+func (k *Keeper) DeleteUnbondingIndex(ctx sdk.Context, id uint64) {
+	store := ctx.KVStore(k.storeKey)
+	store.Delete(types.GetUnbondingIndexKey(id))
+}
+
 // ValidateUnbondAmount validates that a given unbond or redelegation amount is valid based on upon the
 // converted shares. If the amount is valid, the total amount of respective shares is returned,
 // otherwise an error is returned.
@@ -162,18 +168,31 @@ func (k *Keeper) SetUnbondingDelegationEntry(
 	id := k.IncrementUnbondingID(ctx)
 
 	// Either get the existing unbonding delegation, or create a new one
-	ubd, found := k.GetUnbondingDelegation(ctx, data.Delegator, data.Target)
+	ubdType, err := types.GetUnboningDelegationTypeFromTarget(data.Target)
+	if err != nil {
+		return types.UnbondingDelegation{}, err
+	}
+
+	ubd, found := k.GetUnbondingDelegation(ctx, data.Delegator, ubdType, data.Target.GetID())
 	if found {
 		ubd.AddEntry(creationHeight, minTime, balance, id)
 	} else {
 		ubd = data.BuildUnbondingDelegation(data.Delegator, data.Target.GetID(), creationHeight, minTime, balance, id)
 	}
 
-	err := k.SetUnbondingDelegation(ctx, ubd, id)
+	unbondingDelegationKey, err := k.SetUnbondingDelegation(ctx, ubd)
 	if err != nil {
 		return types.UnbondingDelegation{}, err
 	}
 
+	// Set the index allowing to lookup the UnbondingDelegation by the unbondingID of an
+	// UnbondingDelegationEntry that it contains
+	store := ctx.KVStore(k.storeKey)
+	store.Set(types.GetUnbondingIndexKey(id), unbondingDelegationKey)
+
+	// Set the type of the unbonding delegation so that we know how to deserialize id
+	store.Set(types.GetUnbondingTypeKey(id), utils.Uint32ToBytes(ubd.TargetID))
+
 	// Call the hook after the unbonding has been initiated
 	err = k.AfterUnbondingInitiated(ctx, id)
 	if err != nil {
@@ -183,6 +202,79 @@ func (k *Keeper) SetUnbondingDelegationEntry(
 	return ubd, nil
 }
 
+// CompleteUnbonding completes the unbonding of all mature entries in the
+// retrieved unbonding delegation object and returns the total unbonding balance
+// or an error upon failure.
+func (k *Keeper) CompleteUnbonding(ctx sdk.Context, data types.DTData) (sdk.Coins, error) {
+	// Get the unbonding delegation entry
+	ubd, found := k.GetUnbondingDelegation(ctx, data.DelegatorAddress, data.UnbondingDelegationType, data.TargetID)
+	if !found {
+		return nil, types.ErrNoUnbondingDelegation
+	}
+
+	// Get the target of the unbonding delegation
+	target, err := k.getUnbondingDelegationTarget(ctx, ubd)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get the address of the target
+	targetAddress, err := sdk.AccAddressFromBech32(target.GetAddress())
+	if err != nil {
+		return nil, err
+	}
+
+	// Get the address of the delegator
+	delegatorAddress, err := sdk.AccAddressFromBech32(ubd.DelegatorAddress)
+	if err != nil {
+		return nil, err
+	}
+
+	balances := sdk.NewCoins()
+	ctxTime := ctx.BlockHeader().Time
+
+	// Loop through all the entries and complete unbonding mature entries
+	for i := 0; i < len(ubd.Entries); i++ {
+		entry := ubd.Entries[i]
+		if entry.IsMature(ctxTime) {
+			// Remove the entry
+			ubd.RemoveEntry(int64(i))
+			i--
+
+			// Delete the index
+			k.DeleteUnbondingIndex(ctx, entry.UnbondingId)
+
+			// Track undelegation only when remaining or truncated shares are non-zero
+			if !entry.Balance.IsZero() {
+				amount := entry.Balance
+
+				// Send the coins back to the delegator
+				err = k.bankKeeper.SendCoins(ctx, targetAddress, delegatorAddress, amount)
+				if err != nil {
+					return nil, err
+				}
+
+				balances = balances.Add(amount...)
+			}
+		}
+	}
+
+	// Set the unbonding delegation or remove it if there are no more entries
+	if len(ubd.Entries) == 0 {
+		err = k.RemoveUnbondingDelegation(ctx, ubd)
+		if err != nil {
+			return nil, err
+		}
+	} else {
+		_, err = k.SetUnbondingDelegation(ctx, ubd)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return balances, nil
+}
+
 // --------------------------------------------------------------------------------------------------------------------
 // --- Unbonding queue operations
 // --------------------------------------------------------------------------------------------------------------------
@@ -214,7 +306,11 @@ func (k *Keeper) SetUBDQueueTimeSlice(ctx sdk.Context, timestamp time.Time, keys
 // InsertUBDQueue inserts an unbonding delegation to the appropriate timeslice
 // in the unbonding queue.
 func (k *Keeper) InsertUBDQueue(ctx sdk.Context, ubd types.UnbondingDelegation, completionTime time.Time) {
-	dvPair := types.DTData{Type: ubd.Type, DelegatorAddress: ubd.DelegatorAddress, TargetID: ubd.TargetID}
+	dvPair := types.DTData{
+		UnbondingDelegationType: ubd.Type,
+		DelegatorAddress:        ubd.DelegatorAddress,
+		TargetID:                ubd.TargetID,
+	}
 
 	timeSlice := k.GetUBDQueueTimeSlice(ctx, completionTime)
 	if len(timeSlice) == 0 {

x/restaking/module.go

@@ -141,6 +141,12 @@ func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) json.Raw
 // ConsensusVersion implements ConsensusVersion.
 func (AppModule) ConsensusVersion() uint64 { return consensusVersion }
 
+// EndBlock executes all ABCI EndBlock logic respective to the restaking module.
+func (am AppModule) EndBlock(ctx context.Context) error {
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	return EndBlocker(sdkCtx, am.keeper)
+}
+
 func (am AppModule) IsOnePerModuleType() {}
 
 func (am AppModule) IsAppModule() {}

x/restaking/types/errors.go

@@ -11,4 +11,5 @@ var (
 	ErrDelegationNotFound          = errors.Register(ModuleName, 4, "delegation not found")
 	ErrNotEnoughDelegationShares   = errors.Register(ModuleName, 5, "not enough delegation shares")
 	ErrInvalidDelegationType       = errors.Register(ModuleName, 6, "invalid delegation type")
+	ErrNoUnbondingDelegation       = errors.Register(ModuleName, 7, "no unbonding delegation found")
 )

x/restaking/types/events.go

@@ -1,6 +1,7 @@
 package types
 
 const (
+	EventTypeCompleteUnbonding    = "complete_unbonding"
 	EventTypeUpdateOperatorParams = "update_operator_params"
 	EventTypeUpdateServiceParams  = "update_service_params"
 	EventTypeDelegatePool         = "delegate_pool"
@@ -18,4 +19,6 @@ const (
 	AttributeKeyServiceID              = "service_id"
 	AttributeKeyNewShares              = "new_shares"
 	AttributeKeyCompletionTime         = "completion_time"
+	AttributeUnbondingDelegationType   = "unbonding_delegation"
+	AttributeTargetID                  = "target_id"
 )