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

CHANGELOG.md

@@ -58,6 +58,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` requires new argument `VersionManager` which retreives module consensus versions from upgrade's new store.
 
 ### State Machine Breaking
 

simapp/app.go

@@ -320,6 +320,8 @@ func NewSimApp(
 		params.NewAppModule(app.ParamsKeeper),
 		authz.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
 	)
+	// give upgrade keeper the module manager
+	app.UpgradeKeeper.SetVersionManager(app.mm)
 
 	// During begin block slashing happens after distr.BeginBlocker so that
 	// there is nothing left over in the validator fee pool, so as to keep the
@@ -551,19 +553,18 @@ 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, versionMap module.VersionMap) {
+//       return app.RunMigrations(ctx, versionMap)
 //   })
-func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.MigrationMap) error {
+func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.VersionMap) error {
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 
+// Returns a map (VersionMap) of key module name and value module consensus version
+func (app *SimApp) GetConsensusVersions() module.VersionMap {
+	return app.mm.GetConsensusVersions()
+}
+
 // RegisterSwaggerAPI registers swagger route with API Server
 func RegisterSwaggerAPI(ctx client.Context, rtr *mux.Router) {
 	statikFS, err := fs.New()

simapp/app_test.go

@@ -159,7 +159,7 @@ func TestRunMigrations(t *testing.T) {
 			// their latest ConsensusVersion.
 			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(),

types/module/module.go

@@ -333,12 +333,16 @@ 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
+
+type VersionManager interface {
+	GetConsensusVersions() VersionMap
+}
 
 // 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, migrateFromVersions VersionMap) error {
 	c, ok := cfg.(configurator)
 	if !ok {
 		return sdkerrors.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", configurator{}, cfg)
@@ -395,3 +399,15 @@ func (m *Manager) EndBlock(ctx sdk.Context, req abci.RequestEndBlock) abci.Respo
 		Events:           ctx.EventManager().ABCIEvents(),
 	}
 }
+
+// GetConsensusVersions gets consensus version from all modules
+func (m *Manager) GetConsensusVersions() 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

@@ -128,7 +128,7 @@ 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, versionMap module.VersionMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -144,7 +144,7 @@ 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, versionMap module.VersionMap) error { return nil })
 	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
@@ -156,7 +156,7 @@ 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, versionMap module.VersionMap) error { called++; return 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,7 @@ 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(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return 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"
 )
 
@@ -28,6 +29,7 @@ type Keeper struct {
 	storeKey           sdk.StoreKey
 	cdc                codec.BinaryMarshaler
 	upgradeHandlers    map[string]types.UpgradeHandler
+	versionManager     module.VersionManager
 }
 
 // NewKeeper constructs an upgrade Keeper
@@ -48,6 +50,45 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
+func (k *Keeper) SetVersionManager(vm module.VersionManager) {
+	k.versionManager = vm
+}
+
+// SetCurrentConsensusVersions saves the consensus versions retrieved from module.Manager
+func (k Keeper) SetCurrentConsensusVersions(ctx sdk.Context) error {
+	if k.versionManager == nil {
+		return sdkerrors.Wrap(sdkerrors.ErrLogic, "upgrade keeper VersionManager was nil, please call SetVersionManager on the keeper first")
+	}
+	modules := k.versionManager.GetConsensusVersions()
+	store := ctx.KVStore(k.storeKey)
+	versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
+	for modName, ver := range modules {
+		nameBytes := []byte(modName)
+		verBytes := make([]byte, 8)
+		binary.LittleEndian.PutUint64(verBytes, ver)
+		versionStore.Set(nameBytes, verBytes)
+	}
+	return nil
+}
+
+// GetConsensusVersions gets a VersionMap from state
+func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.VersionMap {
+	store := ctx.KVStore(k.storeKey)
+	it := sdk.KVStorePrefixIterator(store, []byte{types.VersionMapByte})
+
+	vermap := 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.LittleEndian.Uint64(it.Value())
+		vermap[name] = moduleVersion
+	}
+
+	return vermap
+}
+
 // 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)
@@ -191,7 +232,15 @@ func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 		panic("ApplyUpgrade should never be called without first checking HasHandler")
 	}
 
-	handler(ctx, plan)
+	err := handler(ctx, plan, k.GetConsensusVersions(ctx))
+	if err != nil {
+		panic(err)
+	}
+
+	err = k.SetCurrentConsensusVersions(ctx)
+	if err != nil {
+		panic(err)
+	}
 
 	// 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,8 @@ 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/bank"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/keeper"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
@@ -135,7 +137,7 @@ 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(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
 				s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, types.Plan{
 					Name:   "all-good",
 					Info:   "some text here",
@@ -152,7 +154,7 @@ func (s *KeeperTestSuite) TestScheduleUpgrade() {
 		s.Run(tc.name, func() {
 			// reset suite
 			s.SetupTest()
-
+			s.app.UpgradeKeeper.SetVersionManager(s.app)
 			// setup test case
 			tc.setup()
 
@@ -211,6 +213,35 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 
 }
 
+// Mock version manager for TestMigrations
+type mockVersionManager struct{}
+
+func (m mockVersionManager) GetConsensusVersions() module.VersionMap {
+	vermap := make(module.VersionMap)
+	vermap["bank"] = 1
+	return vermap
+}
+
+// Tests that the underlying state of x/upgrade is set correctly after
+// an upgrade.
+func (s *KeeperTestSuite) TestMigrations() {
+	mockVM := mockVersionManager{}
+	s.app.UpgradeKeeper.SetVersionManager(mockVM)
+	err := s.app.UpgradeKeeper.SetCurrentConsensusVersions(s.ctx)
+	s.Require().Nil(err)
+	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.VersionMap) error { return nil })
+	dummyPlan := types.Plan{
+		Name: "dummy",
+		Info: "some text here",
+		Time: s.ctx.BlockTime().Add(time.Hour),
+	}
+
+	s.app.UpgradeKeeper.SetVersionManager(s.app)
+	s.app.UpgradeKeeper.ApplyUpgrade(s.ctx, dummyPlan)
+	vermap := s.app.UpgradeKeeper.GetConsensusVersions(s.ctx)
+	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), vermap["bank"])
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/upgrade/spec/01_concepts.md

@@ -48,7 +48,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) 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 little endian `uint64`, and can be accessed with prefix `0x2` appended 
+by the corresponding module name of type `string`. 
 
-The `x/upgrade` module contains no genesis state.
+- Plan: `0x0 -> Plan`
+- Done: `0x1 | byte(plan name)  -> BigEndian(Block Height)`
+- ConsensusVersion: `0x2 | byte(module name)  -> LittleEndian(Module Consensus Version)`
+
+
+The `x/upgrade` module contains no genesis state.

x/upgrade/types/handler.go

@@ -2,7 +2,8 @@ package types
 
 import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
 )
 
 // UpgradeHandler specifies the type of function that is called when an upgrade is applied
-type UpgradeHandler func(ctx sdk.Context, plan Plan)
+type UpgradeHandler func(ctx sdk.Context, plan Plan, versionMap module.VersionMap) error

x/upgrade/types/keys.go

@@ -22,6 +22,9 @@ const (
 	// DoneByte is a prefix for to look up completed upgrade plan by name
 	DoneByte = 0x1
 
+	// VersionMapByte is a prefix to look up module names (key) and versions (value)
+	VersionMapByte = 0x2
+
 	// KeyUpgradedIBCState is the key under which upgraded ibc state is stored in the upgrade store
 	KeyUpgradedIBCState = "upgradedIBCState"