From dd2369daf2dce11d4d234ee11590a03e57812096 Mon Sep 17 00:00:00 2001
From: Matt Kocubinski <mkocubinski@gmail.com>
Date: Thu, 10 Oct 2024 09:29:40 -0500
Subject: [PATCH] refactor(runtime/v2): use AppI.GetStore() to fetch an
 initialized RootStore (#22205)

---
 runtime/v2/module.go             |  3 +++
 runtime/v2/store.go              | 22 ----------------------
 server/v2/cometbft/server.go     | 16 +---------------
 server/v2/server_test.go         | 15 +++------------
 server/v2/store/server.go        | 23 +++++++----------------
 server/v2/types.go               |  2 ++
 simapp/v2/app_di.go              |  3 +--
 simapp/v2/app_test.go            |  2 --
 simapp/v2/simdv2/cmd/commands.go |  7 ++-----
 simapp/v2/simdv2/cmd/root_di.go  |  5 +----
 simapp/v2/simdv2/cmd/testnet.go  | 13 +++++--------
 11 files changed, 25 insertions(+), 86 deletions(-)

diff --git a/runtime/v2/module.go b/runtime/v2/module.go
index 9637871c6aba..dcde5ae48772 100644
--- a/runtime/v2/module.go
+++ b/runtime/v2/module.go
@@ -229,6 +229,7 @@ func ProvideEnvironment(
 // - header.Service
 // - comet.Service
 // - event.Service
+// - store/v2/root.Builder
 //
 // They are all required.  For most use cases these default services bindings should be sufficient.
 // Power users (or tests) may wish to provide their own services bindings, in which case they must
@@ -244,11 +245,13 @@ func DefaultServiceBindings() depinject.Config {
 		cometService  comet.Service = &services.ContextAwareCometInfoService{}
 		headerService               = services.NewGenesisHeaderService(stf.HeaderService{})
 		eventService                = services.NewGenesisEventService(stf.NewEventService())
+		storeBuilder                = root.NewBuilder()
 	)
 	return depinject.Supply(
 		kvServiceFactory,
 		headerService,
 		cometService,
 		eventService,
+		storeBuilder,
 	)
 }
diff --git a/runtime/v2/store.go b/runtime/v2/store.go
index 7f7b2135c809..40912ea41f48 100644
--- a/runtime/v2/store.go
+++ b/runtime/v2/store.go
@@ -3,35 +3,13 @@ package runtime
 import (
 	"errors"
 	"fmt"
-	"sync"
 
 	"cosmossdk.io/core/store"
 	"cosmossdk.io/server/v2/stf"
 	storev2 "cosmossdk.io/store/v2"
 	"cosmossdk.io/store/v2/proof"
-	"cosmossdk.io/store/v2/root"
 )
 
-var (
-	storeBuilderSingleton     root.Builder
-	storeBuilderSingletonOnce sync.Once
-)
-
-// ProvideSingletonScopedStoreBuilder returns a store builder that is a singleton
-// in the scope of the process lifetime.
-func ProvideSingletonScopedStoreBuilder() root.Builder {
-	storeBuilderSingletonOnce.Do(func() {
-		storeBuilderSingleton = root.NewBuilder()
-	})
-	return storeBuilderSingleton
-}
-
-// ResetSingletonScopedStoreBuilder resets the singleton store builder.  Applications
-// should not ever need to call this, but it may be useful in tests.
-func ResetSingletonScopedStoreBuilder() {
-	storeBuilderSingletonOnce = sync.Once{}
-}
-
 // NewKVStoreService creates a new KVStoreService.
 // This wrapper is kept for backwards compatibility.
 // When migrating from runtime to runtime/v2, use runtimev2.NewKVStoreService(storeKey.Name()) instead of runtime.NewKVStoreService(storeKey).
diff --git a/server/v2/cometbft/server.go b/server/v2/cometbft/server.go
index 7fa1a94d8ee6..e5076b897073 100644
--- a/server/v2/cometbft/server.go
+++ b/server/v2/cometbft/server.go
@@ -8,9 +8,6 @@ import (
 	"os"
 	"path/filepath"
 
-	"cosmossdk.io/server/v2/store"
-	"cosmossdk.io/store/v2/root"
-
 	abciserver "github.com/cometbft/cometbft/abci/server"
 	cmtcmd "github.com/cometbft/cometbft/cmd/cometbft/commands"
 	cmtcfg "github.com/cometbft/cometbft/config"
@@ -45,7 +42,6 @@ type CometBFTServer[T transaction.Tx] struct {
 
 	initTxCodec   transaction.Codec[T]
 	logger        log.Logger
-	storeBuilder  root.Builder
 	serverOptions ServerOptions[T]
 	config        Config
 	cfgOptions    []CfgOption
@@ -53,13 +49,11 @@ type CometBFTServer[T transaction.Tx] struct {
 
 func New[T transaction.Tx](
 	txCodec transaction.Codec[T],
-	storeBuilder root.Builder,
 	serverOptions ServerOptions[T],
 	cfgOptions ...CfgOption,
 ) *CometBFTServer[T] {
 	return &CometBFTServer[T]{
 		initTxCodec:   txCodec,
-		storeBuilder:  storeBuilder,
 		serverOptions: serverOptions,
 		cfgOptions:    cfgOptions,
 	}
@@ -106,16 +100,8 @@ func (s *CometBFTServer[T]) Init(appI serverv2.AppI[T], cfg map[string]any, logg
 		indexEvents[e] = struct{}{}
 	}
 
-	storeCfg, err := store.UnmarshalConfig(cfg)
-	if err != nil {
-		return err
-	}
-	rs, err := s.storeBuilder.Build(logger, storeCfg)
-	if err != nil {
-		return err
-	}
-
 	s.logger = logger.With(log.ModuleKey, s.Name())
+	rs := appI.GetStore()
 	consensus := NewConsensus(
 		s.logger,
 		appI.Name(),
diff --git a/server/v2/server_test.go b/server/v2/server_test.go
index c216824fb652..b67fd209e1bd 100644
--- a/server/v2/server_test.go
+++ b/server/v2/server_test.go
@@ -20,7 +20,6 @@ import (
 	"cosmossdk.io/server/v2/appmanager"
 	"cosmossdk.io/server/v2/store"
 	storev2 "cosmossdk.io/store/v2"
-	"cosmossdk.io/store/v2/root"
 )
 
 type mockInterfaceRegistry struct{}
@@ -50,18 +49,10 @@ func (*mockApp[T]) InterfaceRegistry() coreserver.InterfaceRegistry {
 	return &mockInterfaceRegistry{}
 }
 
-var _ root.Builder = &mockStoreBuilder{}
-
-type mockStoreBuilder struct{}
-
-func (m mockStoreBuilder) Build(logger log.Logger, config *root.Config) (storev2.RootStore, error) {
-	return nil, nil
+func (*mockApp[T]) GetStore() storev2.RootStore {
+	return nil
 }
 
-func (m mockStoreBuilder) RegisterKey(string) {}
-
-func (m mockStoreBuilder) Get() storev2.RootStore { return nil }
-
 func TestServer(t *testing.T) {
 	currentDir, err := os.Getwd()
 	require.NoError(t, err)
@@ -78,7 +69,7 @@ func TestServer(t *testing.T) {
 	err = grpcServer.Init(&mockApp[transaction.Tx]{}, cfg, logger)
 	require.NoError(t, err)
 
-	storeServer := store.New[transaction.Tx](&mockStoreBuilder{})
+	storeServer := store.New[transaction.Tx]()
 	err = storeServer.Init(&mockApp[transaction.Tx]{}, cfg, logger)
 	require.NoError(t, err)
 
diff --git a/server/v2/store/server.go b/server/v2/store/server.go
index a793a1912dba..20ed3b15b0ea 100644
--- a/server/v2/store/server.go
+++ b/server/v2/store/server.go
@@ -24,26 +24,17 @@ const ServerName = "store"
 // Server manages store config and contains prune & snapshot commands
 type Server[T transaction.Tx] struct {
 	config  *root.Config
-	builder root.Builder
 	backend storev2.Backend
 }
 
-func New[T transaction.Tx](builder root.Builder) *Server[T] {
-	return &Server[T]{builder: builder}
+func New[T transaction.Tx]() *Server[T] {
+	return &Server[T]{}
 }
 
-func (s *Server[T]) Init(_ serverv2.AppI[T], cfg map[string]any, log log.Logger) error {
-	var err error
-	s.config, err = UnmarshalConfig(cfg)
-	if err != nil {
-		return fmt.Errorf("failed to unmarshal config: %w", err)
-	}
-	s.backend, err = s.builder.Build(log, s.config)
-	if err != nil {
-		return fmt.Errorf("failed to create store backend: %w", err)
-	}
-
-	return nil
+func (s *Server[T]) Init(app serverv2.AppI[T], v map[string]any, _ log.Logger) (err error) {
+	s.backend = app.GetStore()
+	s.config, err = UnmarshalConfig(v)
+	return err
 }
 
 func (s *Server[T]) Name() string {
@@ -90,7 +81,7 @@ func UnmarshalConfig(cfg map[string]any) (*root.Config, error) {
 		Options: root.DefaultStoreOptions(),
 	}
 	if err := serverv2.UnmarshalSubConfig(cfg, ServerName, config); err != nil {
-		return nil, fmt.Errorf("failed to unmarshal config: %w", err)
+		return nil, fmt.Errorf("failed to unmarshal store config: %w", err)
 	}
 	home := cfg[serverv2.FlagHome]
 	if home != nil {
diff --git a/server/v2/types.go b/server/v2/types.go
index a28385922178..8bd04d3221af 100644
--- a/server/v2/types.go
+++ b/server/v2/types.go
@@ -8,6 +8,7 @@ import (
 	"cosmossdk.io/core/transaction"
 	"cosmossdk.io/log"
 	"cosmossdk.io/server/v2/appmanager"
+	"cosmossdk.io/store/v2"
 )
 
 type AppCreator[T transaction.Tx] func(log.Logger, *viper.Viper) AppI[T]
@@ -17,4 +18,5 @@ type AppI[T transaction.Tx] interface {
 	InterfaceRegistry() server.InterfaceRegistry
 	GetAppManager() *appmanager.AppManager[T]
 	GetQueryHandlers() map[string]appmodulev2.Handler
+	GetStore() store.RootStore
 }
diff --git a/simapp/v2/app_di.go b/simapp/v2/app_di.go
index 29f62defcde6..b1dfc26f99ff 100644
--- a/simapp/v2/app_di.go
+++ b/simapp/v2/app_di.go
@@ -66,7 +66,6 @@ func AppConfig() depinject.Config {
 			codec.ProvideAddressCodec,
 			codec.ProvideProtoCodec,
 			codec.ProvideLegacyAmino,
-			runtime.ProvideSingletonScopedStoreBuilder,
 		),
 		depinject.Invoke(
 			std.RegisterInterfaces,
@@ -83,8 +82,8 @@ func NewSimApp[T transaction.Tx](
 	var (
 		app          = &SimApp[T]{}
 		appBuilder   *runtime.AppBuilder[T]
-		storeBuilder root.Builder
 		err          error
+		storeBuilder root.Builder
 
 		// merge the AppConfig and other configuration in one config
 		appConfig = depinject.Configs(
diff --git a/simapp/v2/app_test.go b/simapp/v2/app_test.go
index 8c705565956c..a5d7f352b4d9 100644
--- a/simapp/v2/app_test.go
+++ b/simapp/v2/app_test.go
@@ -18,7 +18,6 @@ import (
 	"cosmossdk.io/core/transaction"
 	"cosmossdk.io/log"
 	sdkmath "cosmossdk.io/math"
-	"cosmossdk.io/runtime/v2"
 	serverv2 "cosmossdk.io/server/v2"
 	serverv2store "cosmossdk.io/server/v2/store"
 	"cosmossdk.io/store/v2/db"
@@ -40,7 +39,6 @@ func NewTestApp(t *testing.T) (*SimApp[transaction.Tx], context.Context) {
 	vp.Set(serverv2store.FlagAppDBBackend, string(db.DBTypeGoLevelDB))
 	vp.Set(serverv2.FlagHome, t.TempDir())
 
-	runtime.ResetSingletonScopedStoreBuilder()
 	app := NewSimApp[transaction.Tx](logger, vp)
 	genesis := app.ModuleManager().DefaultGenesis()
 
diff --git a/simapp/v2/simdv2/cmd/commands.go b/simapp/v2/simdv2/cmd/commands.go
index 6cc233cc9b19..3fcbed114d9e 100644
--- a/simapp/v2/simdv2/cmd/commands.go
+++ b/simapp/v2/simdv2/cmd/commands.go
@@ -19,7 +19,6 @@ import (
 	"cosmossdk.io/server/v2/cometbft"
 	serverstore "cosmossdk.io/server/v2/store"
 	"cosmossdk.io/simapp/v2"
-	"cosmossdk.io/store/v2/root"
 	confixcmd "cosmossdk.io/tools/confix/cmd"
 
 	"github.com/cosmos/cosmos-sdk/client"
@@ -44,7 +43,6 @@ func newApp[T transaction.Tx](logger log.Logger, viper *viper.Viper) serverv2.Ap
 func initRootCmd[T transaction.Tx](
 	rootCmd *cobra.Command,
 	txConfig client.TxConfig,
-	storeBuilder root.Builder,
 	moduleManager *runtimev2.MM[T],
 ) {
 	cfg := sdk.GetConfig()
@@ -54,7 +52,7 @@ func initRootCmd[T transaction.Tx](
 		genutilcli.InitCmd(moduleManager),
 		debug.Cmd(),
 		confixcmd.ConfigCommand(),
-		NewTestnetCmd(storeBuilder, moduleManager),
+		NewTestnetCmd(moduleManager),
 	)
 
 	logger, err := serverv2.NewLogger(viper.New(), rootCmd.OutOrStdout())
@@ -79,12 +77,11 @@ func initRootCmd[T transaction.Tx](
 		initServerConfig(),
 		cometbft.New(
 			&genericTxDecoder[T]{txConfig},
-			storeBuilder,
 			initCometOptions[T](),
 			initCometConfig(),
 		),
 		grpc.New[T](),
-		serverstore.New[T](storeBuilder),
+		serverstore.New[T](),
 		telemetry.New[T](),
 	); err != nil {
 		panic(err)
diff --git a/simapp/v2/simdv2/cmd/root_di.go b/simapp/v2/simdv2/cmd/root_di.go
index 30291b4ecdb3..6b13b7201342 100644
--- a/simapp/v2/simdv2/cmd/root_di.go
+++ b/simapp/v2/simdv2/cmd/root_di.go
@@ -14,7 +14,6 @@ import (
 	"cosmossdk.io/log"
 	"cosmossdk.io/runtime/v2"
 	"cosmossdk.io/simapp/v2"
-	"cosmossdk.io/store/v2/root"
 
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/config"
@@ -32,7 +31,6 @@ func NewRootCmd[T transaction.Tx]() *cobra.Command {
 		autoCliOpts   autocli.AppOptions
 		moduleManager *runtime.MM[T]
 		clientCtx     client.Context
-		storeBuilder  root.Builder
 	)
 
 	if err := depinject.Inject(
@@ -41,7 +39,6 @@ func NewRootCmd[T transaction.Tx]() *cobra.Command {
 			depinject.Provide(ProvideClientContext),
 			depinject.Supply(log.NewNopLogger()),
 		),
-		&storeBuilder,
 		&autoCliOpts,
 		&moduleManager,
 		&clientCtx,
@@ -74,7 +71,7 @@ func NewRootCmd[T transaction.Tx]() *cobra.Command {
 		},
 	}
 
-	initRootCmd(rootCmd, clientCtx.TxConfig, storeBuilder, moduleManager)
+	initRootCmd(rootCmd, clientCtx.TxConfig, moduleManager)
 
 	nodeCmds := nodeservice.NewNodeCommands()
 	autoCliOpts.ModuleOptions = make(map[string]*autocliv1.ModuleOptions)
diff --git a/simapp/v2/simdv2/cmd/testnet.go b/simapp/v2/simdv2/cmd/testnet.go
index 4e24f7083654..f030db8e5956 100644
--- a/simapp/v2/simdv2/cmd/testnet.go
+++ b/simapp/v2/simdv2/cmd/testnet.go
@@ -23,7 +23,6 @@ import (
 	"cosmossdk.io/server/v2/api/grpc"
 	"cosmossdk.io/server/v2/cometbft"
 	"cosmossdk.io/server/v2/store"
-	"cosmossdk.io/store/v2/root"
 	banktypes "cosmossdk.io/x/bank/types"
 	bankv2types "cosmossdk.io/x/bank/v2/types"
 	stakingtypes "cosmossdk.io/x/staking/types"
@@ -88,7 +87,7 @@ func addTestnetFlagsToCmd(cmd *cobra.Command) {
 
 // NewTestnetCmd creates a root testnet command with subcommands to run an in-process testnet or initialize
 // validator configuration files for running a multi-validator testnet in a separate process
-func NewTestnetCmd[T transaction.Tx](sb root.Builder, mm *runtimev2.MM[T]) *cobra.Command {
+func NewTestnetCmd[T transaction.Tx](mm *runtimev2.MM[T]) *cobra.Command {
 	testnetCmd := &cobra.Command{
 		Use:                        "testnet",
 		Short:                      "subcommands for starting or configuring local testnets",
@@ -97,13 +96,13 @@ func NewTestnetCmd[T transaction.Tx](sb root.Builder, mm *runtimev2.MM[T]) *cobr
 		RunE:                       client.ValidateCmd,
 	}
 
-	testnetCmd.AddCommand(testnetInitFilesCmd(sb, mm))
+	testnetCmd.AddCommand(testnetInitFilesCmd(mm))
 
 	return testnetCmd
 }
 
 // testnetInitFilesCmd returns a cmd to initialize all files for CometBFT testnet and application
-func testnetInitFilesCmd[T transaction.Tx](sb root.Builder, mm *runtimev2.MM[T]) *cobra.Command {
+func testnetInitFilesCmd[T transaction.Tx](mm *runtimev2.MM[T]) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "init-files",
 		Short: "Initialize config directories & files for a multi-validator testnet running locally via separate processes (e.g. Docker Compose or similar)",
@@ -144,7 +143,7 @@ Example:
 				return err
 			}
 
-			return initTestnetFiles(clientCtx, sb, cmd, config, mm, args)
+			return initTestnetFiles(clientCtx, cmd, config, mm, args)
 		},
 	}
 
@@ -166,7 +165,6 @@ const nodeDirPerm = 0o755
 // initTestnetFiles initializes testnet files for a testnet to be run in a separate process
 func initTestnetFiles[T transaction.Tx](
 	clientCtx client.Context,
-	sb root.Builder,
 	cmd *cobra.Command,
 	nodeConfig *cmtconfig.Config,
 	mm *runtimev2.MM[T],
@@ -341,11 +339,10 @@ func initTestnetFiles[T transaction.Tx](
 		// Write server config
 		cometServer := cometbft.New[T](
 			&genericTxDecoder[T]{clientCtx.TxConfig},
-			sb,
 			cometbft.ServerOptions[T]{},
 			cometbft.OverwriteDefaultConfigTomlConfig(nodeConfig),
 		)
-		storeServer := store.New[T](sb)
+		storeServer := store.New[T]()
 		grpcServer := grpc.New[T](grpc.OverwriteDefaultConfig(grpcConfig))
 		server := serverv2.NewServer[T](log.NewNopLogger(), serverCfg, cometServer, grpcServer, storeServer)
 		err = server.WriteConfig(filepath.Join(nodeDir, "config"))