feat: support in-place migration ordering

docs/core/upgrade.md

@@ -22,20 +22,6 @@ This document provides steps to use the In-Place Store Migrations upgrade method
 
 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.
 
-## Genesis State
-
-When starting a new chain, the consensus version of each module must be saved to state during the application's genesis. To save the consensus version, add the following line to the `InitChainer` method in `app.go`:
-
-```diff
-func (app *MyApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.ResponseInitChain {
-  ...
-+ app.UpgradeKeeper.SetModuleVersionMap(ctx, app.mm.GetVersionMap())
-  ...
-}
-```
-
-This information is used by the Cosmos SDK to detect when modules with newer versions are introduced to the app.
-
 ### Consensus Version
 
 The consensus version is defined on each app module by the module developer and serves as the breaking change version of the module. The consensus version informs the Cosmos SDK on which modules need to be upgraded. For example, if the bank module was version 2 and an upgrade introduces bank module 3, the Cosmos SDK upgrades the bank module and runs the "version 2 to 3" migration script.
@@ -64,20 +50,38 @@ Migrations are run inside of an `UpgradeHandler` using  `app.mm.RunMigrations(ct
 
 ```go
 cfg := module.NewConfigurator(...)
-app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 
     // ...
     // do upgrade logic
     // ...
 
-    // RunMigrations returns the VersionMap
-    // with the updated module ConsensusVersions
-    return app.mm.RunMigrations(ctx, vm)
+    orderedVersions = module.DefaultMigrationsOrder(fromVM)
+    // returns a VersionMap with the updated module ConsensusVersions
+    return app.mm.RunMigrations(ctx, orderedVersions)
 })
 ```
 
 To learn more about configuring migration scripts for your modules, see the [Module Upgrade Guide](../building-modules/upgrade.md).
 
+### Order Of Migrations
+
+All migrations are run in (priority desc, alphabetical asc) based on the values in the `fromVM` map. By default all modules have the same priority = 100, except `x/auth` which is run last and has priority = 99. 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)).
+
+If you want to change the order of migration then you need to change module priority number. For example, you want to run `foo` last then you can use this code in `SetUpgradeHandler`:
+
+```go
+    max := 0
+    for _, mv := range orderedVersions {
+        if mv.Priority > max {
+            max = mv.Priority
+        }
+    }
+    orderedVersions["foo"].Priority = max + 1
+    return app.mm.RunMigrations(ctx, cfg, orderedVersions)
+})
+```
+
 ## Adding New Modules During Upgrades
 
 You can introduce entirely new modules to the application during an upgrade. New modules are recognized because they have not yet been registered in `x/upgrade`'s `VersionMap` store. In this case, `RunMigrations` calls the `InitGenesis` function from the corresponding module to set up its initial state.
@@ -105,7 +109,35 @@ if upgradeInfo.Name == "my-plan" && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.
 }
 ```
 
-## Overwriting Genesis Functions
+## Genesis State
+
+When starting a new chain, the consensus version of each module MUST be saved to state during the application's genesis. To save the consensus version, add the following line to the `InitChainer` method in `app.go`:
+
+```diff
+func (app *MyApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.ResponseInitChain {
+  ...
++ app.UpgradeKeeper.SetModuleVersionMap(ctx, app.mm.GetVersionMap())
+  ...
+}
+```
+
+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:
+
+```go
+app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+    // ...
+
+    // Set foo's version to the latest ConsensusVersion in the VersionMap.
+    // This will skip running InitGenesis on Foo
+    fromVM[foo.ModuleName] = foo.AppModule{}.ConsensusVersion()
+
+    return app.mm.RunMigrations(ctx, fromVM)
+})
+```
+
+### Overwriting Genesis Functions
 
 The Cosmos SDK offers modules that the application developer can import in their app. These modules often have an `InitGenesis` function already defined.
 
@@ -118,32 +150,17 @@ You MUST manually set the consensus version in the version map passed to the `Up
 ```go
 import foo "github.com/my/module/foo"
 
