refactor: simplify the upgrade module

simapp/app.go

@@ -13,6 +13,7 @@ import (
 	abci "github.com/tendermint/tendermint/abci/types"
 	"github.com/tendermint/tendermint/libs/log"
 	tmos "github.com/tendermint/tendermint/libs/os"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 	dbm "github.com/tendermint/tm-db"
 
 	"github.com/cosmos/cosmos-sdk/baseapp"
@@ -29,7 +30,6 @@ import (
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/testutil/testdata"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/version"
 	"github.com/cosmos/cosmos-sdk/x/auth"
@@ -298,18 +298,23 @@ func NewSimApp(
 	*/
 	app.GroupKeeper = groupkeeper.NewKeeper(keys[group.StoreKey], appCodec, app.MsgServiceRouter(), app.AccountKeeper, groupConfig)
 
-	app.UpgradeKeeper = upgradekeeper.NewKeeper(skipUpgradeHeights, keys[upgradetypes.StoreKey], appCodec, homePath)
-	defer func() {
-		// Register the upgrade plan
-		upgradeHandler := map[string]upgradetypes.UpgradeHandler{
-			// Add specific actions when the upgrade happen
-		}
-		state := app.GetState(0 /*runTxModeCheck*/)
-		err := app.UpgradeKeeper.RegisterUpgradePlan(state.Context(), bApp.AppConfig().Upgrade, upgradeHandler)
-		if err != nil {
-			panic(errors.Wrap(err, "failed to regiter upgrade plan"))
-		}
-	}()
+	// Register the upgrade keeper
+	upgradeHandler := map[string]upgradetypes.UpgradeHandler{
+		// Add specific actions when the upgrade happen
+		// ex.
+		// "BEP111": func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		// 	app.Logger().Info("upgrade to ", plan.Name)
+		// 	return fromVM, nil
+		// },
+	}
+
+	ms := app.CommitMultiStore()
+	ctx := sdk.NewContext(ms, tmproto.Header{ChainID: app.ChainID(), Height: app.LastBlockHeight()}, true, app.Logger())
+	upgradeKeeperOpts := []upgradekeeper.KeeperOption{
+		upgradekeeper.RegisterUpgradePlan(ctx, bApp.AppConfig().Upgrade),
+		upgradekeeper.RegisterUpgradeHandler(upgradeHandler),
+	}
+	app.UpgradeKeeper = upgradekeeper.NewKeeper(skipUpgradeHeights, keys[upgradetypes.StoreKey], appCodec, homePath, upgradeKeeperOpts...)
 
 	// Register the proposal types
 	// Deprecated: Avoid adding new handlers, instead use the new proposal flow

x/upgrade/abci.go

@@ -42,20 +42,6 @@ func BeginBlocker(k keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
 			return
 		}
 
-		// Prepare shutdown if we don't have an upgrade handler for this upgrade name (meaning this software is out of date)
-		if !k.HasHandler(plan.Name) {
-			// Write the upgrade info to disk. The UpgradeStoreLoader uses this info to perform or skip
-			// store migrations.
-			err := k.DumpUpgradeInfoToDisk(ctx.BlockHeight(), plan)
-			if err != nil {
-				panic(fmt.Errorf("unable to write upgrade info to filesystem: %s", err.Error()))
-			}
-
-			upgradeMsg := BuildUpgradeNeededMsg(plan)
-			logger.Error(upgradeMsg)
-			return
-		}
-
 		// We have an upgrade handler for this upgrade name, so apply the upgrade
 		ctx.Logger().Info(fmt.Sprintf("applying upgrade \"%s\" at %s", plan.Name, plan.DueAt()))
 		ctx = ctx.WithBlockGasMeter(sdk.NewInfiniteGasMeter())

x/upgrade/abci_test.go

@@ -87,7 +87,7 @@ func VerifyDoUpgrade(t *testing.T) {
 	newCtx := s.ctx.WithBlockHeight(s.ctx.BlockHeight() + 1).WithBlockTime(time.Now())
 
 	req := abci.RequestBeginBlock{Header: newCtx.BlockHeader()}
-	require.Panics(t, func() {
+	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
 
@@ -347,7 +347,7 @@ func TestUpgradeWithoutSkip(t *testing.T) {
 	err := s.keeper.ScheduleUpgrade(s.ctx, types.Plan{Name: "test", Height: s.ctx.BlockHeight() + 1})
 	require.NoError(t, err)
 	t.Log("Verify if upgrade happens without skip upgrade")
-	require.Panics(t, func() {
+	require.NotPanics(t, func() {
 		s.module.BeginBlock(newCtx, req)
 	})
 

x/upgrade/keeper/keeper.go

@@ -27,21 +27,27 @@ type Keeper struct {
 	storeKey           storetypes.StoreKey             // key to access x/upgrade store
 	cdc                codec.BinaryCodec               // App-wide binary codec
 	upgradeHandlers    map[string]types.UpgradeHandler // map of plan name to upgrade handler
+	upgradeConfig      types.UpgradeConfig
 }
 
 // NewKeeper constructs an upgrade Keeper which requires the following arguments:
 // skipUpgradeHeights - map of heights to skip an upgrade
 // storeKey - a store key with which to access upgrade's store
 // cdc - the app-wide binary codec
 // homePath - root directory of the application's config
-func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string) Keeper {
-	return Keeper{
+func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string, opts ...KeeperOption) Keeper {
+	keeper := Keeper{
 		homePath:           homePath,
 		skipUpgradeHeights: skipUpgradeHeights,
 		storeKey:           storeKey,
 		cdc:                cdc,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
 	}
+
+	for _, opt := range opts {
+		opt(&keeper)
+	}
+	return keeper
 }
 
 // SetUpgradeHandler sets an UpgradeHandler for the upgrade specified by name. This handler will be called when the upgrade
@@ -293,13 +299,15 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger {
 // GetUpgradePlan returns the currently scheduled Plan if any, setting havePlan to true if there is a scheduled
 // upgrade or false if there is none
 func (k Keeper) GetUpgradePlan(ctx sdk.Context) (plan types.Plan, havePlan bool) {
-	store := ctx.KVStore(k.storeKey)
-	bz := store.Get(types.PlanKey())
-	if bz == nil {
-		return plan, false
+	plan, ok := k.upgradeConfig[ctx.BlockHeight()]
+	if !ok {
+		return types.Plan{}, false
+	}
+
+	if k.IsUpgrade(ctx, plan.Name) {
+		return types.Plan{}, false
 	}
 
-	k.cdc.MustUnmarshal(bz, &plan)
 	return plan, true
 }
 
@@ -318,18 +326,16 @@ func (k Keeper) HasHandler(name string) bool {
 // ApplyUpgrade will execute the handler associated with the Plan and mark the plan as done.
 func (k Keeper) ApplyUpgrade(ctx sdk.Context, plan types.Plan) {
 	handler := k.upgradeHandlers[plan.Name]
-	if handler == nil {
-		return
-	}
 
-	updatedVM, err := handler(ctx, plan, k.GetModuleVersionMap(ctx))
-	if err != nil {
-		ctx.Logger().Error("failed to upgrade ["+plan.Name+"]", "err", err)
-		return
+	if handler != nil {
+		updatedVM, err := handler(ctx, plan, k.GetModuleVersionMap(ctx))
+		if err != nil {
+			ctx.Logger().Error("failed to upgrade ["+plan.Name+"]", "err", err)
+			return
+		}
+		k.SetModuleVersionMap(ctx, updatedVM)
 	}
 
-	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.
 	k.ClearIBCState(ctx, plan.Height)

x/upgrade/keeper/keeper_option.go

@@ -0,0 +1,46 @@
+package keeper
+
+import (
+	serverconfig "github.com/cosmos/cosmos-sdk/server/config"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
+)
+
+func convertUpgradeConfig(ctx sdk.Context, plans []serverconfig.UpgradeConfig) types.UpgradeConfig {
+	upgradeConfig := types.NewUpgradeConfig()
+	if ctx.ChainID() == types.MainnetChainID {
+		upgradeConfig = types.MainnetConfig
+	}
+
+	// override by app config
+	for _, plan := range plans {
+		upgradeConfig.SetPlan(types.Plan{
+			Name:   plan.Name,
+			Height: plan.Height,
+			Info:   plan.Info,
+		})
+	}
+
+	return upgradeConfig
+}
+
+// Option function for Keeper
+type KeeperOption func(k *Keeper)
+
+// RegisterUpgradePlan returns a KeeperOption to set the upgrade plan into the upgrade keeper
+func RegisterUpgradePlan(ctx sdk.Context,
+	plans []serverconfig.UpgradeConfig,
+) KeeperOption {
+	return func(k *Keeper) {
+		k.upgradeConfig = convertUpgradeConfig(ctx, plans)
+	}
+}
+
+// RegisterUpgradeHandler returns a KeeperOption to set the upgrade handler into the upgrade keeper
+func RegisterUpgradeHandler(handlers map[string]types.UpgradeHandler) KeeperOption {
+	return func(k *Keeper) {
+		for name, handler := range handlers {
+			k.SetUpgradeHandler(name, handler)
+		}
+	}
+}