fix: Add MigrationModuleManager to handle migration of upgrade module before other modules

CHANGELOG.md

@@ -66,6 +66,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 * [#16547](https://github.com/cosmos/cosmos-sdk/pull/16547) Ensure a transaction's gas limit cannot exceed the block gas limit.
 * (x/auth/types) [#16554](https://github.com/cosmos/cosmos-sdk/pull/16554) `ModuleAccount.Validate` now reports a nil `.BaseAccount` instead of panicking.
+* (migrations) [#16583](https://github.com/cosmos/cosmos-sdk/pull/16583) Add `WithConsensusParamsGetter` to allow other modules to access updated context with consensus parameters within the same block that executes migration.
 * (baseapp) [#16613](https://github.com/cosmos/cosmos-sdk/pull/16613) Ensure each message in a transaction has a registered handler, otherwise `CheckTx` will fail.
 
 ## [v0.50.0-alpha.0](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.50.0-alpha.0) - 2023-06-07

runtime/module.go

@@ -139,7 +139,7 @@ func SetupAppBuilder(inputs AppInputs) {
 	app.config = inputs.Config
 	app.appConfig = inputs.AppConfig
 	app.logger = inputs.Logger
-	app.ModuleManager = module.NewManagerFromMap(inputs.Modules)
+	app.ModuleManager = module.NewManagerFromMap(inputs.Modules).WithConsensusParamsGetter(app)
 
 	for name, mod := range inputs.Modules {
 		if customBasicMod, ok := inputs.CustomModuleBasics[name]; ok {

simapp/app.go

@@ -321,7 +321,7 @@ func NewSimApp(
 		logger,
 	)
 	app.StakingKeeper = stakingkeeper.NewKeeper(
-		appCodec, runtime.NewKVStoreService(keys[stakingtypes.StoreKey]), app.AccountKeeper, app.BankKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
+		appCodec, runtime.NewKVStoreService(keys[stakingtypes.StoreKey]), app.AccountKeeper, app.BankKeeper, app.ConsensusParamsKeeper, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 	app.MintKeeper = mintkeeper.NewKeeper(appCodec, runtime.NewKVStoreService(keys[minttypes.StoreKey]), app.StakingKeeper, app.AccountKeeper, app.BankKeeper, authtypes.FeeCollectorName, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 

types/module/module.go

@@ -36,6 +36,7 @@ import (
 	"sort"
 
 	abci "github.com/cometbft/cometbft/abci/types"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/spf13/cobra"
 	"golang.org/x/exp/maps"
@@ -261,6 +262,11 @@ func (GenesisOnlyAppModule) EndBlock(sdk.Context) ([]abci.ValidatorUpdate, error
 	return []abci.ValidatorUpdate{}, nil
 }
 
+// ConsensusParamsGetter is an interface to retrieve consensus parameters for a given context.
+type ConsensusParamsGetter interface {
+	GetConsensusParams(ctx sdk.Context) tmproto.ConsensusParams
+}
+
 // Manager defines a module manager that provides the high level utility for managing and executing
 // operations for a group of modules
 type Manager struct {
@@ -272,6 +278,7 @@ type Manager struct {
 	OrderPrepareCheckStaters []string
 	OrderPrecommiters        []string
 	OrderMigrations          []string
+	consensusParamsGetter    ConsensusParamsGetter
 }
 
 // NewManager creates a new Manager object.
@@ -318,6 +325,12 @@ func NewManagerFromMap(moduleMap map[string]appmodule.AppModule) *Manager {
 	}
 }
 
+// WithConsensusParamsGetter sets ConsensusParamsGetter for Manager.
+func (m *Manager) WithConsensusParamsGetter(g ConsensusParamsGetter) *Manager {
+	m.consensusParamsGetter = g
+	return m
+}
+
 // SetOrderInitGenesis sets the order of init genesis calls
 func (m *Manager) SetOrderInitGenesis(moduleNames ...string) {
 	m.assertNoForgottenModules("SetOrderInitGenesis", moduleNames, func(moduleName string) bool {
@@ -700,6 +713,17 @@ func (m *Manager) BeginBlock(ctx sdk.Context) (sdk.BeginBlock, error) {
 			if err != nil {
 				return sdk.BeginBlock{}, err
 			}
+			if m.consensusParamsGetter != nil {
+				cp := ctx.ConsensusParams()
+				// Manager skips this step if Block is non-nil since upgrade module is expected to set this params
+				// and consensus parameters should not be overwritten.
+				if cp.Block == nil {
+					cp = m.consensusParamsGetter.GetConsensusParams(ctx)
+					if cp.Block != nil {
+						ctx = ctx.WithConsensusParams(cp)
+					}
+				}
+			}
 		}
 	}
 

types/module/module_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"reflect"
 	"testing"
 
 	abci "github.com/cometbft/cometbft/abci/types"
@@ -482,6 +483,70 @@ func TestCoreAPIManager_BeginBlock(t *testing.T) {
 	require.EqualError(t, err, "some error")
 }
 
+type MockConsensusParamGetter struct {
+	ctrl     *gomock.Controller
+	recorder *MockConsensusParamGetterRecorder
+}
+
+type MockConsensusParamGetterRecorder struct {
+	mock *MockConsensusParamGetter
+}
+
+func NewMockConsensusParamGetter(ctrl *gomock.Controller) *MockConsensusParamGetter {
+	mock := &MockConsensusParamGetter{ctrl: ctrl}
+	mock.recorder = &MockConsensusParamGetterRecorder{mock}
+	return mock
+}
+
+func (m *MockConsensusParamGetter) EXPECT() *MockConsensusParamGetterRecorder {
+	return m.recorder
+}
+
+func (m *MockConsensusParamGetter) GetConsensusParams(arg0 sdk.Context) cmtproto.ConsensusParams {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetConsensusParams", arg0)
+	ret0, _ := ret[0].(cmtproto.ConsensusParams)
+	return ret0
+}
+
+func (mr *MockConsensusParamGetterRecorder) GetConsensusParams(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetConsensusParams", reflect.TypeOf((*MockConsensusParamGetter)(nil).GetConsensusParams), arg0)
+}
+
+func TestManager_BeginBlock_WithConsensusParamsGetter(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	t.Cleanup(mockCtrl.Finish)
+
+	mockUpgradeModule := mock.NewMockCoreAppModule(mockCtrl)
+	mockParamsGetter := NewMockConsensusParamGetter(mockCtrl)
+
+	// Create a Manager with the mock consensusParamsGetter and an upgrade module
+	m := module.NewManagerFromMap(map[string]appmodule.AppModule{
+		"upgrade": mockUpgradeModule,
+	}).WithConsensusParamsGetter(mockParamsGetter)
+
+	// Create a context with a nil consensus params object and an empty event manager
+	ctx := sdk.Context{}
+
+	mockUpgradeModule.EXPECT().BeginBlock(gomock.Any()).Times(1).Return(nil)
+	mockParamsGetter.EXPECT().GetConsensusParams(gomock.Any()).Times(1).Return(cmtproto.ConsensusParams{
+		Block: &cmtproto.BlockParams{MaxBytes: 1000},
+	})
+	_, err := m.BeginBlock(ctx)
+	require.NoError(t, err)
+
+	// Create a context with a consensus params object and an empty event manager
+	ctx = sdk.Context{}.WithConsensusParams(cmtproto.ConsensusParams{
+		Block: &cmtproto.BlockParams{MaxBytes: 1000},
+	})
+
+	mockUpgradeModule.EXPECT().BeginBlock(gomock.Any()).Times(1).Return(nil)
+	mockParamsGetter.EXPECT().GetConsensusParams(gomock.Any()).Times(0)
+	_, err = m.BeginBlock(ctx)
+	require.NoError(t, err)
+}
+
 func TestCoreAPIManager_EndBlock(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	t.Cleanup(mockCtrl.Finish)

x/consensus/types/expected_keepers.go

@@ -0,0 +1,9 @@
+package types
+
+import (
+	context "context"
+)
+
+type ConsensusKeeper interface {
+	Params(ctx context.Context, _ *QueryParamsRequest) (*QueryParamsResponse, error)
+}

x/staking/keeper/keeper.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	consensustypes "github.com/cosmos/cosmos-sdk/x/consensus/types"
 	"github.com/cosmos/cosmos-sdk/x/staking/types"
 )
 
@@ -23,12 +24,13 @@ var _ types.DelegationSet = Keeper{}
 
 // Keeper of the x/staking store
 type Keeper struct {
-	storeService storetypes.KVStoreService
-	cdc          codec.BinaryCodec
-	authKeeper   types.AccountKeeper
-	bankKeeper   types.BankKeeper
-	hooks        types.StakingHooks
-	authority    string
+	storeService    storetypes.KVStoreService
+	cdc             codec.BinaryCodec
+	authKeeper      types.AccountKeeper
+	bankKeeper      types.BankKeeper
+	hooks           types.StakingHooks
+	consensusKeeper consensustypes.ConsensusKeeper
+	authority       string
 }
 
 // NewKeeper creates a new staking Keeper instance
@@ -37,6 +39,7 @@ func NewKeeper(
 	storeService storetypes.KVStoreService,
 	ak types.AccountKeeper,
 	bk types.BankKeeper,
+	ck consensustypes.ConsensusKeeper,
 	authority string,
 ) *Keeper {
 	// ensure bonded and not bonded module accounts are set
@@ -54,12 +57,13 @@ func NewKeeper(
 	}
 
 	return &Keeper{
-		storeService: storeService,
-		cdc:          cdc,
-		authKeeper:   ak,
-		bankKeeper:   bk,
-		hooks:        nil,
-		authority:    authority,
+		storeService:    storeService,
+		cdc:             cdc,
+		authKeeper:      ak,
+		bankKeeper:      bk,
+		consensusKeeper: ck,
+		hooks:           nil,
+		authority:       authority,
 	}
 }
 

x/staking/keeper/keeper_test.go

@@ -57,11 +57,14 @@ func (s *KeeperTestSuite) SetupTest() {
 
 	bankKeeper := stakingtestutil.NewMockBankKeeper(ctrl)
 
+	consensusKeeper := stakingtestutil.NewMockConsensusKeeper(ctrl)
+
 	keeper := stakingkeeper.NewKeeper(
 		encCfg.Codec,
 		storeService,
 		accountKeeper,
 		bankKeeper,
+		consensusKeeper,
 		authtypes.NewModuleAddress(stakingtypes.GovModuleName).String(),
 	)
 	require.NoError(keeper.SetParams(ctx, stakingtypes.DefaultParams()))

x/staking/keeper/msg_server.go

@@ -82,6 +82,13 @@ func (k msgServer) CreateValidator(ctx context.Context, msg *types.MsgCreateVali
 
 	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	cp := sdkCtx.ConsensusParams()
+	if cp.Block == nil {
+		res, err := k.consensusKeeper.Params(ctx, nil)
+		if err != nil {
+			return nil, err
+		}
+		cp = *res.Params
+	}
 	if cp.Validator != nil {
 		pkType := pk.Type()
 		hasKeyType := false

x/staking/module.go

@@ -23,6 +23,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/types/module"
 	simtypes "github.com/cosmos/cosmos-sdk/types/simulation"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	consensusparamkeeper "github.com/cosmos/cosmos-sdk/x/consensus/keeper"
 	"github.com/cosmos/cosmos-sdk/x/staking/client/cli"
 	"github.com/cosmos/cosmos-sdk/x/staking/exported"
 	"github.com/cosmos/cosmos-sdk/x/staking/keeper"
@@ -209,11 +210,12 @@ func init() {
 type ModuleInputs struct {
 	depinject.In
 
-	Config        *modulev1.Module
-	AccountKeeper types.AccountKeeper
-	BankKeeper    types.BankKeeper
-	Cdc           codec.Codec
-	StoreService  store.KVStoreService
+	Config          *modulev1.Module
+	AccountKeeper   types.AccountKeeper
+	BankKeeper      types.BankKeeper
+	consensusKeeper consensusparamkeeper.Keeper
+	Cdc             codec.Codec
+	StoreService    store.KVStoreService
 
 	// LegacySubspace is used solely for migration of x/params managed parameters
 	LegacySubspace exported.Subspace `optional:"true"`
@@ -239,6 +241,7 @@ func ProvideModule(in ModuleInputs) ModuleOutputs {
 		in.StoreService,
 		in.AccountKeeper,
 		in.BankKeeper,
+		in.consensusKeeper,
 		authority.String(),
 	)
 	m := NewAppModule(in.Cdc, k, in.AccountKeeper, in.BankKeeper, in.LegacySubspace)