refactor(runtime/v2): simplify app manager

Author: julienrbrt

runtime/v2/app.go

@@ -1,13 +1,19 @@
 package runtime
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
 
 	runtimev2 "cosmossdk.io/api/cosmos/app/runtime/v2"
 	appmodulev2 "cosmossdk.io/core/appmodule/v2"
 	"cosmossdk.io/core/registry"
+	"cosmossdk.io/core/store"
 	"cosmossdk.io/core/transaction"
 	"cosmossdk.io/log"
+	"cosmossdk.io/runtime/v2/services"
 	"cosmossdk.io/schema/decoding"
 	"cosmossdk.io/server/v2/appmanager"
 	"cosmossdk.io/server/v2/stf"
@@ -23,26 +29,80 @@ import (
 // done declaratively with an app config and the rest of it is done the old way.
 // See simapp/app_v2.go for an example of this setup.
 type App[T transaction.Tx] struct {
-	*appmanager.AppManager[T]
+	// app configuration
+	logger log.Logger
+	config *runtimev2.Module
 
-	// app manager dependencies
+	// state
 	stf                *stf.STF[T]
 	msgRouterBuilder   *stf.MsgRouterBuilder
 	queryRouterBuilder *stf.MsgRouterBuilder
+	appm               *appmanager.AppManager[T]
+	branch             func(state store.ReaderMap) store.WriterMap
 	db                 Store
+	storeLoader        StoreLoader
 
-	// app configuration
-	logger log.Logger
-	config *runtimev2.Module
-
+	// modules
 	interfaceRegistrar registry.InterfaceRegistrar
 	amino              registry.AminoRegistrar
 	moduleManager      *MM[T]
+	queryHandlers      map[string]appmodulev2.Handler // queryHandlers defines the query handlers
+}
+
+// initGenesis initializes the genesis state of the application.
+func (a *App[T]) initGenesis(ctx context.Context, src io.Reader, txHandler func(json.RawMessage) error) (store.WriterMap, error) {
+	// this implementation assumes that the state is a JSON object
+	bz, err := io.ReadAll(src)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read import state: %w", err)
+	}
+	var genesisJSON map[string]json.RawMessage
+	if err = json.Unmarshal(bz, &genesisJSON); err != nil {
+		return nil, err
+	}
+
+	v, zeroState, err := a.db.StateLatest()
+	if err != nil {
+		return nil, fmt.Errorf("unable to get latest state: %w", err)
+	}
+	if v != 0 { // TODO: genesis state may be > 0, we need to set version on store
+		return nil, errors.New("cannot init genesis on non-zero state")
+	}
+	genesisCtx := services.NewGenesisContext(a.branch(zeroState))
+	genesisState, err := genesisCtx.Mutate(ctx, func(ctx context.Context) error {
+		err = a.moduleManager.InitGenesisJSON(ctx, genesisJSON, txHandler)
+		if err != nil {
+			return fmt.Errorf("failed to init genesis: %w", err)
+		}
+		return nil
+	})
+
+	return genesisState, err
+}
 
-	// QueryHandlers defines the query handlers
-	QueryHandlers map[string]appmodulev2.Handler
+// exportGenesis exports the genesis state of the application.
+func (a *App[T]) exportGenesis(ctx context.Context, version uint64) ([]byte, error) {
+	state, err := a.db.StateAt(version)
+	if err != nil {
+		return nil, fmt.Errorf("unable to get state at given version: %w", err)
+	}
+
+	genesisJson, err := a.moduleManager.ExportGenesisForModules(
+		ctx,
+		func() store.WriterMap {
+			return a.branch(state)
+		},
+	)
+	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)
+	}
 
-	storeLoader StoreLoader
+	return bz, nil
 }
 
 // Name returns the app name.
@@ -85,24 +145,26 @@ 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
+// AppManager returns the app's appamanger
+func (a *App[T]) AppManager() *appmanager.AppManager[T] {
+	return a.appm
 }
 
