x/upgrade: added consensus version tracking (part of ADR-041)

CHANGELOG.md

@@ -61,6 +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 `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.
 
 
@@ -79,6 +80,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/auth) [\#8522](https://github.com/cosmos/cosmos-sdk/pull/8522) Allow to query all stored accounts
 * (crypto/types) [\#8600](https://github.com/cosmos/cosmos-sdk/pull/8600) `CompactBitArray`: optimize the `NumTrueBitsBefore` method and add an `Equal` method.
 * (grpc) [\#8815](https://github.com/cosmos/cosmos-sdk/pull/8815) Add orderBy parameter to `TxsByEvents` endpoint.
+* (x/upgrade) [\#8743](https://github.com/cosmos/cosmos-sdk/pull/8743) Add tracking module versions as per ADR-041
 
 ### Bug Fixes
 

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

@@ -82,48 +82,34 @@ Each module's migration functions are specific to the module's store evolutions,
 We introduce a new prefix store in `x/upgrade`'s store. This store will track each module's current version, it can be modelized as a `map[string]uint64` of module name to module ConsensusVersion, and will be used when running the migrations (see next section for details). The key prefix used is `0x1`, and the key/value format is:
 
 ```
-0x2 | {bytes(module_name)} => LittleEndian(module_consensus_version)
+0x2 | {bytes(module_name)} => BigEndian(module_consensus_version)
 ```
+The initial state of the store is set from `app.go`'s `InitChainer` method.
 
-s
-We add a new private field `versionManager` of type `VersionManager` to `x/upgrade`'s keeper, where `VersionManager` is:
-
-```go
-type VersionManager interface {
-    GetConsensusVersions() VersionMap
-}
-
-// Map of module name => new module Consensus Version.
-type VersionMap map[string]uint64
-```
-
-This `versionManager` field can be modified via the `SetVersionManager` field, and will allow the upgrade keeper to know the current versions of loaded modules. `SetVersionManager` MUST be called as early as possible in the app initialization; 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 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) error
++ 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, the current ConsensusVersions of all loaded modules will be stored into 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) {
     // --snip--
 -   handler(ctx, plan)
-+   err := handler(ctx, plan, k.GetConsensusVersions()) // k.GetConsensusVersions() fetches the VersionMap stored in state.
++   updatedVM, err := handler(ctx, plan, k.GetModuleVersionMap(ctx)) // k.GetModuleVersionMap() fetches the VersionMap stored in state.
 +   if err != nil {
 +       return err
 +   }
 +
-+   // Get the current ConsensusVersions of the loaded modules (retrieved from
-+   // `k.versionManager`), and save them to state.
-+   k.SetCurrentConsensusVersions()
++   // Set the updated consensus versions to state
++   k.SetModuleVersionMap(ctx, updatedVM)
 }
 ```
 
-An gRPC query endpoint to query the `VersionMap` stored in `x/upgrade`'s state will also be added, so that app developers can double-check the `VersionMap` before the upgrade handler runs.
+A gRPC query endpoint to query the `VersionMap` stored in `x/upgrade`'s state will also be added, so that app developers can double-check the `VersionMap` before the upgrade handler runs.
 
 ### Running Migrations
 
