cleanup, added versionmap return to RunMigrations

CHANGELOG.md

@@ -61,7 +61,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [\#8629](https://github.com/cosmos/cosmos-sdk/pull/8629) Deprecated `SetFullFundraiserPath` from `Config` in favor of `SetPurpose` and `SetCoinType`.
 * (x/upgrade) [\#8673](https://github.com/cosmos/cosmos-sdk/pull/8673) Remove IBC logic from x/upgrade. Deprecates IBC fields in an Upgrade Plan. IBC upgrade logic moved to 02-client and an IBC UpgradeProposal is added.
 * (x/bank) [\#8517](https://github.com/cosmos/cosmos-sdk/pull/8517) `SupplyI` interface and `Supply` are removed and uses `sdk.Coins` for supply tracking
-* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` includes a new argument `VersionManager` which retreives module consensus versions from upgrade's new store.
+* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) `UpgradeHandler` includes a new argument `VersionMap` which helps facilitate in-place migrations.
 * (x/auth) [\#8129](https://github.com/cosmos/cosmos-sdk/pull/8828) Updated `SigVerifiableTx.GetPubKeys` method signature to return error.
 
 

docs/architecture/adr-041-in-place-store-migrations.md

@@ -85,23 +85,23 @@ We introduce a new prefix store in `x/upgrade`'s store. This store will track ea
 0x2 | {bytes(module_name)} => BigEndian(module_consensus_version)
 ```
 
-We add a new private field `versionMap` of type `VersionMap` to `x/upgrade`'s keeper, where `VersionManager` is:
+We add a new private field `versionMap` of type `VersionMap` to `x/upgrade`'s keeper, where `VersionMap` is:
 
 ```go
 // Map of module name => new module Consensus Version.
 type VersionMap map[string]uint64
 ```
 
-This `versionMap` field can be modified via the `SetInitialVersionMap` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetInitialVersionMap` MUST be called as early as possible in the app initialization; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
+This `versionMap` field can be modified via the `SetInitialVersionMap` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetInitialVersionMap` MUST be called AFTER the `module.NewManager()` is called; in the SDK's `simapp`, it is called in the `NewSimApp` constructor function.
 
-The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return a new `VersionMap` and an error:
+The UpgradeHandler signature needs to be updated to take a `VersionMap`, as well as return an upgraded `VersionMap` and an error:
 
 ```diff
 - type UpgradeHandler func(ctx sdk.Context, plan Plan)
 + type UpgradeHandler func(ctx sdk.Context, plan Plan, versionMap VersionMap) (VersionMap, error)
 ```
 
-To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pass it into the handler. The handler runs the actual migration functions (see next section), and if successful, returns the updated ConsensusVersions of the modules to be stored in state.
+To apply an upgrade, we query the `VersionMap` from the `x/upgrade` store and pass it into the handler. The handler runs the actual migration functions (see next section), and if successful, returns an updated `VersionMap` to be stored in state.
 
 ```diff
 func (k UpgradeKeeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {

docs/run-node/cosmovisor.md

@@ -147,7 +147,7 @@ In `simapp/app.go`, find the line containing the upgrade Keeper initialisation,
 After that line, add the following snippet:
 
  ```
- app.UpgradeKeeper.SetUpgradeHandler("test1", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) {
+ app.UpgradeKeeper.SetUpgradeHandler("test1", func(ctx sdk.Context, plan upgradetypes.Plan) {
 		// Add some coins to a random account
 		addr, err := sdk.AccAddressFromBech32("cosmos18cgkqduwuh253twzmhedesw3l7v3fm37sppt58")
 		if err != nil {

simapp/app.go

@@ -541,7 +541,7 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 //   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap module.VersionMap) {
 //       return app.RunMigrations(ctx, versionMap)
 //   })
-func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) error {
+func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) (module.VersionMap, error) {
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 

simapp/app_test.go

@@ -160,7 +160,7 @@ func TestRunMigrations(t *testing.T) {
 			// Run migrations only for bank. That's why we put the initial
 			// version for bank as 1, and for all other modules, we put as
 			// their latest ConsensusVersion.
-			err = app.RunMigrations(
+			_, err = app.RunMigrations(
 				app.NewContext(true, tmproto.Header{Height: app.LastBlockHeight()}),
 				module.VersionMap{
 					"bank":         1,

types/module/module.go

@@ -338,20 +338,22 @@ type MigrationHandler func(sdk.Context) error
 type VersionMap map[string]uint64
 
 // RunMigrations performs in-place store migrations for all modules.
-func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, vm VersionMap) error {
+func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM VersionMap) (VersionMap, error) {
 	c, ok := cfg.(configurator)
 	if !ok {
-		return sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
+		return nil, sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
 	}
 
+	updatedVM := make(VersionMap)
 	for moduleName, module := range m.Modules {
-		err := c.runModuleMigrations(ctx, moduleName, vm[moduleName], module.ConsensusVersion())
+		err := c.runModuleMigrations(ctx, moduleName, fromVM[moduleName], module.ConsensusVersion())
+		updatedVM[moduleName] = module.ConsensusVersion()
 		if err != nil {
-			return err
+			return nil, err
 		}
 	}
 
-	return nil
+	return updatedVM, nil
 }
 
 // BeginBlock performs begin block functionality for all modules. It creates a

x/upgrade/keeper/keeper.go

@@ -40,7 +40,6 @@ func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey sdk.StoreKey, cdc cod
 		storeKey:           storeKey,
 		cdc:                cdc,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
-		versionMap:         module.VersionMap{},
 	}
 }
 

x/upgrade/spec/02_state.md

@@ -16,4 +16,4 @@ by the corresponding module name of type `string`.
 - ConsensusVersion: `0x2 | byte(module name)  -> BigEndian(Module Consensus Version)`
 
 
-The `x/upgrade` module contains no genesis state.
+The `x/upgrade` module contains no genesis state.