-app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap)  (module.VersionMap, error) {
+app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap)  (module.VersionMap, error) {
 
     // Register the consensus version in the version map
     // to avoid the SDK from triggering the default
     // InitGenesis function.
-    vm["foo"] = foo.AppModule{}.ConsensusVersion()
+    fromVM["foo"] = foo.AppModule{}.ConsensusVersion()
 
     // Run custom InitGenesis for foo
     app.mm["foo"].InitGenesis(ctx, app.appCodec, myCustomGenesisState)
 
-    return app.mm.RunMigrations(ctx, cfg, vm)
-})
-```
-
-If you do not have a custom genesis function and want to skip the module's default genesis function, you can simply register the module with the version map in the `UpgradeHandler` as shown in the example:
-
-```go
-import foo "github.com/my/module/foo"
-
-app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap)  (module.VersionMap, error) {
-
-    // Set foo's version to the latest ConsensusVersion in the VersionMap.
-    // This will skip running InitGenesis on Foo
-    vm["foo"] = foo.AppModule{}.ConsensusVersion()
-
-    return app.mm.RunMigrations(ctx, cfg, vm)
+    return app.mm.RunMigrations(ctx, cfg, fromVM)
 })
 ```
 

types/module/module.go

@@ -344,10 +344,14 @@ func (m *Manager) ExportGenesis(ctx sdk.Context, cdc codec.JSONCodec) map[string
 // 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
-// version from which we should perform the migration for each module.
+// VersionMap is a map of moduleName -> version
 type VersionMap map[string]uint64
 
+type MigrationOrder struct {
+	StateVersion uint64
+	Priority     uint64
+}
+
 // RunMigrations performs in-place store migrations for all modules. This
 // function MUST be called insde an x/upgrade UpgradeHandler.
 //
@@ -372,6 +376,10 @@ 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
+// to change the order then you should run migrations in multiple stages as described in
+// 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
 // foo's module version set to its latest ConsensusVersion. That way, the diff
@@ -393,23 +401,14 @@ type VersionMap map[string]uint64
 //   })
 //
 // Please also refer to docs/core/upgrade.md for more information.
-func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM VersionMap) (VersionMap, error) {
+func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM map[string]MigrationOrder) (VersionMap, error) {
 	c, ok := cfg.(configurator)
 	if !ok {
 		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
 	}
 
-	updatedVM := make(VersionMap)
-	// for deterministic iteration order
-	// (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))
-	for key := range m.Modules {
-		sortedModNames = append(sortedModNames, key)
-	}
-	sort.Strings(sortedModNames)
-
+	sortedModNames := sortMigrations(m.Modules, fromVM)
+	updatedVM := VersionMap{}
 	for _, moduleName := range sortedModNames {
 		module := m.Modules[moduleName]
 		fromVersion, exists := fromVM[moduleName]
@@ -496,3 +495,37 @@ func (m *Manager) GetVersionMap() VersionMap {
 
 	return vermap
 }
+
+// TODO: write tests
+// deterministic migration order
+func sortMigrations(modules map[string]AppModule, fromVM map[string]MigrationOrder) []string {
+	sortedModNames := make([]string, 0, len(modules))
+	for key := range modules {
+		sortedModNames = append(sortedModNames, key)
+	}
+	sort.Slice(sortedModNames, func(i, j int) bool {
+		a, b := sortedModNames[i], sortedModNames[j]
+		aVer, aOK := fromVM[a]
+		bVer, bOK := fromVM[b]
+		// Existing modules should be migrated before new modules
+		if aOK != bOK {
+			return aOK
+		}
+		return aVer.Priority > bVer.Priority || a <= b
+	})
+	return sortedModNames
+}
+
+// Returns a default migrations ordres, where all modules have order=100 except x/auth having
+// order=99 (this way x/auth is run last).
+func DefaultMigrationsOrder(fromVM VersionMap) map[string]MigrationOrder {
+	out := make(map[string]MigrationOrder, len(fromVM))
+	for k, v := range fromVM {
+		if k == "auth" {
+			out[k] = MigrationOrder{v, 99}
+		} else {
+			out[k] = MigrationOrder{v, 100}
+		}
+	}
+	return out
+}