refactor(stf): rm RunWithCtx

core/header/service.go

@@ -13,6 +13,8 @@ type Service interface {
 	HeaderInfo(context.Context) Info
 }
 
+type HeaderServiceFactory func() Service
+
 // Info defines a struct that contains information about the header
 type Info struct {
 	Height  int64     // Height returns the height of the block
@@ -35,7 +37,7 @@ func (i *Info) Bytes() ([]byte, error) {
 
 	// Encode Hash
 	if len(i.Hash) != hashSize {
-		return nil, errors.New("invalid hash size")
+		return nil, errors.New("invalid Hash size")
 	}
 
 	buf = append(buf, i.Hash...)
@@ -47,7 +49,7 @@ func (i *Info) Bytes() ([]byte, error) {
 
 	// Encode AppHash
 	if len(i.AppHash) != hashSize {
-		return nil, errors.New("invalid hash size")
+		return nil, errors.New("invalid AppHash size")
 	}
 	buf = append(buf, i.AppHash...)
 

core/store/service.go

@@ -10,6 +10,8 @@ type KVStoreService interface {
 	OpenKVStore(context.Context) KVStore
 }
 
+type KVStoreServiceFactory func([]byte) KVStoreService
+
 // MemoryStoreService represents a unique, non-forgeable handle to a memory-backed
 // KVStore. It should be provided as a module-scoped dependency by the runtime
 // module being used to build the app.
@@ -18,6 +20,8 @@ type MemoryStoreService interface {
 	OpenMemoryStore(context.Context) KVStore
 }
 
+type MemoryStoreServiceFactory func([]byte) MemoryStoreService
+
 // TransientStoreService represents a unique, non-forgeable handle to a memory-backed
 // KVStore which is reset at the start of every block. It should be provided as
 // a module-scoped dependency by the runtime module being used to build the app.

runtime/v2/builder.go

@@ -3,6 +3,7 @@ package runtime
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"path/filepath"
@@ -45,7 +46,13 @@ func (a *AppBuilder[T]) RegisterModules(modules map[string]appmodulev2.AppModule
 		// if a (legacy) module implements the HasName interface, check that the name matches
 		if mod, ok := appModule.(interface{ Name() string }); ok {
 			if name != mod.Name() {
-				a.app.logger.Warn(fmt.Sprintf("module name %q does not match name returned by HasName: %q", name, mod.Name()))
+				a.app.logger.Warn(
+					fmt.Sprintf(
+						"module name %q does not match name returned by HasName: %q",
+						name,
+						mod.Name(),
+					),
+				)
 			}
 		}
 
@@ -157,20 +164,38 @@ func (a *AppBuilder[T]) Build(opts ...AppBuilderOption[T]) (*App[T], error) {
 		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) error {
+		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 fmt.Errorf("failed to read import state: %w", err)
+				return nil, fmt.Errorf("failed to read import state: %w", err)
 			}
-			var genesisState map[string]json.RawMessage
-			if err = json.Unmarshal(bz, &genesisState); err != nil {
-				return err
+			var genesisJSON map[string]json.RawMessage
+			if err = json.Unmarshal(bz, &genesisJSON); err != nil {
+				return nil, err
 			}
-			if err = a.app.moduleManager.InitGenesisJSON(ctx, genesisState, txHandler); err != nil {
-				return fmt.Errorf("failed to init genesis: %w", err)
+
+			v, zeroState, err := a.app.db.StateLatest()
+			if err != nil {
+				return nil, fmt.Errorf("unable to get latest state: %w", err)
 			}
-			return nil
+			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 := makeGenesisContext(a.branch(zeroState))
+			genesisState, err := genesisCtx.Run(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
 		},
 		ExportGenesis: func(ctx context.Context, version uint64) ([]byte, error) {
 			genesisJson, err := a.app.moduleManager.ExportGenesisForModules(ctx)
@@ -200,23 +225,29 @@ func (a *AppBuilder[T]) Build(opts ...AppBuilderOption[T]) (*App[T], error) {
 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] {
+func AppBuilderWithBranch[T transaction.Tx](
+	branch func(state store.ReaderMap) store.WriterMap,
+) AppBuilderOption[T] {
 	return func(a *AppBuilder[T]) {
 		a.branch = branch
 	}
 }
 
 // AppBuilderWithTxValidator sets the tx validator for the app.
 // It overrides all default tx validators defined by modules.
-func AppBuilderWithTxValidator[T transaction.Tx](txValidators func(ctx context.Context, tx T) error) AppBuilderOption[T] {
+func AppBuilderWithTxValidator[T transaction.Tx](
+	txValidators func(ctx context.Context, tx T) error,
+) AppBuilderOption[T] {
 	return func(a *AppBuilder[T]) {
 		a.txValidator = txValidators
 	}
 }
 
 // AppBuilderWithPostTxExec sets logic that will be executed after each transaction.
 // When not provided, a no-op function will be used.
-func AppBuilderWithPostTxExec[T transaction.Tx](postTxExec func(ctx context.Context, tx T, success bool) error) AppBuilderOption[T] {
+func AppBuilderWithPostTxExec[T transaction.Tx](
+	postTxExec func(ctx context.Context, tx T, success bool) error,
+) AppBuilderOption[T] {
 	return func(a *AppBuilder[T]) {
 		a.postTxExec = postTxExec
 	}

runtime/v2/genesis.go

@@ -0,0 +1,114 @@
+package runtime
+
+import (
+	"context"
+	"fmt"
+
+	"cosmossdk.io/core/header"
+	"cosmossdk.io/core/store"
+)
+
+var _ store.KVStoreService = (*GenesisKVStoreServie)(nil)
+
+type genesisContextKeyType struct{}
+
+var genesisContextKey = genesisContextKeyType{}
+
+type genesisContext struct {
+	state  store.WriterMap
+	didRun bool
+}
+
+func makeGenesisContext(state store.WriterMap) genesisContext {
+	return genesisContext{
+		state: state,
+	}
+}
+
+func (g *genesisContext) Run(
+	ctx context.Context,
+	fn func(ctx context.Context) error,
+) (store.WriterMap, error) {
+	ctx = context.WithValue(ctx, genesisContextKey, g)
+	err := fn(ctx)
+	if err != nil {
+		return nil, err
+	}
+	g.didRun = true
+	return g.state, nil
+}
+
+type GenesisKVStoreServie struct {
+	genesisCapable bool
+	actor          []byte
+	execution      store.KVStoreService
+}
+
+func NewGenesisKVService(
+	actor []byte,
+	execution store.KVStoreService,
+) *GenesisKVStoreServie {
+	return &GenesisKVStoreServie{
+		genesisCapable: true,
+		actor:          actor,
+		execution:      execution,
+	}
+}
+
+// OpenKVStore implements store.KVStoreService.
+func (g *GenesisKVStoreServie) OpenKVStore(ctx context.Context) store.KVStore {
+	if !g.genesisCapable {
+		return g.execution.OpenKVStore(ctx)
+	}
+	v := ctx.Value(genesisContextKey)
+	if v == nil {
+		return g.execution.OpenKVStore(ctx)
+	}
+	genCtx, ok := v.(*genesisContext)
+	if !ok {
+		panic(fmt.Errorf("unexpected genesis context type: %T", v))
+	}
+	if genCtx.didRun {
+		g.genesisCapable = false
+		return g.execution.OpenKVStore(ctx)
+	}
+	state, err := genCtx.state.GetWriter(g.actor)
+	if err != nil {
+		panic(err)
+	}
+	return state
+}
+
+type GenesisHeaderService struct {
+	genesisCapable   bool
+	executionService header.Service
+}
+
+// HeaderInfo implements header.Service.
+func (g *GenesisHeaderService) HeaderInfo(ctx context.Context) header.Info {
+	if !g.genesisCapable {
+		return g.executionService.HeaderInfo(ctx)
+	}
+	v := ctx.Value(genesisContextKey)
+	if v == nil {
+		return g.executionService.HeaderInfo(ctx)
+	}
+	genCtx, ok := v.(*genesisContext)
+	if !ok {
+		panic(fmt.Errorf("unexpected genesis context type: %T", v))
+	}
+	if genCtx.didRun {
+		g.genesisCapable = false
+		return g.executionService.HeaderInfo(ctx)
+	}
+	return header.Info{}
+}
+
+func NewGenesisHeaderService(executionService header.Service) *GenesisHeaderService {
+	return &GenesisHeaderService{
+		genesisCapable:   true,
+		executionService: executionService,
+	}
+}
+
+var _ header.Service = (*GenesisHeaderService)(nil)

runtime/v2/go.mod

@@ -5,6 +5,7 @@ go 1.23
 // server v2 integration
 replace (
 	cosmossdk.io/api => ../../api
+	cosmossdk.io/core => ../../core
 	cosmossdk.io/core/testing => ../../core/testing
 	cosmossdk.io/server/v2/appmanager => ../../server/v2/appmanager
 	cosmossdk.io/server/v2/stf => ../../server/v2/stf

runtime/v2/module.go

@@ -16,6 +16,7 @@ import (
 	reflectionv1 "cosmossdk.io/api/cosmos/reflection/v1"
 	appmodulev2 "cosmossdk.io/core/appmodule/v2"
 	"cosmossdk.io/core/comet"
+	"cosmossdk.io/core/header"
 	"cosmossdk.io/core/registry"
 	"cosmossdk.io/core/server"
 	"cosmossdk.io/core/store"
@@ -96,7 +97,6 @@ func init() {
 			ProvideAppBuilder[transaction.Tx],
 			ProvideEnvironment[transaction.Tx],
 			ProvideModuleManager[transaction.Tx],
-			ProvideCometService,
 		),
 		appconfig.Invoke(SetupAppBuilder),
 	)
@@ -176,6 +176,9 @@ func ProvideEnvironment[T transaction.Tx](
 	config *runtimev2.Module,
 	key depinject.ModuleKey,
 	appBuilder *AppBuilder[T],
+	kvFactory store.KVStoreServiceFactory,
+	memFactory store.MemoryStoreServiceFactory,
+	headerFactory header.HeaderServiceFactory,
 ) (
 	appmodulev2.Environment,
 	store.KVStoreService,
@@ -197,19 +200,19 @@ func ProvideEnvironment[T transaction.Tx](
 		}
 
 		registerStoreKey(appBuilder, kvStoreKey)
-		kvService = stf.NewKVStoreService([]byte(kvStoreKey))
+		kvService = kvFactory([]byte(kvStoreKey))
 
 		memStoreKey := fmt.Sprintf("memory:%s", key.Name())
 		registerStoreKey(appBuilder, memStoreKey)
-		memKvService = stf.NewMemoryStoreService([]byte(memStoreKey))
+		memKvService = memFactory([]byte(memStoreKey))
 	}
 
 	env := appmodulev2.Environment{
 		Logger:             logger,
 		BranchService:      stf.BranchService{},
 		EventService:       stf.NewEventService(),
 		GasService:         stf.NewGasMeterService(),
-		HeaderService:      stf.HeaderService{},
+		HeaderService:      headerFactory(),
 		QueryRouterService: stf.NewQueryRouterService(),
 		MsgRouterService:   stf.NewMsgRouterService([]byte(key.Name())),
 		TransactionService: services.NewContextAwareTransactionService(),
@@ -220,8 +223,8 @@ func ProvideEnvironment[T transaction.Tx](
 	return env, kvService, memKvService
 }
 
-func registerStoreKey[T transaction.Tx](wrapper *AppBuilder[T], key string) {
-	wrapper.app.storeKeys = append(wrapper.app.storeKeys, key)
+func registerStoreKey[T transaction.Tx](builder *AppBuilder[T], key string) {
+	builder.app.storeKeys = append(builder.app.storeKeys, key)
 }
 
 func storeKeyOverride(config *runtimev2.Module, moduleName string) *runtimev2.StoreKeyConfig {
@@ -234,6 +237,24 @@ func storeKeyOverride(config *runtimev2.Module, moduleName string) *runtimev2.St
 	return nil
 }
 
-func ProvideCometService() comet.Service {
-	return &services.ContextAwareCometInfoService{}
+func DefaultServiceBindings() depinject.Config {
+	var (
+		kvServiceFactory store.KVStoreServiceFactory = func(actor []byte) store.KVStoreService {
+			return NewGenesisKVService(
+				actor,
+				stf.NewKVStoreService(actor),
+			)
+		}
+		memStoreServiceFactory store.MemoryStoreServiceFactory = stf.NewMemoryStoreService
+		headerServiceFactory   header.HeaderServiceFactory     = func() header.Service {
+			return NewGenesisHeaderService(stf.HeaderService{})
+		}
+		cometService comet.Service = &services.ContextAwareCometInfoService{}
+	)
+	return depinject.Supply(
+		kvServiceFactory,
+		memStoreServiceFactory,
+		cometService,
+		headerServiceFactory,
+	)
 }