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

CHANGELOG.md

@@ -57,6 +57,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 ### Bug Fixes
 
 * (x/auth/types) [#16554](https://github.com/cosmos/cosmos-sdk/pull/16554) `ModuleAccount.Validate` now reports a nil `.BaseAccount` instead of panicking.
+* (migrations) [#](https://github.com/cosmos/cosmos-sdk/pull/) Add `WithConsensusParamsGetter` to allow other modules to access updated context with consensus parameters within the same block that executes migration.
+
 
 ## [v0.50.0-alpha.0](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.50.0-alpha.0) - 2023-06-07
 

runtime/module.go

@@ -140,7 +140,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 {

types/module/module.go

@@ -42,10 +42,9 @@ import (
 	"github.com/spf13/cobra"
 	"golang.org/x/exp/maps"
 
-	storetypes "cosmossdk.io/store/types"
-
 	errorsmod "cosmossdk.io/errors"
-
+	storetypes "cosmossdk.io/store/types"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/codec/types"
@@ -262,6 +261,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 {
@@ -273,6 +277,7 @@ type Manager struct {
 	OrderPrepareCheckStaters []string
 	OrderPrecommiters        []string
 	OrderMigrations          []string
+	consensusParamsGetter    ConsensusParamsGetter
 }
 
 // NewManager creates a new Manager object.
@@ -319,6 +324,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 {
@@ -701,6 +712,20 @@ func (m *Manager) BeginBlock(ctx sdk.Context) (sdk.BeginBlock, error) {
 			if err != nil {
 				return sdk.BeginBlock{}, err
 			}
+			// Manager retrieves consensus parameters from the consensusParamsGetter and
+			// updates the context with these parameters if current module is upgrade module,
+			// so that Manager uses the updated context to execute BeginBlock.
+			if m.consensusParamsGetter != nil {
+				p := 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 p.Block == nil {
+					cp := m.consensusParamsGetter.GetConsensusParams(ctx)
+					if cp != nil {
+						ctx = ctx.WithConsensusParams(*cp)
+					}
+				}
+			}
 		}
 	}
 

types/module/module_test.go

@@ -5,12 +5,14 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"reflect"
 	"testing"
 
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/log"
 	abci "github.com/cometbft/cometbft/abci/types"
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/golang/mock/gomock"
 	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/require"
@@ -481,6 +483,71 @@ 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) *tmproto.ConsensusParams {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetConsensusParams", arg0)
+	ret0, _ := ret[0].(*tmproto.ConsensusParams)
+	return ret0
+}
+
+// BeginBlock indicates an expected call of BeginBlock.
+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)