@@ -139,8 +125,8 @@ If a required migration is missing (e.g. if it has not been registered in the `C
 In practice, the `RunMigrations` method should be called from inside an `UpgradeHandler`.
 
 ```go
-app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, versionMap VersionMap) error {
-    return app.mm.RunMigrations(ctx, versionMap)
+app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap)  (module.VersionMap, error) {
+    return app.mm.RunMigrations(ctx, vm)
 })
 ```
 

simapp/app.go

@@ -421,6 +421,7 @@ func (app *SimApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.
 	if err := json.Unmarshal(req.AppStateBytes, &genesisState); err != nil {
 		panic(err)
 	}
+	app.UpgradeKeeper.SetModuleVersionMap(ctx, app.mm.GetVersionMap())
 	return app.mm.InitGenesis(ctx, app.appCodec, genesisState)
 }
 
@@ -535,16 +536,10 @@ func (app *SimApp) RegisterTendermintService(clientCtx client.Context) {
 //
 // Example:
 //   cfg := module.NewConfigurator(...)
-//   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan) {
-//       err := app.RunMigrations(ctx, module.MigrationMap{
-//           "bank": 1,     // Migrate x/bank from v1 to current x/bank's ConsensusVersion
-//           "staking": 8,  // Migrate x/staking from v8 to current x/staking's ConsensusVersion
-//      })
-//      if err != nil {
-//           panic(err)
-//      }
+//   app.UpgradeKeeper.SetUpgradeHandler("store-migration", func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) {
+//       return app.RunMigrations(ctx, vm)
 //   })
-func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.MigrationMap) 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,9 +160,9 @@ 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.MigrationMap{
+				module.VersionMap{
 					"bank":         1,
 					"auth":         auth.AppModule{}.ConsensusVersion(),
 					"authz":        authz.AppModule{}.ConsensusVersion(),
@@ -190,3 +190,27 @@ func TestRunMigrations(t *testing.T) {
 		})
 	}
 }
+
+func TestUpgradeStateOnGenesis(t *testing.T) {
+	encCfg := MakeTestEncodingConfig()
+	db := dbm.NewMemDB()
+	app := NewSimApp(log.NewTMLogger(log.NewSyncWriter(os.Stdout)), db, nil, true, map[int64]bool{}, DefaultNodeHome, 0, encCfg, EmptyAppOptions{})
+	genesisState := NewDefaultGenesisState(encCfg.Marshaler)
+	stateBytes, err := json.MarshalIndent(genesisState, "", "  ")
+	require.NoError(t, err)
+
+	// Initialize the chain
+	app.InitChain(
+		abci.RequestInitChain{
+			Validators:    []abci.ValidatorUpdate{},
+			AppStateBytes: stateBytes,
+		},
+	)
+
+	// make sure the upgrade keeper has version map in state
+	ctx := app.NewContext(false, tmproto.Header{})
+	vm := app.UpgradeKeeper.GetModuleVersionMap(ctx)
+	for v, i := range app.mm.Modules {
+		require.Equal(t, vm[v], i.ConsensusVersion())
+	}
+}

types/module/module.go

@@ -333,25 +333,27 @@ func (m *Manager) ExportGenesis(ctx sdk.Context, cdc codec.JSONMarshaler) map[st
 // MigrationHandler is the migration function that each module registers.
 type MigrationHandler func(sdk.Context) error
 
-// MigrationMap is a map of moduleName -> version, where version denotes the
+// VersionMap is a map of moduleName -> version, where version denotes the
 // version from which we should perform the migration for each module.
-type MigrationMap map[string]uint64
+type VersionMap map[string]uint64
 
 // RunMigrations performs in-place store migrations for all modules.
-func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, migrateFromVersions MigrationMap) 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, migrateFromVersions[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
@@ -395,3 +397,15 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) abci.Respo
 		Events:           ctx.EventManager().ABCIEvents(),
 	}
 }
+
+// GetVersionMap gets consensus version from all modules
+func (m *Manager) GetVersionMap() VersionMap {
+	vermap := make(VersionMap)
+	for _, v := range m.Modules {
+		version := v.ConsensusVersion()
+		name := v.Name()
+		vermap[name] = version
+	}
+
+	return vermap
+}

x/upgrade/abci_test.go

@@ -105,7 +105,9 @@ func VerifyDoUpgrade(t *testing.T) {
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan) {})
+	s.keeper.SetUpgradeHandler("test", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -121,7 +123,9 @@ func VerifyDoUpgradeWithCtx(t *testing.T, newCtx sdk.Context, proposalName strin
 	})
 
 	t.Log("Verify that the upgrade can be successfully applied with a handler")
-	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan) {})
+	s.keeper.SetUpgradeHandler(proposalName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return vm, nil
+	})
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -133,7 +137,10 @@ func TestHaltIfTooNew(t *testing.T) {
 	s := setupTest(10, map[int64]bool{})
 	t.Log("Verify that we don't panic with registered plan not in database at all")
 	var called int
-	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan) { called++ })
+	s.keeper.SetUpgradeHandler("future", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		called++
+		return vm, nil
+	})
 
 	newCtx := s.ctx.WithBlockHeight(s.ctx.BlockHeight() + 1).WithBlockTime(time.Now())
 	req := abci.RequestBeginBlock{Header: newCtx.BlockHeader()}

x/upgrade/keeper/grpc_query_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/simapp"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
@@ -110,7 +111,9 @@ func (suite *UpgradeTestSuite) TestAppliedCurrentPlan() {
 				suite.app.UpgradeKeeper.ScheduleUpgrade(suite.ctx, plan)
 
 				suite.ctx = suite.ctx.WithBlockHeight(expHeight)
-				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(ctx sdk.Context, plan types.Plan) {})
+				suite.app.UpgradeKeeper.SetUpgradeHandler(planName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+					return vm, nil
+				})
 				suite.app.UpgradeKeeper.ApplyUpgrade(suite.ctx, plan)
 
 				req = &types.QueryAppliedPlanRequest{Name: planName}

x/upgrade/keeper/keeper.go

@@ -16,6 +16,7 @@ import (
 	store "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
+	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
@@ -48,6 +49,39 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
+// SetModuleVersionMap saves a given version map to state
+func (k Keeper) SetModuleVersionMap(ctx sdk.Context, vm module.VersionMap) {
+	if len(vm) > 0 {
+		store := ctx.KVStore(k.storeKey)
+		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
+		for modName, ver := range vm {
+			nameBytes := []byte(modName)
+			verBytes := make([]byte, 8)
+			binary.BigEndian.PutUint64(verBytes, ver)
+			versionStore.Set(nameBytes, verBytes)
+		}
+	}
+}
+
+// GetModuleVersionMap returns a map of key module name and value module consensus version
+// as defined in ADR-041.
+func (k Keeper) GetModuleVersionMap(ctx sdk.Context) module.VersionMap {
+	store := ctx.KVStore(k.storeKey)
+	it := sdk.KVStorePrefixIterator(store, []byte{types.VersionMapByte})
+
+	vm := make(module.VersionMap)
+	defer it.Close()
+	for ; it.Valid(); it.Next() {
+		moduleBytes := it.Key()
+		// first byte is prefix key, so we remove it here
+		name := string(moduleBytes[1:])
+		moduleVersion := binary.BigEndian.Uint64(it.Value())
+		vm[name] = moduleVersion
+	}
+
+	return vm
+}
+
 // ScheduleUpgrade schedules an upgrade based on the specified plan.
 // If there is another Plan already scheduled, it will overwrite it
 // (implicitly cancelling the current plan)
@@ -187,7 +221,12 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic("ApplyUpgrade should never be called without first checking HasHandler")
 	}
 
-	handler(ctx, plan)
+	updatedVM, err := handler(ctx, plan, k.GetModuleVersionMap(ctx))
+	if err != nil {
+		panic(err)
+	}
+
+	k.SetModuleVersionMap(ctx, updatedVM)
 
 	// Must clear IBC state after upgrade is applied as it is stored separately from the upgrade plan.
 	// This will prevent resubmission of upgrade msg after upgrade is already completed.

x/upgrade/keeper/keeper_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/simapp"
 	store "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/keeper"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
@@ -115,7 +116,9 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 				Height: 123450000,
 			},
 			setup: func() {
-				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(_ sdk.Context, _ types.Plan) {})
+				s.app.UpgradeKeeper.SetUpgradeHandler("all-good", func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+					return vm, nil
+				})
 				s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{
 					Name:   "all-good",
 					Info:   "some text here",
@@ -191,6 +194,29 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 
 }
 
+// Tests that the underlying state of x/upgrade is set correctly after
+// an upgrade.
+func (s *KeeperTestSuite) TestMigrations() {
+	initialVM := module.VersionMap{"bank": uint64(1)}
+	s.app.UpgradeKeeper.SetModuleVersionMap(s.ctx, initialVM)
+	vmBefore := s.app.UpgradeKeeper.GetModuleVersionMap(s.ctx)
+	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		// simulate upgrading the bank module
+		vm["bank"] = vm["bank"] + 1
+		return vm, nil
+	})
+	dummyPlan := types.Plan{
+		Name:   "dummy",
+		Info:   "some text here",
+		Height: 123450000,
+	}
+
+	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
+	vm := s.app.UpgradeKeeper.GetModuleVersionMap(s.ctx)
+	s.Require().Equal(vmBefore["bank"]+1, vm["bank"])
+	s.Require().Greater(vm["bank"], vmBefore["bank"])
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/upgrade/spec/01_concepts.md

@@ -47,7 +47,7 @@ and not defined on a per-module basis. Registering a `Handler` is done via
 `Keeper#SetUpgradeHandler` in the application.
 
 ```go
-type UpgradeHandler func(Context, Plan)
+type UpgradeHandler func(Context, Plan, VersionMap) (VersionMap, error)
 ```
 
 During each `EndBlock` execution, the `x/upgrade` module checks if there exists a

x/upgrade/spec/02_state.md

@@ -5,7 +5,15 @@ order: 2
 # State
 
 The internal state of the `x/upgrade` module is relatively minimal and simple. The
-state only contains the currently active upgrade `Plan` (if one exists) by key
-`0x0` and if a `Plan` is marked as "done" by key `0x1`.
+state contains the currently active upgrade `Plan` (if one exists) by key
+`0x0` and if a `Plan` is marked as "done" by key `0x1`. Additionally, the state
+contains the consensus versions of all app modules in the application. The versions 
+are stored as big endian `uint64`, and can be accessed with prefix `0x2` appended 
+by the corresponding module name of type `string`. 
+
+- Plan: `0x0 -> Plan`
+- Done: `0x1 | byte(plan name)  -> BigEndian(Block Height)`
+- ConsensusVersion: `0x2 | byte(module name)  -> BigEndian(Module Consensus Version)`
+
 
 The `x/upgrade` module contains no genesis state.