feat: export genesis in simapp v2 (#21199)

Co-authored-by: marbar3778 <marbar3778@yahoo.com>
2024-08-16 13:15:45 +02:00 · 2024-08-16 13:15:45 +02:00 · aeeaca64da
commit aeeaca64da
parent e2ec889bb7
15 changed files with 426 additions and 67 deletions
--- a/runtime/v2/app.go
+++ b/runtime/v2/app.go
@ -79,6 +79,11 @@ func (a *App[T]) LoadHeight(height uint64) error {
 	return a.db.LoadVersion(height)
 }

+// LoadLatestHeight loads the latest height.
+func (a *App[T]) LoadLatestHeight() (uint64, error) {
+	return a.db.GetLatestVersion()
+}
+
 // Close is called in start cmd to gracefully cleanup resources.
 func (a *App[T]) Close() error {
 	return nil
--- a/runtime/v2/builder.go
+++ b/runtime/v2/builder.go
@ -69,7 +69,7 @@ func (a *AppBuilder[T]) RegisterModules(modules map[string]appmodulev2.AppModule

 // RegisterStores registers the provided store keys.
 // This method should only be used for registering extra stores
-// wiich is necessary for modules that not registered using the app config.
+// which is necessary for modules that not registered using the app config.
 // To be used in combination of RegisterModules.
 func (a *AppBuilder[T]) RegisterStores(keys ...string) {
 	a.app.storeKeys = append(a.app.storeKeys, keys...)
@ -175,6 +175,19 @@ func (a *AppBuilder[T]) Build(opts ...AppBuilderOption[T]) (*App[T], error) {
 			}
 			return nil
 		},
+		ExportGenesis: func(ctx context.Context, version uint64) ([]byte, error) {
+			genesisJson, err := a.app.moduleManager.ExportGenesisForModules(ctx)
+			if err != nil {
+				return nil, fmt.Errorf("failed to export genesis: %w", err)
+			}
+
+			bz, err := json.Marshal(genesisJson)
+			if err != nil {
+				return nil, fmt.Errorf("failed to marshal genesis: %w", err)
+			}
+
+			return bz, nil
+		},
 	}

 	appManager, err := appManagerBuilder.Build()
--- a/runtime/v2/manager.go
+++ b/runtime/v2/manager.go
@ -203,16 +203,13 @@ func (m *MM[T]) ExportGenesisForModules(
 		return nil, err
 	}

-	type genesisResult struct {
-		bz  json.RawMessage
-		err error
-	}
-
 	type ModuleI interface {
 		ExportGenesis(ctx context.Context) (json.RawMessage, error)
 	}

-	channels := make(map[string]chan genesisResult)
+	genesisData := make(map[string]json.RawMessage)
+
+	// TODO: make async export genesis https://github.com/cosmos/cosmos-sdk/issues/21303
 	for _, moduleName := range modulesToExport {
 		mod := m.modules[moduleName]
 		var moduleI ModuleI
@ -221,27 +218,16 @@ func (m *MM[T]) ExportGenesisForModules(
 			moduleI = module.(ModuleI)
 		} else if module, hasABCIGenesis := mod.(appmodulev2.HasGenesis); hasABCIGenesis {
 			moduleI = module.(ModuleI)
+		} else {
+			continue
 		}

-		channels[moduleName] = make(chan genesisResult)
-		go func(moduleI ModuleI, ch chan genesisResult) {
-			jm, err := moduleI.ExportGenesis(ctx)
-			if err != nil {
-				ch <- genesisResult{nil, err}
-				return
-			}
-			ch <- genesisResult{jm, nil}
-		}(moduleI, channels[moduleName])
-	}
-
-	genesisData := make(map[string]json.RawMessage)
-	for moduleName := range channels {
-		res := <-channels[moduleName]
-		if res.err != nil {
-			return nil, fmt.Errorf("genesis export error in %s: %w", moduleName, res.err)
+		res, err := moduleI.ExportGenesis(ctx)
+		if err != nil {
+			return nil, err
 		}

-		genesisData[moduleName] = res.bz
+		genesisData[moduleName] = res
 	}

 	return genesisData, nil
--- a/server/v2/appmanager/appmanager.go
+++ b/server/v2/appmanager/appmanager.go
@ -89,7 +89,24 @@ func (a AppManager[T]) InitGenesis(

 // ExportGenesis exports the genesis state of the application.
 func (a AppManager[T]) ExportGenesis(ctx context.Context, version uint64) ([]byte, error) {
-	bz, err := a.exportGenesis(ctx, version)
+	zeroState, err := a.db.StateAt(version)
+	if err != nil {
+		return nil, fmt.Errorf("unable to get latest state: %w", err)
+	}
+
+	bz := make([]byte, 0)
+	_, err = a.stf.RunWithCtx(ctx, zeroState, func(ctx context.Context) error {
+		if a.exportGenesis == nil {
+			return errors.New("export genesis function not set")
+		}
+
+		bz, err = a.exportGenesis(ctx, version)
+		if err != nil {
+			return fmt.Errorf("failed to export genesis state: %w", err)
+		}
+
+		return nil
+	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to export genesis state: %w", err)
 	}
--- a/server/v2/cometbft/go.mod
+++ b/server/v2/cometbft/go.mod
@ -26,7 +26,7 @@ require (
 	cosmossdk.io/errors v1.0.1
 	cosmossdk.io/log v1.4.0
 	cosmossdk.io/server/v2 v2.0.0-00010101000000-000000000000
-	cosmossdk.io/server/v2/appmanager v0.0.0-00010101000000-000000000000
+	cosmossdk.io/server/v2/appmanager v0.0.0-20240802110823-cffeedff643d
 	cosmossdk.io/store/v2 v2.0.0-00010101000000-000000000000
 	cosmossdk.io/x/consensus v0.0.0-00010101000000-000000000000
 	github.com/cometbft/cometbft v1.0.0-rc1
--- a/simapp/app.go
+++ b/simapp/app.go
@ -50,7 +50,6 @@ import (
 	circuittypes "cosmossdk.io/x/circuit/types"
 	"cosmossdk.io/x/consensus"
 	consensusparamkeeper "cosmossdk.io/x/consensus/keeper"
-	consensusparamtypes "cosmossdk.io/x/consensus/types"
 	consensustypes "cosmossdk.io/x/consensus/types"
 	distr "cosmossdk.io/x/distribution"
 	distrkeeper "cosmossdk.io/x/distribution/keeper"
@ -265,7 +264,7 @@ func NewSimApp(
 	keys := storetypes.NewKVStoreKeys(
 		authtypes.StoreKey, banktypes.StoreKey, stakingtypes.StoreKey,
 		minttypes.StoreKey, distrtypes.StoreKey, slashingtypes.StoreKey,
-		govtypes.StoreKey, consensusparamtypes.StoreKey, upgradetypes.StoreKey, feegrant.StoreKey,
+		govtypes.StoreKey, consensustypes.StoreKey, upgradetypes.StoreKey, feegrant.StoreKey,
 		evidencetypes.StoreKey, circuittypes.StoreKey,
 		authzkeeper.StoreKey, nftkeeper.StoreKey, group.StoreKey, pooltypes.StoreKey,
 		accounts.StoreKey, epochstypes.StoreKey,
@ -288,7 +287,7 @@ func NewSimApp(
 	cometService := runtime.NewContextAwareCometInfoService()

 	// set the BaseApp's parameter store
-	app.ConsensusParamsKeeper = consensusparamkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[consensusparamtypes.StoreKey]), logger.With(log.ModuleKey, "x/consensus")), authtypes.NewModuleAddress(govtypes.ModuleName).String())
+	app.ConsensusParamsKeeper = consensusparamkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[consensustypes.StoreKey]), logger.With(log.ModuleKey, "x/consensus")), authtypes.NewModuleAddress(govtypes.ModuleName).String())
 	bApp.SetParamStore(app.ConsensusParamsKeeper.ParamsStore)

 	// add keepers
@ -519,7 +518,7 @@ func NewSimApp(
 		group.ModuleName,
 		upgradetypes.ModuleName,
 		vestingtypes.ModuleName,
-		consensusparamtypes.ModuleName,
+		consensustypes.ModuleName,
 		circuittypes.ModuleName,
 		pooltypes.ModuleName,
 		epochstypes.ModuleName,
--- a/simapp/v2/app_di.go
+++ b/simapp/v2/app_di.go
@ -12,7 +12,6 @@ import (
 	"cosmossdk.io/depinject"
 	"cosmossdk.io/log"
 	"cosmossdk.io/runtime/v2"
-	serverv2 "cosmossdk.io/server/v2"
 	"cosmossdk.io/x/accounts"
 	authkeeper "cosmossdk.io/x/auth/keeper"
 	authzkeeper "cosmossdk.io/x/authz/keeper"
@ -92,7 +91,6 @@ func NewSimApp[T transaction.Tx](
 	logger log.Logger,
 	viper *viper.Viper,
 ) *SimApp[T] {
-	viper.Set(serverv2.FlagHome, DefaultNodeHome) // TODO possibly set earlier when viper is created
 	var (
 		app        = &SimApp[T]{}
 		appBuilder *runtime.AppBuilder[T]
--- a/simapp/v2/app_test.go
+++ b/simapp/v2/app_test.go
@ -0,0 +1,155 @@
+package simapp
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/cometbft/cometbft/types"
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/require"
+
+	app2 "cosmossdk.io/core/app"
+	"cosmossdk.io/core/comet"
+	context2 "cosmossdk.io/core/context"
+	"cosmossdk.io/core/store"
+	"cosmossdk.io/core/transaction"
+	"cosmossdk.io/log"
+	sdkmath "cosmossdk.io/math"
+	serverv2 "cosmossdk.io/server/v2"
+	comettypes "cosmossdk.io/server/v2/cometbft/types"
+	"cosmossdk.io/store/v2/db"
+	authtypes "cosmossdk.io/x/auth/types"
+	banktypes "cosmossdk.io/x/bank/types"
+
+	"github.com/cosmos/cosmos-sdk/crypto/keys/secp256k1"
+	"github.com/cosmos/cosmos-sdk/testutil/mock"
+	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func NewTestApp(t *testing.T) (*SimApp[transaction.Tx], context.Context) {
+	t.Helper()
+
+	logger := log.NewTestLogger(t)
+
+	vp := viper.New()
+	vp.Set("store.app-db-backend", string(db.DBTypeGoLevelDB))
+	vp.Set(serverv2.FlagHome, t.TempDir())
+
+	app := NewSimApp[transaction.Tx](logger, vp)
+	genesis := app.ModuleManager().DefaultGenesis()
+
+	privVal := mock.NewPV()
+	pubKey, err := privVal.GetPubKey()
+	require.NoError(t, err)
+
+	// create validator set with single validator
+	validator := types.NewValidator(pubKey, 1)
+	valSet := types.NewValidatorSet([]*types.Validator{validator})
+
+	// generate genesis account
+	senderPrivKey := secp256k1.GenPrivKey()
+	acc := authtypes.NewBaseAccount(senderPrivKey.PubKey().Address().Bytes(), senderPrivKey.PubKey(), 0, 0)
+	balance := banktypes.Balance{
+		Address: acc.GetAddress().String(),
+		Coins:   sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdkmath.NewInt(100000000000000))),
+	}
+
+	genesis, err = simtestutil.GenesisStateWithValSet(
+		app.AppCodec(),
+		genesis,
+		valSet,
+		[]authtypes.GenesisAccount{acc},
+		balance,
+	)
+	require.NoError(t, err)
+
+	genesisBytes, err := json.Marshal(genesis)
+	require.NoError(t, err)
+
+	st := app.GetStore().(comettypes.Store)
+	ci, err := st.LastCommitID()
+	require.NoError(t, err)
+
+	bz := sha256.Sum256([]byte{})
+
+	ctx := context.Background()
+
+	_, newState, err := app.InitGenesis(
+		ctx,
+		&app2.BlockRequest[transaction.Tx]{
+			Time:      time.Now(),
+			Hash:      bz[:],
+			ChainId:   "theChain",
+			AppHash:   ci.Hash,
+			IsGenesis: true,
+		},
+		genesisBytes,
+		nil,
+	)
+	require.NoError(t, err)
+
+	changes, err := newState.GetStateChanges()
+	require.NoError(t, err)
+
+	_, err = st.Commit(&store.Changeset{Changes: changes})
+	require.NoError(t, err)
+
+	return app, ctx
+}
+
+func MoveNextBlock(t *testing.T, app *SimApp[transaction.Tx], ctx context.Context) {
+	t.Helper()
+
+	bz := sha256.Sum256([]byte{})
+
+	st := app.GetStore().(comettypes.Store)
+	ci, err := st.LastCommitID()
+	require.NoError(t, err)
+
+	height, err := app.LoadLatestHeight()
+	require.NoError(t, err)
+
+	// TODO: this is a hack to set the comet info in the context for distribution module dependency.
+	ctx = context.WithValue(ctx, context2.CometInfoKey, comet.Info{
+		Evidence:        nil,
+		ValidatorsHash:  nil,
+		ProposerAddress: nil,
+		LastCommit:      comet.CommitInfo{},
+	})
+
+	_, newState, err := app.DeliverBlock(
+		ctx,
+		&app2.BlockRequest[transaction.Tx]{
+			Height:  height + 1,
+			Time:    time.Now(),
+			Hash:    bz[:],
+			AppHash: ci.Hash,
+		})
+	require.NoError(t, err)
+
+	changes, err := newState.GetStateChanges()
+	require.NoError(t, err)
+
+	_, err = st.Commit(&store.Changeset{Changes: changes})
+	require.NoError(t, err)
+}
+
+func TestSimAppExportAndBlockedAddrs_WithOneBlockProduced(t *testing.T) {
+	app, ctx := NewTestApp(t)
+
+	MoveNextBlock(t, app, ctx)
+
+	_, err := app.ExportAppStateAndValidators(nil)
+	require.NoError(t, err)
+}
+
+func TestSimAppExportAndBlockedAddrs_NoBlocksProduced(t *testing.T) {
+	app, _ := NewTestApp(t)
+
+	_, err := app.ExportAppStateAndValidators(nil)
+	require.NoError(t, err)
+}
--- a/simapp/v2/export.go
+++ b/simapp/v2/export.go
@ -1,10 +1,29 @@
 package simapp

 import (
-	servertypes "github.com/cosmos/cosmos-sdk/server/types"
+	"context"
+
+	v2 "github.com/cosmos/cosmos-sdk/x/genutil/v2"
 )

-// ExportAppStateAndValidators exports the state of the application for a genesis file.
-func (app *SimApp[T]) ExportAppStateAndValidators(forZeroHeight bool, jailAllowedAddrs, modulesToExport []string) (servertypes.ExportedApp, error) {
-	panic("not implemented")
+// ExportAppStateAndValidators exports the state of the application for a genesis
+// file.
+func (app *SimApp[T]) ExportAppStateAndValidators(jailAllowedAddrs []string) (v2.ExportedApp, error) {
+	// as if they could withdraw from the start of the next block
+	ctx := context.Background()
+
+	latestHeight, err := app.LoadLatestHeight()
+	if err != nil {
+		return v2.ExportedApp{}, err
+	}
+
+	genesis, err := app.ExportGenesis(ctx, latestHeight)
+	if err != nil {
+		return v2.ExportedApp{}, err
+	}
+
+	return v2.ExportedApp{
+		AppState: genesis,
+		Height:   int64(latestHeight),
+	}, nil
 }
--- a/simapp/v2/go.mod
+++ b/simapp/v2/go.mod
@ -12,7 +12,7 @@ require (
 	cosmossdk.io/runtime/v2 v2.0.0-00010101000000-000000000000
 	cosmossdk.io/server/v2 v2.0.0-20240718121635-a877e3e8048a
 	cosmossdk.io/server/v2/cometbft v0.0.0-00010101000000-000000000000
-	cosmossdk.io/store/v2 v2.0.0 // indirect
+	cosmossdk.io/store/v2 v2.0.0
 	cosmossdk.io/tools/confix v0.0.0-00010101000000-000000000000
 	cosmossdk.io/x/accounts v0.0.0-20240226161501-23359a0b6d91
 	cosmossdk.io/x/auth v0.0.0-00010101000000-000000000000
@ -32,7 +32,7 @@ require (
 	cosmossdk.io/x/staking v0.0.0-00010101000000-000000000000
 	cosmossdk.io/x/upgrade v0.0.0-20230613133644-0a778132a60f
 	github.com/cometbft/cometbft v1.0.0-rc1
-	github.com/cosmos/cosmos-db v1.0.2
+	github.com/cosmos/cosmos-db v1.0.2 // indirect
 	// this version is not used as it is always replaced by the latest Cosmos SDK version
 	github.com/cosmos/cosmos-sdk v0.53.0
 	github.com/spf13/cobra v1.8.1
@ -42,8 +42,6 @@ require (
 	google.golang.org/protobuf v1.34.2
 )

-require cosmossdk.io/core/testing v0.0.0-00010101000000-000000000000 // indirect
-
 require (
 	buf.build/gen/go/cometbft/cometbft/protocolbuffers/go v1.34.2-20240701160653-fedbb9acfd2f.2 // indirect
 	buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.34.2-20240130113600-88ef6483f90f.2 // indirect
@ -54,10 +52,11 @@ require (
 	cloud.google.com/go/iam v1.1.8 // indirect
 	cloud.google.com/go/storage v1.42.0 // indirect
 	cosmossdk.io/collections v0.4.0 // indirect
+	cosmossdk.io/core/testing v0.0.0-00010101000000-000000000000 // indirect
 	cosmossdk.io/errors v1.0.1 // indirect
 	cosmossdk.io/errors/v2 v2.0.0-20240731132947-df72853b3ca5 // indirect
 	cosmossdk.io/schema v0.1.1 // indirect
-	cosmossdk.io/server/v2/appmanager v0.0.0-00010101000000-000000000000 // indirect
+	cosmossdk.io/server/v2/appmanager v0.0.0-20240802110823-cffeedff643d // indirect
 	cosmossdk.io/server/v2/stf v0.0.0-00010101000000-000000000000 // indirect
 	cosmossdk.io/store v1.1.1-0.20240418092142-896cdf1971bc // indirect
 	cosmossdk.io/x/accounts/defaults/lockup v0.0.0-20240417181816-5e7aae0db1f5 // indirect
--- a/simapp/v2/simdv2/cmd/commands.go
+++ b/simapp/v2/simdv2/cmd/commands.go
@ -3,9 +3,7 @@ package cmd
 import (
 	"errors"
 	"fmt"
-	"io"

-	dbm "github.com/cosmos/cosmos-db"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"

@ -26,17 +24,21 @@ import (
 	"github.com/cosmos/cosmos-sdk/client/keys"
 	"github.com/cosmos/cosmos-sdk/client/rpc"
 	"github.com/cosmos/cosmos-sdk/server"
-	servertypes "github.com/cosmos/cosmos-sdk/server/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/genutil"
 	genutilcli "github.com/cosmos/cosmos-sdk/x/genutil/client/cli"
 	genutiltypes "github.com/cosmos/cosmos-sdk/x/genutil/types"
+	genutilv2 "github.com/cosmos/cosmos-sdk/x/genutil/v2"
+	v2 "github.com/cosmos/cosmos-sdk/x/genutil/v2/cli"
 )

 func newApp[T transaction.Tx](
 	logger log.Logger, viper *viper.Viper,
 ) serverv2.AppI[T] {
-	return serverv2.AppI[T](simapp.NewSimApp[T](logger, viper))
+	viper.Set(serverv2.FlagHome, simapp.DefaultNodeHome)
+
+	return serverv2.AppI[T](
+		simapp.NewSimApp[T](logger, viper))
 }

 func initRootCmd[T transaction.Tx](
@ -84,25 +86,11 @@ func initRootCmd[T transaction.Tx](
 // genesisCommand builds genesis-related `simd genesis` command. Users may provide application specific commands as a parameter
 func genesisCommand[T transaction.Tx](
 	moduleManager *runtimev2.MM[T],
-	appExport func(logger log.Logger,
-		height int64,
-		forZeroHeight bool,
-		jailAllowedAddrs []string,
-		viper *viper.Viper,
-		modulesToExport []string,
-	) (servertypes.ExportedApp, error),
+	appExport genutilv2.AppExporter,
 	cmds ...*cobra.Command,
 ) *cobra.Command {
-	compatAppExporter := func(logger log.Logger, db dbm.DB, traceWriter io.Writer, height int64, forZeroHeight bool, jailAllowedAddrs []string, appOpts servertypes.AppOptions, modulesToExport []string) (servertypes.ExportedApp, error) {
-		viperAppOpts, ok := appOpts.(*viper.Viper)
-		if !ok {
-			return servertypes.ExportedApp{}, errors.New("appOpts is not viper.Viper")
-		}
+	cmd := v2.Commands(moduleManager.Modules()[genutiltypes.ModuleName].(genutil.AppModule), moduleManager, appExport)

-		return appExport(logger, height, forZeroHeight, jailAllowedAddrs, viperAppOpts, modulesToExport)
-	}
-
-	cmd := genutilcli.Commands(moduleManager.Modules()[genutiltypes.ModuleName].(genutil.AppModule), moduleManager, compatAppExporter)
 	for _, subCmd := range cmds {
 		cmd.AddCommand(subCmd)
 	}
@ -159,26 +147,25 @@ func txCommand() *cobra.Command {
 func appExport[T transaction.Tx](
 	logger log.Logger,
 	height int64,
-	forZeroHeight bool,
 	jailAllowedAddrs []string,
 	viper *viper.Viper,
-	modulesToExport []string,
-) (servertypes.ExportedApp, error) {
+) (genutilv2.ExportedApp, error) {
 	// overwrite the FlagInvCheckPeriod
 	viper.Set(server.FlagInvCheckPeriod, 1)
+	viper.Set(serverv2.FlagHome, simapp.DefaultNodeHome)

 	var simApp *simapp.SimApp[T]
 	if height != -1 {
 		simApp = simapp.NewSimApp[T](logger, viper)

 		if err := simApp.LoadHeight(uint64(height)); err != nil {
-			return servertypes.ExportedApp{}, err
+			return genutilv2.ExportedApp{}, err
 		}
 	} else {
 		simApp = simapp.NewSimApp[T](logger, viper)
 	}

-	return simApp.ExportAppStateAndValidators(forZeroHeight, jailAllowedAddrs, modulesToExport)
+	return simApp.ExportAppStateAndValidators(jailAllowedAddrs)
 }

 var _ transaction.Codec[transaction.Tx] = &genericTxDecoder[transaction.Tx]{}
--- a/x/genutil/v2/cli/commands.go
+++ b/x/genutil/v2/cli/commands.go
@ -0,0 +1,47 @@
+package cli
+
+import (
+	"encoding/json"
+
+	"github.com/spf13/cobra"
+
+	banktypes "cosmossdk.io/x/bank/types"
+
+	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/x/genutil"
+	"github.com/cosmos/cosmos-sdk/x/genutil/client/cli"
+	genutiltypes "github.com/cosmos/cosmos-sdk/x/genutil/types"
+	v2 "github.com/cosmos/cosmos-sdk/x/genutil/v2"
+)
+
+type genesisMM interface {
+	DefaultGenesis() map[string]json.RawMessage
+	ValidateGenesis(genesisData map[string]json.RawMessage) error
+}
+
+// Commands adds core sdk's sub-commands into genesis command.
+func Commands(genutilModule genutil.AppModule, genMM genesisMM, appExport v2.AppExporter) *cobra.Command {
+	return CommandsWithCustomMigrationMap(genutilModule, genMM, appExport, genutiltypes.MigrationMap{})
+}
+
+// CommandsWithCustomMigrationMap adds core sdk's sub-commands into genesis command with custom migration map.
+// This custom migration map can be used by the application to add its own migration map.
+func CommandsWithCustomMigrationMap(genutilModule genutil.AppModule, genMM genesisMM, appExport v2.AppExporter, migrationMap genutiltypes.MigrationMap) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:                        "genesis",
+		Short:                      "Application's genesis-related subcommands",
+		DisableFlagParsing:         false,
+		SuggestionsMinimumDistance: 2,
+		RunE:                       client.ValidateCmd,
+	}
+	cmd.AddCommand(
+		cli.GenTxCmd(genMM, banktypes.GenesisBalancesIterator{}),
+		cli.MigrateGenesisCmd(migrationMap),
+		cli.CollectGenTxsCmd(genutilModule.GenTxValidator()),
+		cli.ValidateGenesisCmd(genMM),
+		cli.AddGenesisAccountCmd(),
+		ExportCmd(appExport),
+	)
+
+	return cmd
+}
--- a/x/genutil/v2/cli/export.go
+++ b/x/genutil/v2/cli/export.go
@ -0,0 +1,106 @@
+package cli
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/spf13/cobra"
+
+	"github.com/cosmos/cosmos-sdk/client"
+	"github.com/cosmos/cosmos-sdk/client/flags"
+	"github.com/cosmos/cosmos-sdk/version"
+	genutiltypes "github.com/cosmos/cosmos-sdk/x/genutil/types"
+	v2 "github.com/cosmos/cosmos-sdk/x/genutil/v2"
+)
+
+const (
+	flagHeight           = "height"
+	flagJailAllowedAddrs = "jail-allowed-addrs"
+)
+
+// ExportCmd dumps app state to JSON.
+func ExportCmd(appExporter v2.AppExporter) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "export",
+		Short: "Export state to JSON",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, _ []string) error {
+			config := client.GetConfigFromCmd(cmd)
+			viper := client.GetViperFromCmd(cmd)
+			logger := client.GetLoggerFromCmd(cmd)
+
+			if _, err := os.Stat(config.GenesisFile()); os.IsNotExist(err) {
+				return err
+			}
+
+			if appExporter == nil {
+				if _, err := fmt.Fprintln(cmd.ErrOrStderr(), "WARNING: App exporter not defined. Returning genesis file."); err != nil {
+					return err
+				}
+
+				// Open file in read-only mode so we can copy it to stdout.
+				// It is possible that the genesis file is large,
+				// so we don't need to read it all into memory
+				// before we stream it out.
+				f, err := os.OpenFile(config.GenesisFile(), os.O_RDONLY, 0)
+				if err != nil {
+					return err
+				}
+				defer f.Close()
+
+				if _, err := io.Copy(cmd.OutOrStdout(), f); err != nil {
+					return err
+				}
+
+				return nil
+			}
+
+			height, _ := cmd.Flags().GetInt64(flagHeight)
+			jailAllowedAddrs, _ := cmd.Flags().GetStringSlice(flagJailAllowedAddrs)
+			outputDocument, _ := cmd.Flags().GetString(flags.FlagOutputDocument)
+
+			exported, err := appExporter(logger, height, jailAllowedAddrs, viper)
+			if err != nil {
+				return fmt.Errorf("error exporting state: %w", err)
+			}
+
+			appGenesis, err := genutiltypes.AppGenesisFromFile(config.GenesisFile())
+			if err != nil {
+				return err
+			}
+
+			// set current binary version
+			appGenesis.AppName = version.AppName
+			appGenesis.AppVersion = version.Version
+
+			appGenesis.AppState = exported.AppState
+			appGenesis.InitialHeight = exported.Height
+
+			out, err := json.Marshal(appGenesis)
+			if err != nil {
+				return err
+			}
+
+			if outputDocument == "" {
+				// Copy the entire genesis file to stdout.
+				_, err := io.Copy(cmd.OutOrStdout(), bytes.NewReader(out))
+				return err
+			}
+
+			if err = appGenesis.SaveAs(outputDocument); err != nil {
+				return err
+			}
+
+			return nil
+		},
+	}
+
+	cmd.Flags().Int64(flagHeight, -1, "Export state from a particular height (-1 means latest height)")
+	cmd.Flags().StringSlice(flagJailAllowedAddrs, []string{}, "Comma-separated list of operator addresses of jailed validators to unjail")
+	cmd.Flags().String(flags.FlagOutputDocument, "", "Exported state is written to the given file instead of STDOUT")
+
+	return cmd
+}
--- a/x/genutil/v2/types.go
+++ b/x/genutil/v2/types.go
@ -0,0 +1,27 @@
+package v2
+
+import (
+	"encoding/json"
+
+	"github.com/spf13/viper"
+
+	"cosmossdk.io/log"
+)
+
+// AppExporter is a function that dumps all app state to
+// JSON-serializable structure and returns the current validator set.
+type AppExporter func(
+	logger log.Logger,
+	height int64,
+	jailAllowedAddrs []string,
+	viper *viper.Viper,
+) (ExportedApp, error)
+
+// ExportedApp represents an exported app state, along with
+// validators, consensus params and latest app height.
+type ExportedApp struct {
+	// AppState is the application state as JSON.
+	AppState json.RawMessage
+	// Height is the app's latest block height.
+	Height int64
+}
--- a/x/staking/genesis.go
+++ b/x/staking/genesis.go
@ -1,6 +1,7 @@
 package staking

 import (
+	"context"
 	"fmt"

 	cmttypes "github.com/cometbft/cometbft/types"
@ -14,7 +15,7 @@ import (
 )

 // WriteValidators returns a slice of bonded genesis validators.
-func WriteValidators(ctx sdk.Context, keeper *keeper.Keeper) (vals []cmttypes.GenesisValidator, returnErr error) {
+func WriteValidators(ctx context.Context, keeper *keeper.Keeper) (vals []cmttypes.GenesisValidator, returnErr error) {
 	err := keeper.LastValidatorPower.Walk(ctx, nil, func(key []byte, _ gogotypes.Int64Value) (bool, error) {
 		validator, err := keeper.GetValidator(ctx, key)
 		if err != nil {