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

* -added consensus version tracking to x/upgrade

* -added interface to module manager -added e2e test for migrations using consensus version store in x/upgrade -cleaned up x/upgrade Keeper -handler in apply upgrade now handles errors and setting consensus versions -cleaned up migration map keys -removed init chainer method -simapp now implements GetConsensusVersions to assist with testing

* Changed MigrationMap identifier to VersionMap
removed module_test

* updated docs

* forgot this

* added line to changelog for this PR

* Change set consensus version function to match adr 041 spec

* add documentation

* remove newline from changelog

unnecessary newline removed

* updated example in simapp for RunMigrations, SetCurrentConsensusVersions now returns an error

* switch TestMigrations to use Require instead of t.Fatal

* Update CHANGELOG.md

Co-authored-by: Aaron Craelius <aaron@regen.network>

* docs for SetVersionManager

* -init genesis method added -removed panics/fails from setting consensus versions

* update identifiers to be more go-like

* update docs and UpgradeHandler fnc sig

* Upgrade Keeper now takes a VersionMap instead of a VersionManager interface

* upgrade keeper transition to Version Map

* cleanup, added versionmap return to RunMigrations

* quick fix

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

Co-authored-by: Amaury <1293565+amaurym@users.noreply.github.com>

* remove support for versionmap field on upgrade keeper

* cleanup

* rename get/set version map keeper functions

* update adr doc to match name changes

* remove redudant line

Co-authored-by: technicallyty <48813565+tytech3@users.noreply.github.com>
Co-authored-by: Aaron Craelius <aaron@regen.network>
Co-authored-by: Amaury <1293565+amaurym@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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,28 @@ 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"])
+}
+
 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.

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, vm module.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"