fix: Add MigrationModuleManager to handle migration of upgrade module before other modules (#16583)

Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>
Co-authored-by: Julien Robert <julien@rbrt.fr>

types/module/module.go

@@ -67,6 +67,13 @@ type HasName interface {
 	Name() string
 }
 
+// UpgradeModule is the extension interface that upgrade module should implement to differentiate
+// it from other modules, migration handler need ensure the upgrade module's migration is executed
+// before the rest of the modules.
+type UpgradeModule interface {
+	IsUpgradeModule()
+}
+
 // HasGenesisBasics is the legacy interface for stateless genesis methods.
 type HasGenesisBasics interface {
 	DefaultGenesis(codec.JSONCodec) json.RawMessage
@@ -692,17 +699,31 @@ func (m Manager) RunMigrations(ctx context.Context, cfg Configurator, fromVM Ver
 	return updatedVM, nil
 }
 
-// BeginBlock performs begin block functionality for all modules. It creates a
-// child context with an event manager to aggregate events emitted from all
+// RunMigrationBeginBlock performs begin block functionality for upgrade module.
+// It takes the current context as a parameter and returns a boolean value
+// indicating whether the migration was successfully executed or not.
+func (m *Manager) RunMigrationBeginBlock(ctx sdk.Context) bool {
+	for _, moduleName := range m.OrderBeginBlockers {
+		if mod, ok := m.Modules[moduleName].(appmodule.HasBeginBlocker); ok {
+			if _, ok := mod.(UpgradeModule); ok {
+				return mod.BeginBlock(ctx) == nil
+			}
+		}
+	}
+	return false
+}
+
+// BeginBlock performs begin block functionality for non-upgrade modules. It creates a
+// child context with an event manager to aggregate events emitted from non-upgrade
 // modules.
 func (m *Manager) BeginBlock(ctx sdk.Context) (sdk.BeginBlock, error) {
 	ctx = ctx.WithEventManager(sdk.NewEventManager())
-
 	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
+				}
 			}
 		}
 	}

types/module/module_test.go

@@ -467,6 +467,33 @@ func TestCoreAPIManagerOrderSetters(t *testing.T) {
 	require.Equal(t, []string{"module3", "module2", "module1"}, mm.OrderPrecommiters)
 }
 
+func TestCoreAPIManager_RunMigrationBeginBlock(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	t.Cleanup(mockCtrl.Finish)
+
+	mockAppModule1 := mock.NewMockCoreAppModule(mockCtrl)
+	mockAppModule2 := mock.NewMockUpgradeModule(mockCtrl)
+	mm := module.NewManagerFromMap(map[string]appmodule.AppModule{
+		"module1": mockAppModule1,
+		"module2": mockAppModule2,
+	})
+	require.NotNil(t, mm)
+	require.Equal(t, 2, len(mm.Modules))
+
+	mockAppModule1.EXPECT().BeginBlock(gomock.Any()).Times(0)
+	mockAppModule2.EXPECT().BeginBlock(gomock.Any()).Times(1).Return(nil)
+	success := mm.RunMigrationBeginBlock(sdk.Context{})
+	require.Equal(t, true, success)
+
+	// test false
+	require.Equal(t, false, module.NewManager().RunMigrationBeginBlock(sdk.Context{}))
+
+	// test panic
+	mockAppModule2.EXPECT().BeginBlock(gomock.Any()).Times(1).Return(errors.New("some error"))
+	success = mm.RunMigrationBeginBlock(sdk.Context{})
+	require.Equal(t, false, success)
+}
+
 func TestCoreAPIManager_BeginBlock(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	t.Cleanup(mockCtrl.Finish)