refactor(x/upgrade): migrate upgrade to env var (#19443)

simapp/app.go

@@ -374,7 +374,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, runtime.NewKVStoreService(keys[upgradetypes.StoreKey]), appCodec, homePath, app.BaseApp, authtypes.NewModuleAddress(govtypes.ModuleName).String())
+	app.UpgradeKeeper = upgradekeeper.NewKeeper(runtime.NewEnvironment(runtime.NewKVStoreService(keys[upgradetypes.StoreKey]), logger), skipUpgradeHeights, appCodec, homePath, app.BaseApp, authtypes.NewModuleAddress(govtypes.ModuleName).String())
 
 	// Register the proposal types
 	// Deprecated: Avoid adding new handlers, instead use the new proposal flow

x/upgrade/CHANGELOG.md

@@ -29,6 +29,10 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 * (x/upgrade) [#16244](https://github.com/cosmos/cosmos-sdk/pull/16244) Upgrade module no longer stores the app version but gets and sets the app version stored in the `ParamStore` of baseapp.
 
+### API Breaking Changes
+
+* [#19443](https://github.com/cosmos/cosmos-sdk/pull/19443) Creation of upgrade module receives `appmodule.Environment` instead of individual services
+
 ## [v0.1.1](https://github.com/cosmos/cosmos-sdk/releases/tag/x/upgrade/v0.1.1) - 2023-12-11
 
 ### Improvements
@@ -54,6 +58,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#16511](https://github.com/cosmos/cosmos-sdk/pull/16511) `plan.DownloadUpgrade` does not validate URL anymore. Call `plan.ValidateURL` before calling `plan.DownloadUpgrade` to validate the URL.
 * [#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".
 
+
 ### Bug Fixes
 
 * [#17421](https://github.com/cosmos/cosmos-sdk/pull/17421) Replace `BeginBlock` by `PreBlock`. Read [ADR-68](https://github.com/cosmos/cosmos-sdk/blob/main/docs/architecture/adr-068-preblock.md) for more information.

x/upgrade/depinject.go

@@ -6,7 +6,6 @@ import (
 	modulev1 "cosmossdk.io/api/cosmos/upgrade/module/v1"
 	"cosmossdk.io/core/address"
 	"cosmossdk.io/core/appmodule"
-	"cosmossdk.io/core/store"
 	"cosmossdk.io/depinject"
 	"cosmossdk.io/depinject/appconfig"
 	authtypes "cosmossdk.io/x/auth/types"
@@ -37,7 +36,7 @@ type ModuleInputs struct {
 	depinject.In
 
 	Config             *modulev1.Module
-	StoreService       store.KVStoreService
+	Environment        appmodule.Environment
 	Cdc                codec.Codec
 	AddressCodec       address.Codec
 	AppVersionModifier baseapp.AppVersionModifier
@@ -78,7 +77,7 @@ func ProvideModule(in ModuleInputs) ModuleOutputs {
 	}
 
 	// set the governance module account as the authority for conducting upgrades
-	k := keeper.NewKeeper(skipUpgradeHeights, in.StoreService, in.Cdc, homePath, in.AppVersionModifier, auth)
+	k := keeper.NewKeeper(in.Environment, skipUpgradeHeights, in.Cdc, homePath, in.AppVersionModifier, auth)
 	m := NewAppModule(k, in.AddressCodec)
 
 	return ModuleOutputs{UpgradeKeeper: k, Module: m}

x/upgrade/abci.go → x/upgrade/keeper/abci.go

@@ -1,4 +1,4 @@
-package upgrade
+package keeper
 
 import (
 	"context"
@@ -8,7 +8,6 @@ import (
 
 	"cosmossdk.io/core/appmodule"
 	storetypes "cosmossdk.io/store/types"
-	"cosmossdk.io/x/upgrade/keeper"
 	"cosmossdk.io/x/upgrade/types"
 
 	"github.com/cosmos/cosmos-sdk/telemetry"
@@ -23,17 +22,17 @@ 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 PreBlocker(ctx context.Context, k *keeper.Keeper) (appmodule.ResponsePreBlock, error) {
+func (k Keeper) PreBlocker(ctx context.Context) (appmodule.ResponsePreBlock, error) {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
 
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	blockHeight := sdkCtx.HeaderInfo().Height
+	blockHeight := k.environment.HeaderService.GetHeaderInfo(ctx).Height
 	plan, err := k.GetUpgradePlan(ctx)
 	if err != nil && !errors.Is(err, types.ErrNoUpgradePlanFound) {
 		return nil, err
 	}
 	found := err == nil
 
+	sdkCtx := sdk.UnwrapSDKContext(ctx) // TODO remove with consensus messages
 	if !k.DowngradeVerified() {
 		k.SetDowngradeVerified(true)
 		// This check will make sure that we are using a valid binary.

x/upgrade/abci_test.go → x/upgrade/keeper/abci_test.go

@@ -1,4 +1,4 @@
-package upgrade_test
+package keeper_test
 
 import (
 	"context"
@@ -114,6 +114,7 @@ func setupTest(t *testing.T, height int64, skip map[int64]bool) *TestSuite {
 	s.encCfg = moduletestutil.MakeTestEncodingConfig(upgrade.AppModuleBasic{})
 	key := storetypes.NewKVStoreKey(types.StoreKey)
 	storeService := runtime.NewKVStoreService(key)
+	env := runtime.NewEnvironment(storeService, log.NewNopLogger())
 	testCtx := testutil.DefaultContextWithDB(t, key, storetypes.NewTransientStoreKey("transient_test"))
 
 	s.baseApp = baseapp.NewBaseApp(
@@ -128,7 +129,7 @@ func setupTest(t *testing.T, height int64, skip map[int64]bool) *TestSuite {
 	authority, err := addresscodec.NewBech32Codec("cosmos").BytesToString(authtypes.NewModuleAddress(govModuleName))
 	require.NoError(t, err)
 
-	s.keeper = keeper.NewKeeper(skip, storeService, s.encCfg.Codec, t.TempDir(), s.baseApp, authority)
+	s.keeper = keeper.NewKeeper(env, skip, s.encCfg.Codec, t.TempDir(), s.baseApp, authority)
 
 	s.ctx = testCtx.Ctx.WithHeaderInfo(header.Info{Time: time.Now(), Height: height})
 
@@ -460,6 +461,7 @@ func TestDowngradeVerification(t *testing.T) {
 	encCfg := moduletestutil.MakeTestEncodingConfig(upgrade.AppModuleBasic{})
 	key := storetypes.NewKVStoreKey(types.StoreKey)
 	storeService := runtime.NewKVStoreService(key)
+	env := runtime.NewEnvironment(storeService, log.NewNopLogger())
 	testCtx := testutil.DefaultContextWithDB(t, key, storetypes.NewTransientStoreKey("transient_test"))
 	ctx := testCtx.Ctx.WithHeaderInfo(header.Info{Time: time.Now(), Height: 10})
 
@@ -468,7 +470,7 @@ func TestDowngradeVerification(t *testing.T) {
 	authority, err := addresscodec.NewBech32Codec("cosmos").BytesToString(authtypes.NewModuleAddress(govModuleName))
 	require.NoError(t, err)
 
-	k := keeper.NewKeeper(skip, storeService, encCfg.Codec, t.TempDir(), nil, authority)
+	k := keeper.NewKeeper(env, skip, encCfg.Codec, t.TempDir(), nil, authority)
 	m := upgrade.NewAppModule(k, addresscodec.NewBech32Codec("cosmos"))
 
 	// submit a plan.
@@ -517,7 +519,7 @@ func TestDowngradeVerification(t *testing.T) {
 		require.NoError(t, err)
 
 		// downgrade. now keeper does not have the handler.
-		k := keeper.NewKeeper(skip, storeService, encCfg.Codec, t.TempDir(), nil, authority)
+		k := keeper.NewKeeper(env, skip, encCfg.Codec, t.TempDir(), nil, authority)
 		m := upgrade.NewAppModule(k, addresscodec.NewBech32Codec("cosmos"))
 
 		// assertions
@@ -541,22 +543,3 @@ func TestDowngradeVerification(t *testing.T) {
 		}
 	}
 }
-
-type paramStore struct {
-	params cmtproto.ConsensusParams
-}
-
-var _ baseapp.ParamStore = (*paramStore)(nil)
-
-func (ps *paramStore) Set(_ context.Context, value cmtproto.ConsensusParams) error {
-	ps.params = value
-	return nil
-}
-
-func (ps paramStore) Has(_ context.Context) (bool, error) {
-	return true, nil
-}
-
-func (ps paramStore) Get(_ context.Context) (cmtproto.ConsensusParams, error) {
-	return ps.params, nil
-}

x/upgrade/keeper/grpc_query_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/suite"
 
 	"cosmossdk.io/core/header"
+	"cosmossdk.io/log"
 	storetypes "cosmossdk.io/store/types"
 	authtypes "cosmossdk.io/x/auth/types"
 	"cosmossdk.io/x/upgrade"
@@ -37,14 +38,15 @@ func (suite *UpgradeTestSuite) SetupTest() {
 	suite.encCfg = moduletestutil.MakeTestEncodingConfig(upgrade.AppModuleBasic{})
 	key := storetypes.NewKVStoreKey(types.StoreKey)
 	storeService := runtime.NewKVStoreService(key)
+	env := runtime.NewEnvironment(storeService, log.NewNopLogger())
 	testCtx := testutil.DefaultContextWithDB(suite.T(), key, storetypes.NewTransientStoreKey("transient_test"))
 	suite.ctx = testCtx.Ctx
 
 	skipUpgradeHeights := make(map[int64]bool)
 	authority, err := addresscodec.NewBech32Codec("cosmos").BytesToString(authtypes.NewModuleAddress(types.GovModuleName))
 	suite.Require().NoError(err)
 	suite.encodedAuthority = authority
-	suite.upgradeKeeper = keeper.NewKeeper(skipUpgradeHeights, storeService, suite.encCfg.Codec, suite.T().TempDir(), nil, authority)
+	suite.upgradeKeeper = keeper.NewKeeper(env, skipUpgradeHeights, suite.encCfg.Codec, suite.T().TempDir(), nil, authority)
 	err = suite.upgradeKeeper.SetModuleVersionMap(suite.ctx, module.VersionMap{
 		"bank": 0,
 	})

x/upgrade/keeper/keeper.go

@@ -14,7 +14,7 @@ import (
 
 	"github.com/hashicorp/go-metrics"
 
-	corestore "cosmossdk.io/core/store"
+	"cosmossdk.io/core/appmodule"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/log"
 	"cosmossdk.io/store/prefix"
@@ -25,16 +25,15 @@ import (
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/runtime"
 	"github.com/cosmos/cosmos-sdk/telemetry"
-	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/types/kv"
 	"github.com/cosmos/cosmos-sdk/types/module"
 )
 
 type Keeper struct {
-	homePath           string                          // root directory of app config
-	skipUpgradeHeights map[int64]bool                  // map of heights to skip for an upgrade
-	storeService       corestore.KVStoreService        // key to access x/upgrade store
+	homePath           string         // root directory of app config
+	skipUpgradeHeights map[int64]bool // map of heights to skip for an upgrade
+	environment        appmodule.Environment
 	cdc                codec.BinaryCodec               // App-wide binary codec
 	upgradeHandlers    map[string]types.UpgradeHandler // map of plan name to upgrade handler
 	versionModifier    xp.AppVersionModifier           // implements setting the protocol version field on BaseApp
@@ -49,11 +48,11 @@ type Keeper struct {
 // cdc - the app-wide binary codec
 // homePath - root directory of the application's config
 // vs - the interface implemented by baseapp which allows setting baseapp's protocol version field
-func NewKeeper(skipUpgradeHeights map[int64]bool, storeService corestore.KVStoreService, cdc codec.BinaryCodec, homePath string, vs xp.AppVersionModifier, authority string) *Keeper {
+func NewKeeper(env appmodule.Environment, skipUpgradeHeights map[int64]bool, cdc codec.BinaryCodec, homePath string, vs xp.AppVersionModifier, authority string) *Keeper {
 	k := &Keeper{
 		homePath:           homePath,
 		skipUpgradeHeights: skipUpgradeHeights,
-		storeService:       storeService,
+		environment:        env,
 		cdc:                cdc,
 		upgradeHandlers:    map[string]types.UpgradeHandler{},
 		versionModifier:    vs,
@@ -89,7 +88,7 @@ func (k Keeper) SetUpgradeHandler(name string, upgradeHandler types.UpgradeHandl
 // SetModuleVersionMap saves a given version map to state
 func (k Keeper) SetModuleVersionMap(ctx context.Context, vm module.VersionMap) error {
 	if len(vm) > 0 {
-		store := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
+		store := runtime.KVStoreAdapter(k.environment.KVStoreService.OpenKVStore(ctx))
 		versionStore := prefix.NewStore(store, []byte{types.VersionMapByte})
 		// Even though the underlying store (cachekv) store is sorted, we still
 		// prefer a deterministic iteration order of the map, to avoid undesired
@@ -116,7 +115,7 @@ func (k Keeper) SetModuleVersionMap(ctx context.Context, vm module.VersionMap) e
 // GetModuleVersionMap returns a map of key module name and value module consensus version
 // as defined in ADR-041.
 func (k Keeper) GetModuleVersionMap(ctx context.Context) (module.VersionMap, error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	prefix := []byte{types.VersionMapByte}
 	it, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix))
 	if err != nil {
@@ -138,7 +137,7 @@ func (k Keeper) GetModuleVersionMap(ctx context.Context) (module.VersionMap, err
 
 // GetModuleVersions gets a slice of module consensus versions
 func (k Keeper) GetModuleVersions(ctx context.Context) ([]*types.ModuleVersion, error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	prefix := []byte{types.VersionMapByte}
 	it, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix))
 	if err != nil {
@@ -163,7 +162,7 @@ func (k Keeper) GetModuleVersions(ctx context.Context) ([]*types.ModuleVersion,
 // getModuleVersion gets the version for a given module. If it doesn't exist it returns ErrNoModuleVersionFound, other
 // errors may be returned if there is an error reading from the store.
 func (k Keeper) getModuleVersion(ctx context.Context, name string) (uint64, error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	prefix := []byte{types.VersionMapByte}
 	it, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix))
 	if err != nil {
@@ -193,8 +192,7 @@ func (k Keeper) ScheduleUpgrade(ctx context.Context, plan types.Plan) error {
 
 	// NOTE: allow for the possibility of chains to schedule upgrades in begin block of the same block
 	// as a strategy for emergency hard fork recoveries
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	if plan.Height < sdkCtx.HeaderInfo().Height {
+	if plan.Height < k.environment.HeaderService.GetHeaderInfo(ctx).Height {
 		return errorsmod.Wrap(sdkerrors.ErrInvalidRequest, "upgrade cannot be scheduled in the past")
 	}
 
@@ -207,7 +205,7 @@ func (k Keeper) ScheduleUpgrade(ctx context.Context, plan types.Plan) error {
 		return errorsmod.Wrapf(sdkerrors.ErrInvalidRequest, "upgrade with name %s has already been completed", plan.Name)
 	}
 
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 
 	// clear any old IBC state stored by previous plan
 	oldPlan, err := k.GetUpgradePlan(ctx)
@@ -240,14 +238,14 @@ func (k Keeper) ScheduleUpgrade(ctx context.Context, plan types.Plan) error {
 
 // SetUpgradedClient sets the expected upgraded client for the next version of this chain at the last height the current chain will commit.
 func (k Keeper) SetUpgradedClient(ctx context.Context, planHeight int64, bz []byte) error {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	return store.Set(types.UpgradedClientKey(planHeight), bz)
 }
 
 // GetUpgradedClient gets the expected upgraded client for the next version of this chain. If not found it returns
 // ErrNoUpgradedClientFound, but other errors may be returned if there is an error reading from the store.
 func (k Keeper) GetUpgradedClient(ctx context.Context, height int64) ([]byte, error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	bz, err := store.Get(types.UpgradedClientKey(height))
 	if err != nil {
 		return nil, err
@@ -263,14 +261,14 @@ func (k Keeper) GetUpgradedClient(ctx context.Context, height int64) ([]byte, er
 // SetUpgradedConsensusState sets the expected upgraded consensus state for the next version of this chain
 // using the last height committed on this chain.
 func (k Keeper) SetUpgradedConsensusState(ctx context.Context, planHeight int64, bz []byte) error {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	return store.Set(types.UpgradedConsStateKey(planHeight), bz)
 }
 
 // GetUpgradedConsensusState gets the expected upgraded consensus state for the next version of this chain. If not found
 // it returns ErrNoUpgradedConsensusStateFound, but other errors may be returned if there is an error reading from the store.
 func (k Keeper) GetUpgradedConsensusState(ctx context.Context, lastHeight int64) ([]byte, error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	bz, err := store.Get(types.UpgradedConsStateKey(lastHeight))
 	if err != nil {
 		return nil, err
@@ -285,7 +283,7 @@ func (k Keeper) GetUpgradedConsensusState(ctx context.Context, lastHeight int64)
 
 // GetLastCompletedUpgrade returns the last applied upgrade name and height.
 func (k Keeper) GetLastCompletedUpgrade(ctx context.Context) (string, int64, error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	prefix := []byte{types.DoneByte}
 	it, err := store.ReverseIterator(prefix, storetypes.PrefixEndBytes(prefix))
 	if err != nil {
@@ -320,7 +318,7 @@ func encodeDoneKey(name string, height int64) []byte {
 
 // GetDoneHeight returns the height at which the given upgrade was executed
 func (k Keeper) GetDoneHeight(ctx context.Context, name string) (int64, error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	prefix := []byte{types.DoneByte}
 	it, err := store.Iterator(prefix, storetypes.PrefixEndBytes(prefix))
 	if err != nil {
@@ -341,7 +339,7 @@ func (k Keeper) GetDoneHeight(ctx context.Context, name string) (int64, error) {
 // ClearIBCState clears any planned IBC state
 func (k Keeper) ClearIBCState(ctx context.Context, lastHeight int64) error {
 	// delete IBC client and consensus state from store if this is IBC plan
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	err := store.Delete(types.UpgradedClientKey(lastHeight))
 	if err != nil {
 		return err
@@ -367,20 +365,19 @@ func (k Keeper) ClearUpgradePlan(ctx context.Context) error {
 		return err
 	}
 
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	return store.Delete(types.PlanKey())
 }
 
 // Logger returns a module-specific logger.
 func (k Keeper) Logger(ctx context.Context) log.Logger {
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	return sdkCtx.Logger().With("module", "x/"+types.ModuleName)
+	return k.environment.Logger.With("module", "x/"+types.ModuleName)
 }
 
 // GetUpgradePlan returns the currently scheduled Plan if any. If not found it returns
 // ErrNoUpgradePlanFound, but other errors may be returned if there is an error reading from the store.
 func (k Keeper) GetUpgradePlan(ctx context.Context) (plan types.Plan, err error) {
-	store := k.storeService.OpenKVStore(ctx)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 	bz, err := store.Get(types.PlanKey())
 	if err != nil {
 		return plan, err
@@ -400,11 +397,11 @@ func (k Keeper) GetUpgradePlan(ctx context.Context) (plan types.Plan, err error)
 
 // setDone marks this upgrade name as being done so the name can't be reused accidentally
 func (k Keeper) setDone(ctx context.Context, name string) error {
-	store := k.storeService.OpenKVStore(ctx)
-	sdkCtx := sdk.UnwrapSDKContext(ctx)
-	k.Logger(ctx).Debug("setting done", "height", sdkCtx.HeaderInfo().Height, "name", name)
+	store := k.environment.KVStoreService.OpenKVStore(ctx)
 
-	return store.Set(encodeDoneKey(name, sdkCtx.HeaderInfo().Height), []byte{1})
+	k.Logger(ctx).Debug("setting done", "height", k.environment.HeaderService.GetHeaderInfo(ctx).Height, "name", name)
+
+	return store.Set(encodeDoneKey(name, k.environment.HeaderService.GetHeaderInfo(ctx).Height), []byte{1})
 }
 
 // HasHandler returns true iff there is a handler registered for this name