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

simapp/app.go

@@ -374,6 +374,10 @@ func NewSimApp(
 		authz.NewAppModule(appCodec, app.AuthzKeeper, app.AccountKeeper, app.BankKeeper, app.interfaceRegistry),
 		transferModule,
 	)
+	// give upgrade keeper the module manager
+	app.UpgradeKeeper.SetVersionManager(app.mm)
+	// pass the updated keeper to the module manager
+	app.mm.Modules[upgradetypes.ModuleName] = upgrade.NewAppModule(app.UpgradeKeeper)
 
 	// 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
@@ -391,7 +395,7 @@ func NewSimApp(
 	// so that other modules that want to create or claim capabilities afterwards in InitChain
 	// can do so safely.
 	app.mm.SetOrderInitGenesis(
-		capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
+		upgradetypes.ModuleName, capabilitytypes.ModuleName, authtypes.ModuleName, banktypes.ModuleName, distrtypes.ModuleName, stakingtypes.ModuleName,
 		slashingtypes.ModuleName, govtypes.ModuleName, minttypes.ModuleName, crisistypes.ModuleName,
 		ibchost.ModuleName, genutiltypes.ModuleName, evidencetypes.ModuleName, authztypes.ModuleName, ibctransfertypes.ModuleName,
 		feegranttypes.ModuleName,
@@ -627,6 +631,11 @@ func (app *SimApp) RunMigrations(ctx sdk.Context, migrateFromVersions module.Mig
 	return app.mm.RunMigrations(ctx, app.configurator, migrateFromVersions)
 }
 
+// Returns a map (MigrationMap) of key module name and value module consensus version
+func (app *SimApp) GetConsensusVersions() module.MigrationMap {
+	return app.mm.GetConsensusVersions()
+}
+
 // RegisterSwaggerAPI registers swagger route with API Server
 func RegisterSwaggerAPI(ctx client.Context, rtr *mux.Router) {
 	statikFS, err := fs.New()

types/module/module.go

@@ -337,6 +337,10 @@ type MigrationHandler func(sdk.Context) error
 // version from which we should perform the migration for each module.
 type MigrationMap map[string]uint64
 
+type VersionManager interface {
+	GetConsensusVersions() MigrationMap
+}
+
 // RunMigrations performs in-place store migrations for all modules.
 func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, migrateFromVersions MigrationMap) error {
 	c, ok := cfg.(configurator)
@@ -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() MigrationMap {
+	migmap := make(MigrationMap)
+	for _, v := range m.Modules {
+		version := v.ConsensusVersion()
+		name := v.Name()
+		migmap[name] = version
+	}
+
+	return migmap
+}

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, migrationMap module.MigrationMap) 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, migrationMap module.MigrationMap) 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, migrationMap module.MigrationMap) 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.MigrationMap) 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,40 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 	k.upgradeHandlers[name] = upgradeHandler
 }
 
+func (k *Keeper) SetVersionManager(vm module.VersionManager) {
+	k.versionManager = vm
+}
+
+// SetConsensusVersions saves the consensus versions retrieved from module.Manager
+func (k Keeper) SetConsensusVersions(ctx sdk.Context) {
+	modules := k.versionManager.GetConsensusVersions()
+	store := ctx.KVStore(k.storeKey)
+	migrationStore := prefix.NewStore(store, []byte{types.MigrationMapByte})
+	for modName, ver := range modules {
+		nameBytes := []byte(modName)
+		verBytes := make([]byte, 8)
+		binary.LittleEndian.PutUint64(verBytes, ver)
+		migrationStore.Set(nameBytes, verBytes)
+	}
+}
+
+// GetConsensusVersions gets a MigrationMap from state
+func (k Keeper) GetConsensusVersions(ctx sdk.Context) module.MigrationMap {
+	store := ctx.KVStore(k.storeKey)
+	it := sdk.KVStorePrefixIterator(store, []byte{types.MigrationMapByte})
+
+	migmap := make(module.MigrationMap)
+	defer it.Close()
+	for ; it.Valid(); it.Next() {
+		moduleBytes := it.Key()
+		name := string(moduleBytes[1:])
+		moduleVersion := binary.LittleEndian.Uint64(it.Value())
+		migmap[name] = moduleVersion
+	}
+
+	return migmap
+}
+
 // 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 +227,12 @@ 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)
+	}
+
+	k.SetConsensusVersions(ctx)
 
 	// 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.MigrationMap) 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,34 @@ func (s *KeeperTestSuite) TestSetUpgradedClient() {
 
 }
 
+// Mock version manager for TestMigrations
+type MockVersionManager struct{}
+
+func (m MockVersionManager) GetConsensusVersions() module.MigrationMap {
+	migmap := make(module.MigrationMap)
+	migmap["bank"] = 1
+	return migmap
+}
+
+// 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)
+	s.app.UpgradeKeeper.SetConsensusVersions(s.ctx)
+	s.app.UpgradeKeeper.SetUpgradeHandler("dummy", func(_ sdk.Context, _ types.Plan, _ module.MigrationMap) 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)
+	migmap := s.app.UpgradeKeeper.GetConsensusVersions(s.ctx)
+	s.Require().Equal(bank.AppModule{}.ConsensusVersion(), migmap["bank"])
+}
+
 func TestKeeperTestSuite(t *testing.T) {
 	suite.Run(t, new(KeeperTestSuite))
 }

x/upgrade/module.go

@@ -101,7 +101,7 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 }
 
 // InitGenesis is ignored, no sense in serializing future upgrades
-func (am AppModule) InitGenesis(_ sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
+func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
 	return []abci.ValidatorUpdate{}
 }
 

x/upgrade/module_test.go

@@ -0,0 +1,24 @@
+package upgrade_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	abcitypes "github.com/tendermint/tendermint/abci/types"
+
+	"github.com/cosmos/cosmos-sdk/simapp"
+)
+
+func TestItGetsModuleVersions(t *testing.T) {
+	app := simapp.Setup(false)
+
+	app.InitChain(
+		abcitypes.RequestInitChain{
+			AppStateBytes: []byte("{}"),
+			ChainId:       "test-chain-id",
+		},
+	)
+
+	migmap := app.GetConsensusVersions()
+	require.NotEmpty(t, migmap)
+}

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, migrationMap module.MigrationMap) 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
 
+	// MigrationMapByte is a prefix to look up module names (key) and versions (value)
+	MigrationMapByte = 0x2
+
 	// KeyUpgradedIBCState is the key under which upgraded ibc state is stored in the upgrade store
 	KeyUpgradedIBCState = "upgradedIBCState"