feat: support in-place migration ordering (backport #10614) (#10890)

* feat: support in-place migration ordering (#10614)

## Description

Closes: #10604

---

### Author Checklist

*All items are required. Please add a note to the item if the item is not applicable and
please add links to any relevant follow up issues.*

I have...

- [x] included the correct [type prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json) in the PR title
- [ ] added `!` to the type prefix if API or client breaking change
- [ ] targeted the correct branch (see [PR Targeting](https://github.com/cosmos/cosmos-sdk/blob/master/CONTRIBUTING.md#pr-targeting))
- [ ] provided a link to the relevant issue or specification
- [ ] followed the guidelines for [building modules](https://github.com/cosmos/cosmos-sdk/blob/master/docs/building-modules)
- [ ] included the necessary unit and integration [tests](https://github.com/cosmos/cosmos-sdk/blob/master/CONTRIBUTING.md#testing)
- [ ] added a changelog entry to `CHANGELOG.md`
- [ ] included comments for [documenting Go code](https://blog.golang.org/godoc)
- [ ] updated the relevant documentation or specification
- [ ] reviewed "Files changed" and left comments if necessary
- [ ] confirmed all CI checks have passed

### Reviewers Checklist

*All items are required. Please add a note if the item is not applicable and please add
your handle next to the items reviewed if you only reviewed selected items.*

I have...

- [ ] confirmed the correct [type prefix](https://github.com/commitizen/conventional-commit-types/blob/v3.0.0/index.json) in the PR title
- [ ] confirmed `!` in the type prefix if API or client breaking change
- [ ] confirmed all author checklist items have been addressed
- [ ] reviewed state machine logic
- [ ] reviewed API design and naming
- [ ] reviewed documentation is accurate
- [ ] reviewed tests and test coverage
- [ ] manually tested (if applicable)

(cherry picked from commit da929211d6)

# Conflicts:
#	CHANGELOG.md
#	docs/core/upgrade.md
#	types/errors/errors.go
#	types/module/module.go

* conflict fix

* Add Features section

* fix conflicts

Co-authored-by: Robert Zaremba <robert@zaremba.ch>

CHANGELOG.md

@@ -39,7 +39,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### State Machine Breaking
 
-- [#10833](https://github.com/cosmos/cosmos-sdk/pull/10833) fix reported tx gas used when block gas limit exceeded.
+* [#10833](https://github.com/cosmos/cosmos-sdk/pull/10833) fix reported tx gas used when block gas limit exceeded.
 * (auth) [\#10536](https://github.com/cosmos/cosmos-sdk/pull/10536]) Enable `SetSequence` for `ModuleAccount`.
 * (store) [#10218](https://github.com/cosmos/cosmos-sdk/pull/10218) Charge gas even when there are no entries while seeking.
 * (store) [#10247](https://github.com/cosmos/cosmos-sdk/pull/10247) Charge gas for the key length in gas meter.
@@ -48,6 +48,10 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 * (auth) [\#10022](https://github.com/cosmos/cosmos-sdk/pull/10022) `AuthKeeper` interface in `x/auth` now includes a function `HasAccount`.
 
+### Features
+
+* [\#10614](https://github.com/cosmos/cosmos-sdk/pull/10614) Support in-place migration ordering
+
 ### Improvements
 
 * [\#10486](https://github.com/cosmos/cosmos-sdk/pull/10486) store/cachekv's `Store.Write` conservatively

docs/core/upgrade.md

@@ -20,7 +20,7 @@ This document provides steps to use the In-Place Store Migrations upgrade method
 
 ## Tracking Module Versions
 
-Each module gets assigned a consensus version by the module developer. The consensus version serves as the breaking change version of the module. The SDK keeps track of all module consensus versions in the x/upgrade `VersionMap` store. During an upgrade, the difference between the old `VersionMap` stored in state and the new `VersionMap` is calculated by the Cosmos SDK. For each identified difference, the module-specific migrations are run and the respective consensus version of each upgraded module is incremented.
+Each module gets assigned a consensus version by the module developer. The consensus version serves as the breaking change version of the module. The Cosmos SDK keeps track of all module consensus versions in the x/upgrade `VersionMap` store. During an upgrade, the difference between the old `VersionMap` stored in state and the new `VersionMap` is calculated by the Cosmos SDK. For each identified difference, the module-specific migrations are run and the respective consensus version of each upgraded module is incremented.
 
 ### Consensus Version
 
@@ -28,7 +28,7 @@ The consensus version is defined on each app module by the module developer and
 
 ### Version Map
 
-The version map is a mapping of module names to consensus versions. The map is persisted to x/upgrade's state for use during in-place migrations. When migrations finish, the updated version map is persisted to state.
+The version map is a mapping of module names to consensus versions. The map is persisted to x/upgrade's state for use during in-place migrations. When migrations finish, the updated version map is persisted in the state.
 
 ## Upgrade Handlers
 
@@ -46,14 +46,14 @@ Inside these functions, you must perform any upgrade logic to include in the pro
 
 ## Running Migrations
 
-Migrations are run inside of an `UpgradeHandler` using  `app.mm.RunMigrations(ctx, cfg, vm)`. The `UpgradeHandler` functions describe the functionality to occur during an upgrade. The `RunMigration` function loops through the `VersionMap` argument and runs the migration scripts for all versions that are less than the versions of the new binary app module. After the migrations are finished, a new `VersionMap` is returned to persist the upgraded module versions to state.
+Migrations are run inside of an `UpgradeHandler` using `app.mm.RunMigrations(ctx, cfg, vm)`. The `UpgradeHandler` functions describe the functionality to occur during an upgrade. The `RunMigration` function loops through the `VersionMap` argument and runs the migration scripts for all versions that are less than the versions of the new binary app module. After the migrations are finished, a new `VersionMap` is returned to persist the upgraded module versions to state.
 
 ```go
 cfg := module.NewConfigurator(...)
 app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 
     // ...
-    // do upgrade logic
+    // additional upgrade logic
     // ...
 
     // returns a VersionMap with the updated module ConsensusVersions
@@ -65,24 +65,9 @@ To learn more about configuring migration scripts for your modules, see the [Mod
 
 ### Order Of Migrations
 
-All migrations are run in alphabetical order, except `x/auth` which is run the last, because of state dependencies between modules (you can read more in [issue #10606](https://github.com/cosmos/cosmos-sdk/issues/10606)). If you want to change the order of migration then you can run migrations in multiple stages. __Please beware that this is hacky, and make sure you understand what you are doing before running such migrations in production_. For example, you want to run `foo` last:
+By default, all migrations are run in module name alphabetical ascending order, except `x/auth` which is run last. The reason is state dependencies between x/auth and other modules (you can read more in [issue #10606](https://github.com/cosmos/cosmos-sdk/issues/10606)).
 
-```go
+If you want to change the order of migration then you should call `app.mm.SetOrderMigrations(module1, module2, ...)` in your app.go file. The function will panic if you forget to include a module in the argument list.
-app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap)  (module.VersionMap, error) {
-
-    fooFrom := fromVM["foo"]
-    fromVM["foo"] = foo.AppModule{}.ConsensusVersion()
-    toVM, err := app.mm.RunMigrations(ctx, cfg, fromVM)
-    if err != nil {
-        return toVM, err
-    }
-
-    stage2 := module.VersionMap{"foo": fooFrom}
-    _, err = app.mm.RunMigrations(ctx, cfg, stage2)
-
-    return toVM, err
-})
-```
 
 ## Adding New Modules During Upgrades
 
@@ -125,7 +110,7 @@ func (app *MyApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.R
 
 This information is used by the Cosmos SDK to detect when modules with newer versions are introduced to the app.
 
-For a new module `foo`, `InitGenesis` is called by the `RunMigration` only when there is a new module registered in the module manager and there is no `foo` entry in the `fromVM` registered in the state. Therefore, if you want to skip `InitGenesis` when a new module is added to the app, then you should set its module version in `fromVM` to the module package consensus version:
+For a new module `foo`, `InitGenesis` is called by `RunMigration` only when `foo` is registered in the module manager but it's not set in the `fromVM`. Therefore, if you want to skip `InitGenesis` when a new module is added to the app, then you should set its module version in `fromVM` to the module consensus version:
 
 ```go
 app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {

simapp/app.go

@@ -358,6 +358,9 @@ func NewSimApp(
 		paramstypes.ModuleName, upgradetypes.ModuleName, vestingtypes.ModuleName,
 	)
 
+	// Uncomment if you want to set a custom migration order here.
+	// app.mm.SetOrderMigrations(custom order)
+
 	app.mm.RegisterInvariants(&app.CrisisKeeper)
 	app.mm.RegisterRoutes(app.Router(), app.QueryRouter(), encodingConfig.Amino)
 	app.configurator = module.NewConfigurator(app.appCodec, app.MsgServiceRouter(), app.GRPCQueryRouter())

types/errors/errors.go

@@ -141,12 +141,12 @@ var (
 	// Examples: not DB domain error, file writing etc...
 	ErrIO = Register(RootCodespace, 39, "Internal IO error")
 
+	// ErrAppConfig defines an error occurred if min-gas-prices field in BaseConfig is empty.
+	ErrAppConfig = Register(RootCodespace, 40, "error in app.toml")
+
 	// ErrPanic is only set when we recover from a panic, so we know to
 	// redact potentially sensitive system info
 	ErrPanic = Register(UndefinedCodespace, 111222, "panic")
-
-	// ErrAppConfig defines an error occurred if min-gas-prices field in BaseConfig is empty.
-	ErrAppConfig = Register(RootCodespace, 40, "error in app.toml")
 )
 
 // Register returns an error instance that should be used as the base for

types/module/module.go

@@ -230,6 +230,7 @@ type Manager struct {
 	OrderExportGenesis []string
 	OrderBeginBlockers []string
 	OrderEndBlockers   []string
+	OrderMigrations    []string
 }
 
 // NewManager creates a new Manager object
@@ -253,28 +254,35 @@ func NewManager(modules ...AppModule) *Manager {
 
 // SetOrderInitGenesis sets the order of init genesis calls
 func (m *Manager) SetOrderInitGenesis(moduleNames ...string) {
-	m.checkForgottenModules("SetOrderInitGenesis", moduleNames)
+	m.assertNoForgottenModules("SetOrderInitGenesis", moduleNames)
 	m.OrderInitGenesis = moduleNames
 }
 
 // SetOrderExportGenesis sets the order of export genesis calls
 func (m *Manager) SetOrderExportGenesis(moduleNames ...string) {
-	m.checkForgottenModules("SetOrderExportGenesis", moduleNames)
+	m.assertNoForgottenModules("SetOrderExportGenesis", moduleNames)
 	m.OrderExportGenesis = moduleNames
 }
 
 // SetOrderBeginBlockers sets the order of set begin-blocker calls
 func (m *Manager) SetOrderBeginBlockers(moduleNames ...string) {
-	m.checkForgottenModules("SetOrderBeginBlockers", moduleNames)
+	m.assertNoForgottenModules("SetOrderBeginBlockers", moduleNames)
 	m.OrderBeginBlockers = moduleNames
 }
 
 // SetOrderEndBlockers sets the order of set end-blocker calls
 func (m *Manager) SetOrderEndBlockers(moduleNames ...string) {
-	m.checkForgottenModules("SetOrderEndBlockers", moduleNames)
+	m.assertNoForgottenModules("SetOrderEndBlockers", moduleNames)
 	m.OrderEndBlockers = moduleNames
 }
 
+// SetOrderMigrations sets the order of migrations to be run. If not set
+// then migrations will be run with an order defined in `DefaultMigrationsOrder`.
+func (m *Manager) SetOrderMigrations(moduleNames ...string) {
+	m.assertNoForgottenModules("SetOrderMigrations", moduleNames)
+	m.OrderMigrations = moduleNames
+}
+
 // RegisterInvariants registers all module invariants
 func (m *Manager) RegisterInvariants(ir sdk.InvariantRegistry) {
 	for _, module := range m.Modules {
@@ -336,24 +344,29 @@ func (m *Manager) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) map[string
 	return genesisData
 }
 
-// checkForgottenModules checks that we didn't forget any modules in the
+// assertNoForgottenModules checks that we didn't forget any modules in the
 // SetOrder* functions.
-func (m *Manager) checkForgottenModules(setOrderFnName string, moduleNames []string) {
+func (m *Manager) assertNoForgottenModules(setOrderFnName string, moduleNames []string) {
-	setOrderMap := map[string]struct{}{}
+	ms := make(map[string]bool)
 	for _, m := range moduleNames {
-		setOrderMap[m] = struct{}{}
+		ms[m] = true
 	}
-
+	var missing []string
-	if len(setOrderMap) != len(m.Modules) {
+	for m := range m.Modules {
-		panic(fmt.Sprintf("got %d modules in the module manager, but %d modules in %s", len(m.Modules), len(setOrderMap), setOrderFnName))
+		if !ms[m] {
+			missing = append(missing, m)
+		}
+	}
+	if len(missing) != 0 {
+		panic(fmt.Sprintf(
+			"%s: all modules must be defined when setting SetOrderMigrations, missing: %v", setOrderFnName, missing))
 	}
 }
 
 // MigrationHandler is the migration function that each module registers.
 type MigrationHandler func(sdk.Context) error
 
-// VersionMap is a map of moduleName -> version, where version denotes the
+// VersionMap is a map of moduleName -> version
-// version from which we should perform the migration for each module.
 type VersionMap map[string]uint64
 
 // RunMigrations performs in-place store migrations for all modules. This
@@ -380,9 +393,8 @@ type VersionMap map[string]uint64
 //      `InitGenesis` on that module.
 // - return the `updatedVM` to be persisted in the x/upgrade's store.
 //
-// Migrations are run in an alphabetical order, except x/auth which is run last. If you want
+// Migrations are run in an order defined by `Manager.OrderMigrations` or (if not set) defined by
-// to change the order then you should run migrations in multiple stages as described in
+// `DefaultMigrationsOrder` function.
-// docs/core/upgrade.md.
 //
 // As an app developer, if you wish to skip running InitGenesis for your new
 // module "foo", you need to manually pass a `fromVM` argument to this function
@@ -410,30 +422,13 @@ func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM Version
 	if !ok {
 		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
 	}
-
+	var modules = m.OrderMigrations
-	updatedVM := make(VersionMap)
+	if modules == nil {
-	// for deterministic iteration order
+		modules = DefaultMigrationsOrder(m.ModuleNames())
-	// (as some migrations depend on other modules
-	// and the order of executing migrations matters)
-	// TODO: make the order user-configurable?
-	sortedModNames := make([]string, 0, len(m.Modules))
-	hasAuth := false
-	const authModulename = "auth" // using authtypes.ModuleName causes import cycle.
-	for key := range m.Modules {
-		if key != authModulename {
-			sortedModNames = append(sortedModNames, key)
-		} else {
-			hasAuth = true
-		}
-	}
-	sort.Strings(sortedModNames)
-	// auth module must be pushed at the end because it might depend on the state from
-	// other modules, eg x/staking
-	if hasAuth {
-		sortedModNames = append(sortedModNames, authModulename)
 	}
 
-	for _, moduleName := range sortedModNames {
+	updatedVM := VersionMap{}
+	for _, moduleName := range modules {
 		module := m.Modules[moduleName]
 		fromVersion, exists := fromVM[moduleName]
 		toVersion := module.ConsensusVersion()
@@ -526,3 +521,35 @@ func (m *Manager) GetVersionMap() VersionMap {
 
 	return vermap
 }
+
+// ModuleNames returns list of all module names, without any particular order.
+func (m *Manager) ModuleNames() []string {
+	ms := make([]string, len(m.Modules))
+	i := 0
+	for m := range m.Modules {
+		ms[i] = m
+		i++
+	}
+	return ms
+}
+
+// DefaultMigrationsOrder returns a default migrations order: ascending alphabetical by module name,
+// except x/auth which will run last, see:
+// https://github.com/cosmos/cosmos-sdk/issues/10591
+func DefaultMigrationsOrder(modules []string) []string {
+	const authName = "auth"
+	out := make([]string, 0, len(modules))
+	hasAuth := false
+	for _, m := range modules {
+		if m == authName {
+			hasAuth = true
+		} else {
+			out = append(out, m)
+		}
+	}
+	sort.Strings(out)
+	if hasAuth {
+		out = append(out, authName)
+	}
+	return out
+}

types/module/module_int_test.go

@@ -0,0 +1,57 @@
+package module
+
+import (
+	"sort"
+	"testing"
+
+	"github.com/stretchr/testify/suite"
+)
+
+func TestModuleIntSuite(t *testing.T) {
+	suite.Run(t, new(TestSuite))
+}
+
+type TestSuite struct {
+	suite.Suite
+}
+
+func (s TestSuite) TestAssertNoForgottenModules() {
+	m := Manager{
+		Modules: map[string]AppModule{"a": nil, "b": nil},
+	}
+	tcs := []struct {
+		name     string
+		positive bool
+		modules  []string
+	}{
+		{"same modules", true, []string{"a", "b"}},
+		{"more modules", true, []string{"a", "b", "c"}},
+	}
+
+	for _, tc := range tcs {
+		if tc.positive {
+			m.assertNoForgottenModules("x", tc.modules)
+		} else {
+			s.Panics(func() { m.assertNoForgottenModules("x", tc.modules) })
+		}
+	}
+}
+
+func (s TestSuite) TestModuleNames() {
+	m := Manager{
+		Modules: map[string]AppModule{"a": nil, "b": nil},
+	}
+	ms := m.ModuleNames()
+	sort.Strings(ms)
+	s.Require().Equal([]string{"a", "b"}, ms)
+}
+
+func (s TestSuite) TestDefaultMigrationsOrder() {
+	require := s.Require()
+	require.Equal(
+		[]string{"auth2", "d", "z", "auth"},
+		DefaultMigrationsOrder([]string{"d", "auth", "auth2", "z"}), "alphabetical, but auth should be last")
+	require.Equal(
+		[]string{"auth2", "d", "z"},
+		DefaultMigrationsOrder([]string{"d", "auth2", "z"}), "alphabetical")
+}