refactor: rollback EndBlock

custom/staking/abci.go

@@ -1,22 +1,21 @@
 package bank
 
 import (
+	"context"
 	"time"
 
 	abci "github.com/cometbft/cometbft/abci/types"
 
 	"github.com/cosmos/cosmos-sdk/telemetry"
-	sdk "github.com/cosmos/cosmos-sdk/types"
-
 	// "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 
 	customstakingkeeper "github.com/notional-labs/composable/v6/custom/staking/keeper"
 )
 
-// Called every block, update validator set
-func EndBlocker(ctx sdk.Context, k *customstakingkeeper.Keeper) []abci.ValidatorUpdate {
+// EndBlocker returns the end blocker for the staking module.
+func EndBlocker(ctx context.Context, k *customstakingkeeper.Keeper) ([]abci.ValidatorUpdate, error) {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyEndBlocker)
 
-	return k.BlockValidatorUpdates(ctx, ctx.BlockHeight())
+	return k.BlockValidatorUpdates(ctx)
 }

custom/staking/keeper/keeper.go

@@ -29,7 +29,7 @@ type Keeper struct {
 	authKeeper        minttypes.AccountKeeper
 }
 
-func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, height int64) []abcicometbft.ValidatorUpdate {
+func (k Keeper) BlockValidatorUpdates(ctx context.Context) ([]abcicometbft.ValidatorUpdate, error) {
 	// Calculate validator set changes.
 	//
 	// NOTE: ApplyAndReturnValidatorSetUpdates has to come before
@@ -39,7 +39,9 @@ func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, height int64) []abcicomet
 	// unbonded after the Endblocker (go from Bonded -> Unbonding during
 	// ApplyAndReturnValidatorSetUpdates and then Unbonding -> Unbonded during
 	// UnbondAllMatureValidatorQueue).
-	params := k.Stakingmiddleware.GetParams(ctx)
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	params := k.Stakingmiddleware.GetParams(sdkCtx)
+	height := sdkCtx.BlockHeight()
 	shouldExecuteBatch := (height % int64(params.BlocksPerEpoch)) == 0
 	var validatorUpdates []abcicometbft.ValidatorUpdate
 	if shouldExecuteBatch {
@@ -52,17 +54,20 @@ func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, height int64) []abcicomet
 	}
 
 	// unbond all mature validators from the unbonding queue
-	k.UnbondAllMatureValidators(ctx)
+	err := k.UnbondAllMatureValidators(ctx)
+	if err != nil {
+		return nil, err
+	}
 
 	// Remove all mature unbonding delegations from the ubd queue.
-	matureUnbonds, err := k.DequeueAllMatureUBDQueue(ctx, ctx.BlockHeader().Time)
+	matureUnbonds, err := k.DequeueAllMatureUBDQueue(ctx, sdkCtx.BlockHeader().Time)
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	for _, dvPair := range matureUnbonds {
 		addr, err := sdk.ValAddressFromBech32(dvPair.ValidatorAddress)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 		delegatorAddress := sdk.MustAccAddressFromBech32(dvPair.DelegatorAddress)
 
@@ -71,7 +76,7 @@ func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, height int64) []abcicomet
 			continue
 		}
 
-		ctx.EventManager().EmitEvent(
+		sdkCtx.EventManager().EmitEvent(
 			sdk.NewEvent(
 				types.EventTypeCompleteUnbonding,
 				sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
@@ -82,18 +87,18 @@ func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, height int64) []abcicomet
 	}
 
 	// Remove all mature redelegations from the red queue.
-	matureRedelegations, err := k.DequeueAllMatureRedelegationQueue(ctx, ctx.BlockHeader().Time)
+	matureRedelegations, err := k.DequeueAllMatureRedelegationQueue(ctx, sdkCtx.BlockHeader().Time)
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 	for _, dvvTriplet := range matureRedelegations {
 		valSrcAddr, err := sdk.ValAddressFromBech32(dvvTriplet.ValidatorSrcAddress)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 		valDstAddr, err := sdk.ValAddressFromBech32(dvvTriplet.ValidatorDstAddress)
 		if err != nil {
-			panic(err)
+			return nil, err
 		}
 		delegatorAddress := sdk.MustAccAddressFromBech32(dvvTriplet.DelegatorAddress)
 
@@ -107,7 +112,7 @@ func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, height int64) []abcicomet
 			continue
 		}
 
-		ctx.EventManager().EmitEvent(
+		sdkCtx.EventManager().EmitEvent(
 			sdk.NewEvent(
 				types.EventTypeCompleteRedelegation,
 				sdk.NewAttribute(sdk.AttributeKeyAmount, balances.String()),
@@ -118,7 +123,7 @@ func (k Keeper) BlockValidatorUpdates(ctx sdk.Context, height int64) []abcicomet
 		)
 	}
 
-	return validatorUpdates
+	return validatorUpdates, nil
 }
 
 func NewKeeper(

custom/staking/module.go

@@ -1,8 +1,10 @@
 package bank
 
 import (
-	"fmt"
+	"context"
+	abci "github.com/cometbft/cometbft/abci/types"
 
+	"fmt"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	stakingmodule "github.com/cosmos/cosmos-sdk/x/staking"
@@ -60,3 +62,7 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 		panic(fmt.Sprintf("failed to migrate x/staking from version 4 to 5: %v", err))
 	}
 }
+
+func (am AppModule) EndBlock(ctx context.Context) ([]abci.ValidatorUpdate, error) {
+	return EndBlocker(ctx, &am.keeper)
+}