feat(server): Allow calling back into the application struct in PostSetup. (#19455)

Devon Bear · tac0turtle · web-flow · commit 23359a0b6d91 · 2024-02-26T16:15:01.000Z
Co-authored-by: Marko &lt;marbar3778@yahoo.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -56,6 +56,7 @@ Every module contains its own CHANGELOG.md. Please refer to the module you are i
 
 ### Improvements
 
+* (server) [#19455](https://github.com/cosmos/cosmos-sdk/pull/19455) Allow calling back into the application struct in PostSetup.
 * (types) [#19512](https://github.com/cosmos/cosmos-sdk/pull/19512) The notion of basic manager does not exist anymore.
     * The module manager now can do everything that the basic manager was doing.
     * `AppModuleBasic` has been deprecated for extension interfaces. Modules can now implement `HasRegisterInterfaces`, `HasGRPCGateway` and `HasAminoCodec` when relevant.
diff --git a/client/pruning/main.go b/client/pruning/main.go
@@ -22,7 +22,7 @@ const FlagAppDBBackend = "app-db-backend"
 
 // Cmd prunes the sdk root multi store history versions based on the pruning options
 // specified by command flags.
-func Cmd(appCreator servertypes.AppCreator) *cobra.Command {
+func Cmd[T servertypes.Application](appCreator servertypes.AppCreator[T]) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "prune [pruning-method]",
 		Short: "Prune app history states by keeping the recent heights and deleting old heights",
diff --git a/client/snapshot/cmd.go b/client/snapshot/cmd.go
@@ -7,7 +7,7 @@ import (
 )
 
 // Cmd returns the snapshots group command
-func Cmd(appCreator servertypes.AppCreator) *cobra.Command {
+func Cmd[T servertypes.Application](appCreator servertypes.AppCreator[T]) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "snapshots",
 		Short: "Manage local snapshots",
diff --git a/client/snapshot/export.go b/client/snapshot/export.go
@@ -10,7 +10,7 @@ import (
 )
 
 // ExportSnapshotCmd returns a command to take a snapshot of the application state
-func ExportSnapshotCmd(appCreator servertypes.AppCreator) *cobra.Command {
+func ExportSnapshotCmd[T servertypes.Application](appCreator servertypes.AppCreator[T]) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "export",
 		Short: "Export app state to snapshot store",
diff --git a/client/snapshot/restore.go b/client/snapshot/restore.go
@@ -14,7 +14,7 @@ import (
 )
 
 // RestoreSnapshotCmd returns a command to restore a snapshot
-func RestoreSnapshotCmd(appCreator servertypes.AppCreator) *cobra.Command {
+func RestoreSnapshotCmd[T servertypes.Application](appCreator servertypes.AppCreator[T]) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "restore <height> <format>",
 		Short: "Restore app state from local snapshot",
diff --git a/server/cmt_cmds.go b/server/cmt_cmds.go
@@ -356,7 +356,7 @@ func parseOptionalHeight(heightStr string) (*int64, error) {
 	return &tmp, nil
 }
 
-func BootstrapStateCmd(appCreator types.AppCreator) *cobra.Command {
+func BootstrapStateCmd[T types.Application](appCreator types.AppCreator[T]) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "bootstrap-state",
 		Short: "Bootstrap CometBFT state at an arbitrary block height using a light client",
diff --git a/server/rollback.go b/server/rollback.go
@@ -10,7 +10,7 @@ import (
 )
 
 // NewRollbackCmd creates a command to rollback CometBFT and multistore state by one height.
-func NewRollbackCmd(appCreator types.AppCreator) *cobra.Command {
+func NewRollbackCmd[T types.Application](appCreator types.AppCreator[T]) *cobra.Command {
 	var removeBlock bool
 
 	cmd := &cobra.Command{
diff --git a/server/start.go b/server/start.go
@@ -108,26 +108,28 @@ const (
 )
 
 // StartCmdOptions defines options that can be customized in `StartCmdWithOptions`,
-type StartCmdOptions struct {
+type StartCmdOptions[T types.Application] struct {
 	// DBOpener can be used to customize db opening, for example customize db options or support different db backends,
 	// default to the builtin db opener.
 	DBOpener func(rootDir string, backendType dbm.BackendType) (dbm.DB, error)
 	// PostSetup can be used to setup extra services under the same cancellable context,
 	// it's not called in stand-alone mode, only for in-process mode.
-	PostSetup func(svrCtx *Context, clientCtx client.Context, ctx context.Context, g *errgroup.Group) error
+	PostSetup func(app T, svrCtx *Context, clientCtx client.Context, ctx context.Context, g *errgroup.Group) error
+	// PostSetupStandalone can be used to setup extra services under the same cancellable context,
+	PostSetupStandalone func(app T, svrCtx *Context, clientCtx client.Context, ctx context.Context, g *errgroup.Group) error
 	// AddFlags add custom flags to start cmd
 	AddFlags func(cmd *cobra.Command)
 }
 
 // StartCmd runs the service passed in, either stand-alone or in-process with
 // CometBFT.
-func StartCmd(appCreator types.AppCreator) *cobra.Command {
-	return StartCmdWithOptions(appCreator, StartCmdOptions{})
+func StartCmd[T types.Application](appCreator types.AppCreator[T]) *cobra.Command {
+	return StartCmdWithOptions(appCreator, StartCmdOptions[T]{})
 }
 
 // StartCmdWithOptions runs the service passed in, either stand-alone or in-process with
 // CometBFT.
-func StartCmdWithOptions(appCreator types.AppCreator, opts StartCmdOptions) *cobra.Command {
+func StartCmdWithOptions[T types.Application](appCreator types.AppCreator[T], opts StartCmdOptions[T]) *cobra.Command {
 	if opts.DBOpener == nil {
 		opts.DBOpener = OpenDB
 	}
@@ -199,13 +201,13 @@ is performed. Note, when enabled, gRPC will also be automatically enabled.
 	return cmd
 }
 
-func start(svrCtx *Context, clientCtx client.Context, appCreator types.AppCreator, withCmt bool, opts StartCmdOptions) error {
+func start[T types.Application](svrCtx *Context, clientCtx client.Context, appCreator types.AppCreator[T], withCmt bool, opts StartCmdOptions[T]) error {
 	svrCfg, err := getAndValidateConfig(svrCtx)
 	if err != nil {
 		return err
 	}
 
-	app, appCleanupFn, err := startApp(svrCtx, appCreator, opts)
+	app, appCleanupFn, err := startApp[T](svrCtx, appCreator, opts)
 	if err != nil {
 		return err
 	}
@@ -219,12 +221,12 @@ func start(svrCtx *Context, clientCtx client.Context, appCreator types.AppCreato
 	emitServerInfoMetrics()
 
 	if !withCmt {
-		return startStandAlone(svrCtx, svrCfg, clientCtx, app, metrics, opts)
+		return startStandAlone[T](svrCtx, svrCfg, clientCtx, app, metrics, opts)
 	}
-	return startInProcess(svrCtx, svrCfg, clientCtx, app, metrics, opts)
+	return startInProcess[T](svrCtx, svrCfg, clientCtx, app, metrics, opts)
 }
 
-func startStandAlone(svrCtx *Context, svrCfg serverconfig.Config, clientCtx client.Context, app types.Application, metrics *telemetry.Metrics, opts StartCmdOptions) error {
+func startStandAlone[T types.Application](svrCtx *Context, svrCfg serverconfig.Config, clientCtx client.Context, app T, metrics *telemetry.Metrics, opts StartCmdOptions[T]) error {
 	addr := svrCtx.Viper.GetString(flagAddress)
 	transport := svrCtx.Viper.GetString(flagTransport)
 
@@ -271,6 +273,12 @@ func startStandAlone(svrCtx *Context, svrCfg serverconfig.Config, clientCtx clie
 		return err
 	}
 
+	if opts.PostSetup != nil {
+		if err := opts.PostSetupStandalone(app, svrCtx, clientCtx, ctx, g); err != nil {
+			return err
+		}
+	}
+
 	g.Go(func() error {
 		if err := svr.Start(); err != nil {
 			svrCtx.Logger.Error("failed to start out-of-process ABCI server", "err", err)
@@ -287,8 +295,8 @@ func startStandAlone(svrCtx *Context, svrCfg serverconfig.Config, clientCtx clie
 	return g.Wait()
 }
 
-func startInProcess(svrCtx *Context, svrCfg serverconfig.Config, clientCtx client.Context, app types.Application,
-	metrics *telemetry.Metrics, opts StartCmdOptions,
+func startInProcess[T types.Application](svrCtx *Context, svrCfg serverconfig.Config, clientCtx client.Context, app T,
+	metrics *telemetry.Metrics, opts StartCmdOptions[T],
 ) error {
 	cmtCfg := svrCtx.Config
 	home := cmtCfg.RootDir
@@ -334,7 +342,7 @@ func startInProcess(svrCtx *Context, svrCfg serverconfig.Config, clientCtx clien
 	}
 
 	if opts.PostSetup != nil {
-		if err := opts.PostSetup(svrCtx, clientCtx, ctx, g); err != nil {
+		if err := opts.PostSetup(app, svrCtx, clientCtx, ctx, g); err != nil {
 			return err
 		}
 	}
@@ -585,7 +593,7 @@ func getCtx(svrCtx *Context, block bool) (*errgroup.Group, context.Context) {
 	return g, ctx
 }
 
-func startApp(svrCtx *Context, appCreator types.AppCreator, opts StartCmdOptions) (app types.Application, cleanupFn func(), err error) {
+func startApp[T types.Application](svrCtx *Context, appCreator types.AppCreator[T], opts StartCmdOptions[T]) (app T, cleanupFn func(), err error) {
 	traceWriter, traceCleanupFn, err := SetupTraceWriter(svrCtx.Logger, svrCtx.Viper.GetString(flagTraceStore))
 	if err != nil {
 		return app, traceCleanupFn, err
@@ -598,10 +606,12 @@ func startApp(svrCtx *Context, appCreator types.AppCreator, opts StartCmdOptions
 	}
 
 	if isTestnet, ok := svrCtx.Viper.Get(KeyIsTestnet).(bool); ok && isTestnet {
-		app, err = testnetify(svrCtx, home, appCreator, db, traceWriter)
+		var appPtr *T
+		appPtr, err = testnetify[T](svrCtx, home, appCreator, db, traceWriter)
 		if err != nil {
 			return app, traceCleanupFn, err
 		}
+		app = *appPtr
 	} else {
 		app = appCreator(svrCtx.Logger, db, traceWriter, svrCtx.Viper)
 	}
@@ -618,8 +628,8 @@ func startApp(svrCtx *Context, appCreator types.AppCreator, opts StartCmdOptions
 // InPlaceTestnetCreator utilizes the provided chainID and operatorAddress as well as the local private validator key to
 // control the network represented in the data folder. This is useful to create testnets nearly identical to your
 // mainnet environment.
-func InPlaceTestnetCreator(testnetAppCreator types.AppCreator) *cobra.Command {
-	opts := StartCmdOptions{}
+func InPlaceTestnetCreator[T types.Application](testnetAppCreator types.AppCreator[T]) *cobra.Command {
+	opts := StartCmdOptions[T]{}
 	if opts.DBOpener == nil {
 		opts.DBOpener = OpenDB
 	}
@@ -711,7 +721,7 @@ you want to test the upgrade handler itself.
 
 // testnetify modifies both state and blockStore, allowing the provided operator address and local validator key to control the network
 // that the state in the data folder represents. The chainID of the local genesis file is modified to match the provided chainID.
-func testnetify(ctx *Context, home string, testnetAppCreator types.AppCreator, db dbm.DB, traceWriter io.WriteCloser) (types.Application, error) {
+func testnetify[T types.Application](ctx *Context, home string, testnetAppCreator types.AppCreator[T], db dbm.DB, traceWriter io.WriteCloser) (*T, error) {
 	config := ctx.Config
 
 	newChainID, ok := ctx.Viper.Get(KeyNewChainID).(string)
@@ -933,11 +943,11 @@ func testnetify(ctx *Context, home string, testnetAppCreator types.AppCreator, d
 		return nil, err
 	}
 
-	return testnetApp, err
+	return &testnetApp, err
 }
 
 // addStartNodeFlags should be added to any CLI commands that start the network.
-func addStartNodeFlags(cmd *cobra.Command, opts StartCmdOptions) {
+func addStartNodeFlags[T types.Application](cmd *cobra.Command, opts StartCmdOptions[T]) {
 	cmd.Flags().Bool(flagWithComet, true, "Run abci app embedded in-process with CometBFT")
 	cmd.Flags().String(flagAddress, "tcp://127.0.0.1:26658", "Listen address")
 	cmd.Flags().String(flagTransport, "socket", "Transport protocol: socket, grpc")
diff --git a/server/types/app.go b/server/types/app.go
@@ -66,7 +66,7 @@ type (
 
 	// AppCreator is a function that allows us to lazily initialize an
 	// application using various configurations.
-	AppCreator func(log.Logger, dbm.DB, io.Writer, AppOptions) Application
+	AppCreator[T Application] func(log.Logger, dbm.DB, io.Writer, AppOptions) T
 
 	// ModuleInitFlags takes a start command and adds modules specific init flags.
 	ModuleInitFlags func(startCmd *cobra.Command)
diff --git a/server/util.go b/server/util.go
@@ -324,7 +324,7 @@ func interceptConfigs(rootViper *viper.Viper, customAppTemplate string, customCo
 }
 
 // add server commands
-func AddCommands(rootCmd *cobra.Command, appCreator types.AppCreator, addStartFlags types.ModuleInitFlags) {
+func AddCommands[T types.Application](rootCmd *cobra.Command, appCreator types.AppCreator[T], addStartFlags types.ModuleInitFlags) {
 	cometCmd := &cobra.Command{
 		Use:     "comet",
 		Aliases: []string{"cometbft", "tendermint"},
@@ -355,7 +355,7 @@ func AddCommands(rootCmd *cobra.Command, appCreator types.AppCreator, addStartFl
 }
 
 // AddTestnetCreatorCommand allows chains to create a testnet from the state existing in their node's data directory.
-func AddTestnetCreatorCommand(rootCmd *cobra.Command, appCreator types.AppCreator, addStartFlags types.ModuleInitFlags) {
+func AddTestnetCreatorCommand[T types.Application](rootCmd *cobra.Command, appCreator types.AppCreator[T], addStartFlags types.ModuleInitFlags) {
 	testnetCreateCmd := InPlaceTestnetCreator(appCreator)
 	addStartFlags(testnetCreateCmd)
 	rootCmd.AddCommand(testnetCreateCmd)
diff --git a/server/util_test.go b/server/util_test.go
@@ -50,7 +50,7 @@ func TestGetAppDBBackend(t *testing.T) {
 
 func TestInterceptConfigsPreRunHandlerCreatesConfigFilesWhenMissing(t *testing.T) {
 	tempDir := t.TempDir()
-	cmd := server.StartCmd(nil)
+	cmd := server.StartCmd[servertypes.Application](nil)
 	cmd.PersistentFlags().String(flags.FlagHome, "/foobar", "")
 	if err := cmd.PersistentFlags().Set(flags.FlagHome, tempDir); err != nil {
 		t.Fatalf("Could not set home flag [%T] %v", err, err)
@@ -129,7 +129,7 @@ func TestInterceptConfigsPreRunHandlerReadsConfigToml(t *testing.T) {
 		t.Fatalf("Failed closing config.toml: %v", err)
 	}
 
-	cmd := server.StartCmd(nil)
+	cmd := server.StartCmd[servertypes.Application](nil)
 	cmd.PersistentFlags().String(flags.FlagHome, "/foobar", "")
 	if err := cmd.PersistentFlags().Set(flags.FlagHome, tempDir); err != nil {
 		t.Fatalf("Could not set home flag [%T] %v", err, err)
@@ -172,7 +172,7 @@ func TestInterceptConfigsPreRunHandlerReadsAppToml(t *testing.T) {
 	if err := writer.Close(); err != nil {
 		t.Fatalf("Failed closing app.toml: %v", err)
 	}
-	cmd := server.StartCmd(nil)
+	cmd := server.StartCmd[servertypes.Application](nil)
 	cmd.PersistentFlags().String(flags.FlagHome, tempDir, "")
 
 	cmd.PreRunE = preRunETestImpl
@@ -194,7 +194,7 @@ func TestInterceptConfigsPreRunHandlerReadsAppToml(t *testing.T) {
 func TestInterceptConfigsPreRunHandlerReadsFlags(t *testing.T) {
 	const testAddr = "tcp://127.1.2.3:12345"
 	tempDir := t.TempDir()
-	cmd := server.StartCmd(nil)
+	cmd := server.StartCmd[servertypes.Application](nil)
 	cmd.PersistentFlags().String(flags.FlagHome, "/foobar", "")
 	if err := cmd.PersistentFlags().Set(flags.FlagHome, tempDir); err != nil {
 		t.Fatalf("Could not set home flag [%T] %v", err, err)
@@ -224,7 +224,7 @@ func TestInterceptConfigsPreRunHandlerReadsFlags(t *testing.T) {
 func TestInterceptConfigsPreRunHandlerReadsEnvVars(t *testing.T) {
 	const testAddr = "tcp://127.1.2.3:12345"
 	tempDir := t.TempDir()
-	cmd := server.StartCmd(nil)
+	cmd := server.StartCmd[servertypes.Application](nil)
 	cmd.PersistentFlags().String(flags.FlagHome, "/foobar", "")
 	if err := cmd.PersistentFlags().Set(flags.FlagHome, tempDir); err != nil {
 		t.Fatalf("Could not set home flag [%T] %v", err, err)
@@ -314,7 +314,7 @@ func newPrecedenceCommon(t *testing.T) precedenceCommon {
 	})
 
 	// Set up the command object that is used in this test
-	retval.cmd = server.StartCmd(nil)
+	retval.cmd = server.StartCmd[servertypes.Application](nil)
 	retval.cmd.PersistentFlags().String(flags.FlagHome, tempDir, "")
 	retval.cmd.PreRunE = preRunETestImpl
 
@@ -431,7 +431,7 @@ func TestInterceptConfigsWithBadPermissions(t *testing.T) {
 	if err := os.Mkdir(subDir, 0o600); err != nil {
 		t.Fatalf("Failed to create sub directory: %v", err)
 	}
-	cmd := server.StartCmd(nil)
+	cmd := server.StartCmd[servertypes.Application](nil)
 	cmd.PersistentFlags().String(flags.FlagHome, "/foobar", "")
 	if err := cmd.PersistentFlags().Set(flags.FlagHome, subDir); err != nil {
 		t.Fatalf("Could not set home flag [%T] %v", err, err)
@@ -471,7 +471,7 @@ func TestEmptyMinGasPrices(t *testing.T) {
 	require.NoError(t, err)
 
 	// Run StartCmd.
-	cmd = server.StartCmd(nil)
+	cmd = server.StartCmd[servertypes.Application](nil)
 	cmd.PersistentFlags().String(flags.FlagHome, tempDir, "")
 	cmd.PreRunE = func(cmd *cobra.Command, _ []string) error {
 		ctx, err := server.InterceptConfigsAndCreateContext(cmd, "", nil, cmtcfg.DefaultConfig())

Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ import (`
`7`	`7`	`)`
`8`	`8`
`9`	`9`	`// Cmd returns the snapshots group command`
`10`		`-func Cmd(appCreator servertypes.AppCreator) *cobra.Command {`
	`10`	`+func Cmd[T servertypes.Application](appCreator servertypes.AppCreator[T]) *cobra.Command {`
`11`	`11`	`cmd := &cobra.Command{`
`12`	`12`	`Use: "snapshots",`
`13`	`13`	`Short: "Manage local snapshots",`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@ import (`
`10`	`10`	`)`
`11`	`11`
`12`	`12`	`// ExportSnapshotCmd returns a command to take a snapshot of the application state`
`13`		`-func ExportSnapshotCmd(appCreator servertypes.AppCreator) *cobra.Command {`
	`13`	`+func ExportSnapshotCmd[T servertypes.Application](appCreator servertypes.AppCreator[T]) *cobra.Command {`
`14`	`14`	`cmd := &cobra.Command{`
`15`	`15`	`Use: "export",`
`16`	`16`	`Short: "Export app state to snapshot store",`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ import (`
`14`	`14`	`)`
`15`	`15`
`16`	`16`	`// RestoreSnapshotCmd returns a command to restore a snapshot`
`17`		`-func RestoreSnapshotCmd(appCreator servertypes.AppCreator) *cobra.Command {`
	`17`	`+func RestoreSnapshotCmd[T servertypes.Application](appCreator servertypes.AppCreator[T]) *cobra.Command {`
`18`	`18`	`cmd := &cobra.Command{`
`19`	`19`	`Use: "restore <height> <format>",`
`20`	`20`	`Short: "Restore app state from local snapshot",`
Original file line number	Diff line number	Diff line change
`@@ -356,7 +356,7 @@ func parseOptionalHeight(heightStr string) (*int64, error) {`
`356`	`356`	`return &tmp, nil`
`357`	`357`	`}`
`358`	`358`
`359`		`-func BootstrapStateCmd(appCreator types.AppCreator) *cobra.Command {`
	`359`	`+func BootstrapStateCmd[T types.Application](appCreator types.AppCreator[T]) *cobra.Command {`
`360`	`360`	`cmd := &cobra.Command{`
`361`	`361`	`Use: "bootstrap-state",`
`362`	`362`	`Short: "Bootstrap CometBFT state at an arbitrary block height using a light client",`