refactor(x/upgrade)!: use KVStoreService and context.Context

CHANGELOG.md

@@ -130,6 +130,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (module) [#16227](https://github.com/cosmos/cosmos-sdk/issues/16227) `manager.RunMigrations()` now take a `context.Context` instead of a `sdk.Context`.
 * (x/crisis) [#16216](https://github.com/cosmos/cosmos-sdk/issues/16216) `NewKeeper` now takes a `KVStoreService` instead of a `StoreKey`, methods in the `Keeper` now take a `context.Context` instead of a `sdk.Context` and return an `error` instead of panicking.
 * (x/gov) [#15988](https://github.com/cosmos/cosmos-sdk/issues/15988) `NewKeeper` now takes a `KVStoreService` instead of a `StoreKey`, methods in the `Keeper` now take a `context.Context` instead of a `sdk.Context` and return an `error` (instead of panicking or returning a `found bool`). Iterators callback functions now return an error instead of a `bool`.
 * (x/auth) [#15985](https://github.com/cosmos/cosmos-sdk/pull/15985) The `AccountKeeper` does not expose the `QueryServer` and `MsgServer` APIs anymore.

UPGRADING.md

@@ -83,6 +83,7 @@ The following modules `NewKeeper` function now take a `KVStoreService` instead o
 * `x/feegrant`
 * `x/gov`
 * `x/nft`
+* `x/upgrade`
 
 User manually wiring their chain need to use the `runtime.NewKVStoreService` method to create a `KVStoreService` from a `StoreKey`:
 
@@ -103,6 +104,7 @@ The following modules' `Keeper` methods now take in a `context.Context` instead
 * `x/distribution`
 * `x/evidence`
 * `x/gov`
+* `x/upgrade`
 
 **Users using depinject do not need any changes, this is automatically done for them.**
 

simapp/app.go

@@ -319,7 +319,7 @@ func NewSimApp(
 	}
 	homePath := cast.ToString(appOpts.Get(flags.FlagHome))
 	// set the governance module account as the authority for conducting upgrades
-	app.UpgradeKeeper = upgradekeeper.NewKeeper(skipUpgradeHeights, keys[upgradetypes.StoreKey], appCodec, homePath, app.BaseApp, authtypes.NewModuleAddress(govtypes.ModuleName).String())
+	app.UpgradeKeeper = upgradekeeper.NewKeeper(skipUpgradeHeights, runtime.NewKVStoreService(keys[upgradetypes.StoreKey]), appCodec, homePath, app.BaseApp, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
 	// Register the proposal types
 	// Deprecated: Avoid adding new handlers, instead use the new proposal flow

simapp/app_test.go

@@ -263,7 +263,8 @@ func TestUpgradeStateOnGenesis(t *testing.T) {
 
 	// make sure the upgrade keeper has version map in state
 	ctx := app.NewContext(false, cmtproto.Header{})
-	vm := app.UpgradeKeeper.GetModuleVersionMap(ctx)
+	vm, err := app.UpgradeKeeper.GetModuleVersionMap(ctx)
+	require.NoError(t, err)
 	for v, i := range app.ModuleManager.Modules {
 		if i, ok := i.(module.HasConsensusVersion); ok {
 			require.Equal(t, vm[v], i.ConsensusVersion())

simapp/upgrades.go

@@ -1,10 +1,11 @@
 package simapp
 
 import (
+	"context"
+
 	storetypes "cosmossdk.io/store/types"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
 
-	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 )
 
@@ -19,7 +20,7 @@ const UpgradeName = "v047-to-v048"
 func (app SimApp) RegisterUpgradeHandlers() {
 	app.UpgradeKeeper.SetUpgradeHandler(
 		UpgradeName,
-		func(ctx sdk.Context, _ upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		func(ctx context.Context, _ upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 			return app.ModuleManager.RunMigrations(ctx, app.Configurator(), fromVM)
 		},
 	)

tests/e2e/upgrade/suite.go

@@ -76,10 +76,13 @@ func (s *E2ETestSuite) TestModuleVersionsCLI() {
 	// avoid printing as yaml from CLI command
 	clientCtx.OutputFormat = "JSON"
 
-	vm := s.upgradeKeeper.GetModuleVersionMap(s.ctx)
-	mv := s.upgradeKeeper.GetModuleVersions(s.ctx)
+	vm, err := s.upgradeKeeper.GetModuleVersionMap(s.ctx)
+	s.Require().NoError(err)
 	s.Require().NotEmpty(vm)
 
+	mv, err := s.upgradeKeeper.GetModuleVersions(s.ctx)
+	s.Require().NoError(err)
+
 	for _, tc := range testCases {
 		s.Run(fmt.Sprintf("Case %s", tc.msg), func() {
 			expect := mv

types/module/module.go

@@ -609,7 +609,7 @@ type VersionMap map[string]uint64
 // Example:
 //
 //	cfg := module.NewConfigurator(...)
-//	app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+//	app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 //	    return app.mm.RunMigrations(ctx, cfg, fromVM)
 //	})
 //
@@ -636,7 +636,7 @@ type VersionMap map[string]uint64
 // Example:
 //
 //	cfg := module.NewConfigurator(...)
-//	app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx sdk.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+//	app.UpgradeKeeper.SetUpgradeHandler("my-plan", func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 //	    // Assume "foo" is a new module.
 //	    // `fromVM` is fetched from existing x/upgrade store. Since foo didn't exist
 //	    // before this upgrade, `v, exists := fromVM["foo"]; exists == false`, and RunMigration will by default
@@ -649,7 +649,7 @@ type VersionMap map[string]uint64
 //	})
 //
 // Please also refer to docs/core/upgrade.md for more information.
-func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM VersionMap) (VersionMap, error) {
+func (m Manager) RunMigrations(ctx context.Context, cfg Configurator, fromVM VersionMap) (VersionMap, error) {
 	c, ok := cfg.(*configurator)
 	if !ok {
 		return nil, errorsmod.Wrapf(sdkerrors.ErrInvalidType, "expected %T, got %T", &configurator{}, cfg)
@@ -659,6 +659,7 @@ func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM Version
 		modules = DefaultMigrationsOrder(m.ModuleNames())
 	}
 
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
 	updatedVM := VersionMap{}
 	for _, moduleName := range modules {
 		module := m.Modules[moduleName]
@@ -677,14 +678,14 @@ func (m Manager) RunMigrations(ctx sdk.Context, cfg Configurator, fromVM Version
 		// 2. An existing chain is upgrading from version < 0.43 to v0.43+ for the first time.
 		// In this case, all modules have yet to be added to x/upgrade's VersionMap store.
 		if exists {
-			err := c.runModuleMigrations(ctx, moduleName, fromVersion, toVersion)
+			err := c.runModuleMigrations(sdkCtx, moduleName, fromVersion, toVersion)
 			if err != nil {
 				return nil, err
 			}
 		} else {
-			ctx.Logger().Info(fmt.Sprintf("adding a new module: %s", moduleName))
+			sdkCtx.Logger().Info(fmt.Sprintf("adding a new module: %s", moduleName))
 			if module, ok := m.Modules[moduleName].(HasGenesis); ok {
-				moduleValUpdates := module.InitGenesis(ctx, c.cdc, module.DefaultGenesis(c.cdc))
+				moduleValUpdates := module.InitGenesis(sdkCtx, c.cdc, module.DefaultGenesis(c.cdc))
 				// The module manager assumes only one module will update the
 				// validator set, and it can't be a new module.
 				if len(moduleValUpdates) > 0 {

x/upgrade/CHANGELOG.md

@@ -29,3 +29,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 * [#14880](https://github.com/cosmos/cosmos-sdk/pull/14880) Switch from using gov v1beta1 to gov v1 in upgrade CLIs.
 * [#14764](https://github.com/cosmos/cosmos-sdk/pull/14764) The `x/upgrade` module is extracted to have a separate go.mod file which allows it be a standalone module.
+
+### API Breaking
+
+* (x/upgrade) [#16227](https://github.com/cosmos/cosmos-sdk/issues/16227) `NewKeeper` now takes a `KVStoreService` instead of a `StoreKey`, methods in the `Keeper` now take a `context.Context` instead of a `sdk.Context` and return an `error`. `UpgradeHandler` now receives a `context.Context`. `GetUpgradedClient`, `GetUpgradedConsensusState`, `GetUpgradePlan` now return a specific error for "not found".

x/upgrade/abci.go

@@ -1,6 +1,8 @@
 package upgrade
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"time"
 
@@ -20,10 +22,16 @@ import (
 // The purpose is to ensure the binary is switched EXACTLY at the desired block, and to allow
 // a migration to be executed if needed upon this switch (migration defined in the new binary)
 // skipUpgradeHeightArray is a set of block heights for which the upgrade must be skipped
-func BeginBlocker(k *keeper.Keeper, ctx sdk.Context) {
+func BeginBlocker(ctx context.Context, k *keeper.Keeper) error {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
 
-	plan, found := k.GetUpgradePlan(ctx)
+	sdkCtx := sdk.UnwrapSDKContext(ctx)
+	blockHeight := sdkCtx.BlockHeight()
+	plan, err := k.GetUpgradePlan(ctx)
+	if err != nil && !errors.Is(err, types.ErrNoUpgradePlanFound) {
+		return err
+	}
+	found := err == nil
 
 	if !k.DowngradeVerified() {
 		k.SetDowngradeVerified(true)
@@ -32,66 +40,75 @@ func BeginBlocker(k *keeper.Keeper, ctx sdk.Context) {
 		// 1. If there is no scheduled upgrade.
 		// 2. If the plan is not ready.
 		// 3. If the plan is ready and skip upgrade height is set for current height.
-		if !found || !plan.ShouldExecute(ctx) || (plan.ShouldExecute(ctx) && k.IsSkipHeight(ctx.BlockHeight())) {
-			lastAppliedPlan, _ := k.GetLastCompletedUpgrade(ctx)
+		if !found || !plan.ShouldExecute(blockHeight) || (plan.ShouldExecute(blockHeight) && k.IsSkipHeight(blockHeight)) {
+			lastAppliedPlan, _, err := k.GetLastCompletedUpgrade(ctx)
+			if err != nil {
+				return err
+			}
+
 			if lastAppliedPlan != "" && !k.HasHandler(lastAppliedPlan) {
 				var appVersion uint64
 
-				cp := ctx.ConsensusParams()
+				cp := sdkCtx.ConsensusParams()
 				if cp.Version != nil {
 					appVersion = cp.Version.App
 				}
 
-				panic(fmt.Sprintf("Wrong app version %d, upgrade handler is missing for %s upgrade plan", appVersion, lastAppliedPlan))
+				return fmt.Errorf("wrong app version %d, upgrade handler is missing for %s upgrade plan", appVersion, lastAppliedPlan)
 			}
 		}
 	}
 
 	if !found {
-		return
+		return nil
 	}
-	logger := ctx.Logger()
+
+	logger := k.Logger(ctx)
 
 	// To make sure clear upgrade is executed at the same block
-	if plan.ShouldExecute(ctx) {
+	if plan.ShouldExecute(blockHeight) {
 		// If skip upgrade has been set for current height, we clear the upgrade plan
-		if k.IsSkipHeight(ctx.BlockHeight()) {
+		if k.IsSkipHeight(blockHeight) {
 			skipUpgradeMsg := fmt.Sprintf("UPGRADE \"%s\" SKIPPED at %d: %s", plan.Name, plan.Height, plan.Info)
 			logger.Info(skipUpgradeMsg)
 
 			// Clear the upgrade plan at current height
-			k.ClearUpgradePlan(ctx)
-			return
+			return k.ClearUpgradePlan(ctx)
 		}
 
 		// 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)
+			err := k.DumpUpgradeInfoToDisk(blockHeight, plan)
 			if err != nil {
-				panic(fmt.Errorf("unable to write upgrade info to filesystem: %s", err.Error()))
+				return fmt.Errorf("unable to write upgrade info to filesystem: %s", err.Error())
 			}
 
 			upgradeMsg := BuildUpgradeNeededMsg(plan)
 			logger.Error(upgradeMsg)
-			panic(upgradeMsg)
+
+			// Returning an error will end up in a panic
+			return fmt.Errorf(upgradeMsg)
 		}
 
 		// 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(storetypes.NewInfiniteGasMeter())
-		k.ApplyUpgrade(ctx, plan)
-		return
+		logger.Info(fmt.Sprintf("applying upgrade \"%s\" at %s", plan.Name, plan.DueAt()))
+		sdkCtx = sdkCtx.WithBlockGasMeter(storetypes.NewInfiniteGasMeter())
+		return k.ApplyUpgrade(sdkCtx, plan)
 	}
 
 	// if we have a pending upgrade, but it is not yet time, make sure we did not
 	// set the handler already
 	if k.HasHandler(plan.Name) {
 		downgradeMsg := fmt.Sprintf("BINARY UPDATED BEFORE TRIGGER! UPGRADE \"%s\" - in binary but not executed on chain. Downgrade your binary", plan.Name)
-		ctx.Logger().Error(downgradeMsg)
-		panic(downgradeMsg)
+		logger.Error(downgradeMsg)
+
+		// Returning an error will end up in a panic
+		return fmt.Errorf(downgradeMsg)
 	}
+
+	return nil
 }
 
 // BuildUpgradeNeededMsg prints the message that notifies that an upgrade is needed.