-func (a *App[T]) GetAppManager() *appmanager.AppManager[T] {
-	return a.AppManager
+// GetQueryHandlers returns the query handlers.
+func (a *App[T]) QueryHandlers() map[string]appmodulev2.Handler {
+	return a.queryHandlers
 }
 
-func (a *App[T]) GetQueryHandlers() map[string]appmodulev2.Handler {
-	return a.QueryHandlers
-}
-
-// GetSchemaDecoderResolver returns the module schema resolver.
-func (a *App[T]) GetSchemaDecoderResolver() decoding.DecoderResolver {
+// SchemaDecoderResolver returns the module schema resolver.
+func (a *App[T]) SchemaDecoderResolver() decoding.DecoderResolver {
 	moduleSet := map[string]any{}
 	for moduleName, module := range a.moduleManager.Modules() {
 		moduleSet[moduleName] = module
 	}
 	return decoding.ModuleSetDecoderResolver(moduleSet)
 }
+
+// Close is called in start cmd to gracefully cleanup resources.
+func (a *App[T]) Close() error {
+	return nil
+}

runtime/v2/builder.go

@@ -3,15 +3,12 @@ package runtime
 import (
 	"context"
 	"encoding/json"
-	"errors"
 	"fmt"
-	"io"
 
 	"cosmossdk.io/core/appmodule"
 	appmodulev2 "cosmossdk.io/core/appmodule/v2"
 	"cosmossdk.io/core/store"
 	"cosmossdk.io/core/transaction"
-	"cosmossdk.io/runtime/v2/services"
 	"cosmossdk.io/server/v2/appmanager"
 	"cosmossdk.io/server/v2/stf"
 	"cosmossdk.io/server/v2/stf/branch"
@@ -26,7 +23,6 @@ type AppBuilder[T transaction.Tx] struct {
 	storeBuilder root.Builder
 
 	// the following fields are used to overwrite the default
-	branch      func(state store.ReaderMap) store.WriterMap
 	txValidator func(ctx context.Context, tx T) error
 	postTxExec  func(ctx context.Context, tx T, success bool) error
 }
@@ -71,8 +67,8 @@ func (a *AppBuilder[T]) Build(opts ...AppBuilderOption[T]) (*App[T], error) {
 	}
 
 	// default branch
-	if a.branch == nil {
-		a.branch = branch.DefaultNewWriterMap
+	if a.app.branch == nil {
+		a.app.branch = branch.DefaultNewWriterMap
 	}
 
 	// default tx validator
@@ -108,82 +104,28 @@ func (a *AppBuilder[T]) Build(opts ...AppBuilderOption[T]) (*App[T], error) {
 		a.txValidator,
 		valUpdate,
 		a.postTxExec,
-		a.branch,
+		a.app.branch,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create STF: %w", err)
 	}
 	a.app.stf = stf
 
-	appManagerBuilder := appmanager.Builder[T]{
-		STF:                a.app.stf,
-		DB:                 a.app.db,
-		ValidateTxGasLimit: a.app.config.GasConfig.ValidateTxGasLimit,
-		QueryGasLimit:      a.app.config.GasConfig.QueryGasLimit,
-		SimulationGasLimit: a.app.config.GasConfig.SimulationGasLimit,
-		InitGenesis: func(
-			ctx context.Context,
-			src io.Reader,
-			txHandler func(json.RawMessage) error,
-		) (store.WriterMap, error) {
-			// this implementation assumes that the state is a JSON object
-			bz, err := io.ReadAll(src)
-			if err != nil {
-				return nil, fmt.Errorf("failed to read import state: %w", err)
-			}
-			var genesisJSON map[string]json.RawMessage
-			if err = json.Unmarshal(bz, &genesisJSON); err != nil {
-				return nil, err
-			}
-
-			v, zeroState, err := a.app.db.StateLatest()
-			if err != nil {
-				return nil, fmt.Errorf("unable to get latest state: %w", err)
-			}
-			if v != 0 { // TODO: genesis state may be > 0, we need to set version on store
-				return nil, errors.New("cannot init genesis on non-zero state")
-			}
-			genesisCtx := services.NewGenesisContext(a.branch(zeroState))
-			genesisState, err := genesisCtx.Mutate(ctx, func(ctx context.Context) error {
-				err = a.app.moduleManager.InitGenesisJSON(ctx, genesisJSON, txHandler)
-				if err != nil {
-					return fmt.Errorf("failed to init genesis: %w", err)
-				}
-				return nil
-			})
-
-			return genesisState, err
+	appManager, err := appmanager.NewAppManager[T](
+		appmanager.Config{
+			ValidateTxGasLimit: a.app.config.GasConfig.ValidateTxGasLimit,
+			QueryGasLimit:      a.app.config.GasConfig.QueryGasLimit,
+			SimulationGasLimit: a.app.config.GasConfig.SimulationGasLimit,
 		},
-		ExportGenesis: func(ctx context.Context, version uint64) ([]byte, error) {
-			state, err := a.app.db.StateAt(version)
-			if err != nil {
-				return nil, fmt.Errorf("unable to get state at given version: %w", err)
-			}
-
-			genesisJson, err := a.app.moduleManager.ExportGenesisForModules(
-				ctx,
-				func() store.WriterMap {
-					return a.branch(state)
-				},
-			)
-			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.app.db,
+		a.app.stf,
+		a.app.initGenesis,
+		a.app.exportGenesis,
+	)
 	if err != nil {
-		return nil, fmt.Errorf("failed to build app manager: %w", err)
+		return nil, fmt.Errorf("failed to create AppManager: %w", err)
 	}
-	a.app.AppManager = appManager
+	a.app.appm = appManager
 
 	return a.app, nil
 }
@@ -194,7 +136,7 @@ type AppBuilderOption[T transaction.Tx] func(*AppBuilder[T])
 // AppBuilderWithBranch sets a custom branch implementation for the app.
 func AppBuilderWithBranch[T transaction.Tx](branch func(state store.ReaderMap) store.WriterMap) AppBuilderOption[T] {
 	return func(a *AppBuilder[T]) {
-		a.branch = branch
+		a.app.branch = branch
 	}
 }
 

runtime/v2/manager.go

@@ -636,7 +636,7 @@ func registerServices[T transaction.Tx](s appmodulev2.AppModule, app *App[T], re
 
 		// merge maps
 		for path, decoder := range c.queryHandlers {
-			app.QueryHandlers[path] = decoder
+			app.queryHandlers[path] = decoder
 		}
 	}
 
@@ -655,7 +655,7 @@ func registerServices[T transaction.Tx](s appmodulev2.AppModule, app *App[T], re
 		module.RegisterQueryHandlers(&wrapper)
 
 		for path, handler := range wrapper.handlers {
-			app.QueryHandlers[path] = handler
+			app.queryHandlers[path] = handler
 		}
 	}
 

runtime/v2/module.go

@@ -131,7 +131,7 @@ func ProvideAppBuilder[T transaction.Tx](
 		amino:              amino,
 		msgRouterBuilder:   msgRouterBuilder,
 		queryRouterBuilder: stf.NewMsgRouterBuilder(), // TODO dedicated query router
-		QueryHandlers:      map[string]appmodulev2.Handler{},
+		queryHandlers:      map[string]appmodulev2.Handler{},
 		storeLoader:        DefaultStoreLoader,
 	}
 	appBuilder := &AppBuilder[T]{app: app, storeBuilder: storeBuilder}

server/v2/api/grpc/server.go

@@ -57,14 +57,14 @@ func (s *Server[T]) Init(appI serverv2.AppI[T], cfg map[string]any, logger log.L
 			return fmt.Errorf("failed to unmarshal config: %w", err)
 		}
 	}
-	methodsMap := appI.GetQueryHandlers()
+	methodsMap := appI.QueryHandlers()
 
 	grpcSrv := grpc.NewServer(
 		grpc.ForceServerCodec(newProtoCodec(appI.InterfaceRegistry()).GRPCCodec()),
 		grpc.MaxSendMsgSize(serverCfg.MaxSendMsgSize),
 		grpc.MaxRecvMsgSize(serverCfg.MaxRecvMsgSize),
 		grpc.UnknownServiceHandler(
-			makeUnknownServiceHandler(methodsMap, appI.GetAppManager()),
+			makeUnknownServiceHandler(methodsMap, appI.AppManager()),
 		),
 	)
 

server/v2/api/rest/server.go

@@ -45,7 +45,7 @@ func (s *Server[T]) Init(appI serverv2.AppI[T], cfg map[string]any, logger log.L
 	}
 
 	s.router = http.NewServeMux()
-	s.router.Handle("/", NewDefaultHandler(appI.GetAppManager()))
+	s.router.Handle("/", NewDefaultHandler(appI.AppManager()))
 	s.config = serverCfg
 
 	return nil

server/v2/appmanager/appmanager.go

@@ -26,14 +26,34 @@ type Store interface {
 
 // AppManager is a coordinator for all things related to an application
 type AppManager[T transaction.Tx] struct {
+	// Gas limits for validating, querying, and simulating transactions.
 	config Config
-
+	// InitGenesis is a function that initializes the application state from a genesis file.
+	// It takes a context, a source reader for the genesis file, and a transaction handler function.
+	initGenesis InitGenesis
+	// ExportGenesis is a function that exports the application state to a genesis file.
+	// It takes a context and a version number for the genesis file.
+	exportGenesis ExportGenesis
+	// The database for storing application data.
 	db Store
+	// The state transition function for processing transactions.
+	stf StateTransitionFunction[T]
+}
 
-	initGenesis   InitGenesis
-	exportGenesis ExportGenesis
+func NewAppManager[T transaction.Tx](
+	config Config,
+	db Store,
+	stf StateTransitionFunction[T],
+	initGenesisImpl InitGenesis,
+	exportGenesisImpl ExportGenesis,
+) (*AppManager[T], error) {
+	appManager := &AppManager[T]{
+		config: config,
+		db:     db,
+		stf:    stf,
+	}
 
-	stf StateTransitionFunction[T]
+	return appManager, nil
 }
 
 // InitGenesis initializes the genesis state of the application.