feat: single binary upgrade mods

client/app/app.go

@@ -3,9 +3,8 @@ package app
 import (
 	"cosmossdk.io/depinject"
 	"cosmossdk.io/log"
+	storetypes "cosmossdk.io/store/types"
 	upgradekeeper "cosmossdk.io/x/upgrade/keeper"
-
-	abci "github.com/cometbft/cometbft/abci/types"
 	dbm "github.com/cosmos/cosmos-db"
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/client"
@@ -14,15 +13,19 @@ import (
 	"github.com/cosmos/cosmos-sdk/runtime"
 	servertypes "github.com/cosmos/cosmos-sdk/server/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/module"
+	"github.com/cosmos/cosmos-sdk/version"
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
 	distrkeeper "github.com/cosmos/cosmos-sdk/x/distribution/keeper"
+	paramstypes "github.com/cosmos/cosmos-sdk/x/params/types"
 	slashingkeeper "github.com/cosmos/cosmos-sdk/x/slashing/keeper"
 	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
+	"io"
 
 	"github.com/piplabs/story/client/app/keepers"
+	"github.com/piplabs/story/client/app/module"
 	"github.com/piplabs/story/client/comet"
+	appv1 "github.com/piplabs/story/client/pkg/appconsts/v1"
 	evmstakingkeeper "github.com/piplabs/story/client/x/evmstaking/keeper"
 	"github.com/piplabs/story/lib/errors"
 	"github.com/piplabs/story/lib/ethclient"
@@ -43,6 +46,11 @@ import (
 
 const Name = "story"
 
+const (
+	v1                    = appv1.Version
+	DefaultInitialVersion = v1
+)
+
 var (
 	_ runtime.AppI            = (*App)(nil)
 	_ servertypes.Application = (*App)(nil)
@@ -59,13 +67,21 @@ type App struct {
 	interfaceRegistry codectypes.InterfaceRegistry
 
 	Keepers keepers.Keepers
+
+	keyVersions map[uint64][]string
+	keys        map[string]*storetypes.KVStoreKey
+
+	// override the runtime baseapp's module manager to use the custom module manager
+	ModuleManager *module.Manager
+	configurator  module.Configurator
 }
 
 // newApp returns a reference to an initialized App.
 func newApp(
 	logger log.Logger,
 	db dbm.DB,
 	engineCl ethclient.EngineClient,
+	traceStore io.Writer,
 	baseAppOpts ...func(*baseapp.BaseApp),
 ) (*App, error) {
 	depCfg := depinject.Configs(
@@ -75,15 +91,20 @@ func newApp(
 		),
 	)
 
+	encodingConfig := MakeEncodingConfig(ModuleEncodingRegisters...)
+	appCodec := encodingConfig.Codec
+	txConfig := encodingConfig.TxConfig
+	interfaceRegistry := encodingConfig.InterfaceRegistry
+
 	var (
 		app        = new(App)
 		appBuilder = new(runtime.AppBuilder)
 	)
 	if err := depinject.Inject(depCfg,
 		&appBuilder,
-		&app.appCodec,
-		&app.txConfig,
-		&app.interfaceRegistry,
+		&appCodec,
+		&txConfig,
+		&interfaceRegistry,
 		&app.Keepers.AccountKeeper,
 		&app.Keepers.BankKeeper,
 		&app.Keepers.SignalKeeper,
@@ -92,56 +113,46 @@ func newApp(
 		&app.Keepers.DistrKeeper,
 		&app.Keepers.ConsensusParamsKeeper,
 		&app.Keepers.GovKeeper,
-		&app.Keepers.UpgradeKeeper,
+		// Story modules
 		&app.Keepers.EpochsKeeper,
 		&app.Keepers.EvmStakingKeeper,
 		&app.Keepers.EVMEngKeeper,
+		&app.Keepers.SignalKeeper,
 	); err != nil {
 		return nil, errors.Wrap(err, "dep inject")
 	}
 
-	baseAppOpts = append(baseAppOpts, func(bapp *baseapp.BaseApp) {
+	prepareOpt := func(bapp *baseapp.BaseApp) {
 		// Use evm engine to create block proposals.
 		// Note that we do not check MaxTxBytes since all EngineEVM transaction MUST be included since we cannot
 		// postpone them to the next block. Nit: we could drop some vote extensions though...?
 		bapp.SetPrepareProposal(app.Keepers.EVMEngKeeper.PrepareProposal)
 
 		// Route proposed messages to keepers for verification and external state updates.
 		bapp.SetProcessProposal(makeProcessProposalHandler(makeProcessProposalRouter(app), app.txConfig))
-	})
-
-	app.App = appBuilder.Build(db, nil, baseAppOpts...)
 
-	// Override the preblockers with custom PreBlocker function, which handles forks.
-	{
-		app.ModuleManager.SetOrderPreBlockers(preBlockers...)
-		app.SetPreBlocker(app.PreBlocker)
+		// This is to set the Cosmos SDK version used by the app.
+		// The app's version is set with bapp.SetProtocolVersion()
+		bapp.SetVersion(version.Version)
 	}
+	baseAppOpts = append(baseAppOpts, prepareOpt)
 
-	// Set "OrderEndBlockers" directly instead of using "SetOrderEndBlockers," which will panic since the staking module
-	// is missing in the "endBlockers", which is an intended behavior in Story. The panic message is:
-	// `panic: all modules must be defined when setting SetOrderEndBlockers, missing: [staking]`
-	{
-		app.ModuleManager.OrderEndBlockers = endBlockers
-		app.SetEndBlocker(app.EndBlocker)
-	}
+	app.App = appBuilder.Build(db, traceStore, baseAppOpts...)
+	app.keys = storetypes.NewKVStoreKeys(allStoreKeys()...) // TODO: ensure DI injected keys are matched here
+	app.keyVersions = versionedStoreKeys()
 
-	// Need to manually set the module version map, otherwise dep inject will NOT call `SetModuleVersionMap` for
-	// whatever reason that needs to be investigated. Since `SetModuleVersionMap` is not called, `fromVM` will have
-	// no entries (i.e. does not know about each module's consensus version) and will try to "add" modules during an
-	// upgrade. Specifically, the upgrade module will try to add all modules as new from version 0 to the latest version
-	// of each module since `fromVM` is empty on the very first upgrade.
-	app.SetInitChainer(func(ctx sdk.Context, req *abci.RequestInitChain) (*abci.ResponseInitChain, error) {
-		err := app.Keepers.UpgradeKeeper.SetModuleVersionMap(ctx, app.ModuleManager.GetVersionMap())
-		if err != nil {
-			return nil, errors.Wrap(err, "set module version map")
-		}
+	//app.ModuleManager.RegisterInvariants(&app.CrisisKeeper)
+	app.configurator = module.NewConfigurator(app.appCodec, app.MsgServiceRouter(), app.GRPCQueryRouter())
+	app.ModuleManager.RegisterServices(app.configurator)
 
-		return app.App.InitChainer(ctx, req)
-	})
+	// NOTE: Modules can't be modified or else must be passed by reference to the module manager
+	err := app.setupModuleManager()
+	if err != nil {
+		panic(err)
+	}
 
-	app.setupUpgradeHandlers()
-	app.setupUpgradeStoreLoaders()
+	// override module orders after DI
+	app.setModuleOrder()
 
 	if err := app.Load(true); err != nil {
 		return nil, errors.Wrap(err, "load app")
@@ -150,14 +161,19 @@ func newApp(
 	return app, nil
 }
 
-// PreBlocker application updates every pre block.
-func (a *App) PreBlocker(ctx sdk.Context, _ *abci.RequestFinalizeBlock) (*sdk.ResponsePreBlock, error) {
-	// All forks should be executed at their planned upgrade heights before any modules.
-	a.scheduleForkUpgrade(ctx)
-
-	res, err := a.ModuleManager.PreBlock(ctx)
+// EndBlocker executes application updates at the end of every block.
+func (a *App) EndBlocker(ctx sdk.Context) (sdk.EndBlock, error) {
+	res, err := a.ModuleManager.EndBlock(ctx)
 	if err != nil {
-		return nil, errors.Wrap(err, "module manager preblocker")
+		return sdk.EndBlock{}, errors.Wrap(err, "module manager endblocker")
+	}
+
+	currentVersion := a.AppVersion()
+	if shouldUpgrade, newVersion := a.Keepers.SignalKeeper.ShouldUpgrade(ctx); shouldUpgrade {
+		// Version changes must be increasing. Downgrades are not permitted
+		if newVersion > currentVersion {
+			a.SetProtocolVersion(newVersion)
+		}
 	}
 
 	return res, nil
@@ -181,6 +197,46 @@ func (a App) SetCometAPI(api comet.API) {
 	a.Keepers.EVMEngKeeper.SetCometAPI(api)
 }
 
+// LoadHeight loads a particular height
+func (a *App) LoadHeight(height int64) error {
+	return a.LoadVersion(height)
+}
+
+// SupportedVersions returns all the state machines that the
+// application supports
+func (a *App) SupportedVersions() []uint64 {
+	return a.ModuleManager.SupportedVersions()
+}
+
+// versionedKeys returns a map from moduleName to KV store key for the given app
+// version.
+func (a *App) versionedKeys(appVersion uint64) map[string]*storetypes.KVStoreKey {
+	output := make(map[string]*storetypes.KVStoreKey)
+	if keys, exists := a.keyVersions[appVersion]; exists {
+		for _, moduleName := range keys {
+			if key, exists := a.keys[moduleName]; exists {
+				output[moduleName] = key
+			}
+		}
+	}
+	return output
+}
+
+// baseKeys returns the base keys that are mounted to every version
+func (app *App) baseKeys() map[string]*storetypes.KVStoreKey {
+	return map[string]*storetypes.KVStoreKey{
+		// we need to know the app version to know what stores to mount
+		// thus the paramstore must always be a store that is mounted
+		paramstypes.StoreKey: app.keys[paramstypes.StoreKey],
+	}
+}
+
+// migrateModules performs migrations on existing modules that have registered migrations
+// between versions and initializes the state of new modules for the specified app version.
+func (a App) migrateModules(ctx sdk.Context, fromVersion, toVersion uint64) error {
+	return a.ModuleManager.RunMigrations(ctx, a.configurator, fromVersion, toVersion)
+}
+
 func (a App) GetEvmStakingKeeper() *evmstakingkeeper.Keeper {
 	return a.Keepers.EvmStakingKeeper
 }

client/app/app_config.go

@@ -13,11 +13,8 @@ import (
 	slashingmodulev1 "cosmossdk.io/api/cosmos/slashing/module/v1"
 	stakingmodulev1 "cosmossdk.io/api/cosmos/staking/module/v1"
 	txconfigv1 "cosmossdk.io/api/cosmos/tx/config/v1"
-	upgrademodulev1 "cosmossdk.io/api/cosmos/upgrade/module/v1"
 	"cosmossdk.io/core/appconfig"
 	"cosmossdk.io/depinject"
-	upgradetypes "cosmossdk.io/x/upgrade/types"
-
 	"github.com/cosmos/cosmos-sdk/runtime"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
@@ -86,19 +83,13 @@ var (
 		govtypes.ModuleName,
 		minttypes.ModuleName,
 		genutiltypes.ModuleName,
-		upgradetypes.ModuleName,
 		// Story modules
 		epochstypes.ModuleName,
 		evmenginetypes.ModuleName,
 		evmstakingtypes.ModuleName,
 		signaltypes.ModuleName,
 	}
 
-	// NOTE: upgrade module must come first, as upgrades might break state schema.
-	preBlockers = []string{
-		upgradetypes.ModuleName,
-	}
-
 	// 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
 	// CanWithdrawInvariant invariant.
@@ -204,10 +195,6 @@ var (
 				Name:   stakingtypes.ModuleName,
 				Config: appconfig.WrapAny(&stakingmodulev1.Module{}),
 			},
-			{
-				Name:   upgradetypes.ModuleName,
-				Config: appconfig.WrapAny(&upgrademodulev1.Module{}),
-			},
 			{
 				Name:   epochstypes.ModuleName,
 				Config: appconfig.WrapAny(&epochsmodule.Module{}),

client/app/encoding.go

@@ -0,0 +1,50 @@
+package app
+
+import (
+	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/codec"
+	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	"github.com/cosmos/cosmos-sdk/std"
+	"github.com/cosmos/cosmos-sdk/x/auth/tx"
+)
+
+type ModuleRegister interface {
+	RegisterLegacyAminoCodec(*codec.LegacyAmino)
+	RegisterInterfaces(codectypes.InterfaceRegistry)
+}
+
+// Config specifies the concrete encoding types to use for a given app.
+// This is provided for compatibility between protobuf and amino implementations.
+type EncodingConfig struct {
+	InterfaceRegistry codectypes.InterfaceRegistry
+	Codec             codec.Codec
+	TxConfig          client.TxConfig
+	Amino             *codec.LegacyAmino
+}
+
+// MakeConfig returns an encoding config for the app.
+func MakeEncodingConfig(moduleRegisters ...ModuleRegister) EncodingConfig {
+	interfaceRegistry := codectypes.NewInterfaceRegistry()
+	amino := codec.NewLegacyAmino()
+
+	// Register the standard types from the Cosmos SDK on interfaceRegistry and
+	// amino.
+	std.RegisterInterfaces(interfaceRegistry)
+	std.RegisterLegacyAminoCodec(amino)
+
+	// Register types from the moduleRegisters on interfaceRegistry and amino.
+	for _, moduleRegister := range moduleRegisters {
+		moduleRegister.RegisterInterfaces(interfaceRegistry)
+		moduleRegister.RegisterLegacyAminoCodec(amino)
+	}
+
+	protoCodec := codec.NewProtoCodec(interfaceRegistry)
+	txConfig := tx.NewTxConfig(protoCodec, tx.DefaultSignModes)
+
+	return EncodingConfig{
+		InterfaceRegistry: interfaceRegistry,
+		Codec:             protoCodec,
+		TxConfig:          txConfig,
+		Amino:             amino,
+	}
+}