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

CHANGELOG.md

@@ -61,6 +61,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/group) [#17146](https://github.com/cosmos/cosmos-sdk/pull/17146) Rename x/group legacy ORM package's error codespace from "orm" to "legacy_orm", preventing collisions with ORM's error codespace "orm".
 * (x/bank) [#17170](https://github.com/cosmos/cosmos-sdk/pull/17170) Avoid empty spendable error message on send coins.
 * (x/distribution) [#17236](https://github.com/cosmos/cosmos-sdk/pull/17236) Using "validateCommunityTax" in "Params.ValidateBasic", preventing panic when field "CommunityTax" is nil.
+* (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.
 
 ### API Breaking Changes
 

runtime/module.go

@@ -142,7 +142,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/interfaces/consensus_params.go

@@ -0,0 +1,12 @@
+package interfaces
+
+import (
+	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+// ConsensusParamsGetter is an interface to retrieve consensus parameters for a given context.
+type ConsensusParamsGetter interface {
+	GetConsensusParams(ctx sdk.Context) tmproto.ConsensusParams
+}

types/module/module.go

@@ -50,6 +50,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/cosmos/cosmos-sdk/types/module/interfaces"
 )
 
 // AppModuleBasic is the standard form for basic non-dependant elements of an application module.
@@ -67,6 +68,11 @@ type HasName interface {
 	Name() string
 }
 
+// UpgradeModule is the extension interface that upgrade module
+type UpgradeModule interface {
+	IsUpgradeModule()
+}
+
 // HasGenesisBasics is the legacy interface for stateless genesis methods.
 type HasGenesisBasics interface {
 	DefaultGenesis(codec.JSONCodec) json.RawMessage
@@ -276,6 +282,7 @@ type Manager struct {
 	OrderPrepareCheckStaters []string
 	OrderPrecommiters        []string
 	OrderMigrations          []string
+	consensusParamsGetter    interfaces.ConsensusParamsGetter
 }
 
 // NewManager creates a new Manager object.
@@ -322,6 +329,12 @@ func NewManagerFromMap(moduleMap map[string]appmodule.AppModule) *Manager {
 	}
 }
 
+// WithConsensusParamsGetter sets ConsensusParamsGetter for Manager.
+func (m *Manager) WithConsensusParamsGetter(g interfaces.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,9 +713,31 @@ func (m *Manager) BeginBlock(ctx sdk.Context) (sdk.BeginBlock, error) {
 
 	for _, moduleName := range m.OrderBeginBlockers {
 		if module, ok := m.Modules[moduleName].(appmodule.HasBeginBlocker); ok {
-			err := module.BeginBlock(ctx)
-			if err != nil {
-				return sdk.BeginBlock{}, err
+			if _, ok := module.(UpgradeModule); ok {
+				if err := module.BeginBlock(ctx); 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 {
+						if cp = m.consensusParamsGetter.GetConsensusParams(ctx); cp.Block != nil {
+							ctx = ctx.WithConsensusParams(cp)
+						}
+					}
+				}
+				break
+			}
+		}
+	}
+
+	for _, moduleName := range m.OrderBeginBlockers {
+		if module, ok := m.Modules[moduleName].(appmodule.HasBeginBlocker); ok {
+			if _, ok := module.(UpgradeModule); !ok {
+				if err := module.BeginBlock(ctx); err != nil {
+					return sdk.BeginBlock{}, err
+				}
 			}
 		}
 	}

types/module/module_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"io"
+	"reflect"
 	"testing"
 
 	abci "github.com/cometbft/cometbft/abci/types"
@@ -491,6 +492,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.NewMockUpgradeModule(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/upgrade/module.go

@@ -162,6 +162,9 @@ func (am AppModule) BeginBlock(ctx context.Context) error {
 	return BeginBlocker(ctx, am.keeper)
 }
 
+// IsUpgradeModule implements the appmodule.AppModule interface.
+func (am AppModuleBasic) IsUpgradeModule() {}
+
 //
 // App Wiring Setup
 //