refactor(stf): remove RunWithCtx

core/header/service.go

@@ -35,7 +35,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 +47,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,11 @@ type KVStoreService interface {
 	OpenKVStore(context.Context) KVStore
 }
 
+// KVStoreServiceFactory is a function that creates a new KVStoreService.
+// It can be used to override the default KVStoreService bindings for cases
+// where an application must supply a custom stateful backend.
+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.

runtime/v2/builder.go

@@ -3,6 +3,7 @@ package runtime
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"path/filepath"
@@ -12,6 +13,7 @@ import (
 	"cosmossdk.io/core/server"
 	"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"
@@ -157,23 +159,51 @@ 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 := services.NewGenesisContext(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)
+			_, state, err := a.app.db.StateLatest()
+			if err != nil {
+				return nil, fmt.Errorf("unable to get latest state: %w", err)
+			}
+			genesisCtx := services.NewGenesisContext(a.branch(state))
+
+			var genesisJson map[string]json.RawMessage
+			_, err = genesisCtx.Run(ctx, func(ctx context.Context) error {
+				genesisJson, err = a.app.moduleManager.ExportGenesisForModules(ctx)
+				return err
+			})
 			if err != nil {
 				return nil, fmt.Errorf("failed to export genesis: %w", err)
 			}

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/go.sum

@@ -2,8 +2,6 @@ buf.build/gen/go/cometbft/cometbft/protocolbuffers/go v1.34.2-20240701160653-fed
 buf.build/gen/go/cometbft/cometbft/protocolbuffers/go v1.34.2-20240701160653-fedbb9acfd2f.2/go.mod h1:1+3gJj2NvZ1mTLAtHu+lMhOjGgQPiCKCeo+9MBww0Eo=
 buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.34.2-20240130113600-88ef6483f90f.2 h1:b7EEYTUHmWSBEyISHlHvXbJPqtKiHRuUignL1tsHnNQ=
 buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.34.2-20240130113600-88ef6483f90f.2/go.mod h1:HqcXMSa5qnNuakaMUo+hWhF51mKbcrZxGl9Vp5EeJXc=
-cosmossdk.io/core v1.0.0-alpha.2 h1:epU0Xwces4Rgl5bMhHHkXGaGDcyucNGlC/JDH+Suckg=
-cosmossdk.io/core v1.0.0-alpha.2/go.mod h1:abgLjeFLhtuKIYZWSPlVUgQBrKObO7ULV35KYfexE90=
 cosmossdk.io/depinject v1.0.0 h1:dQaTu6+O6askNXO06+jyeUAnF2/ssKwrrszP9t5q050=
 cosmossdk.io/depinject v1.0.0/go.mod h1:zxK/h3HgHoA/eJVtiSsoaRaRA2D5U4cJ5thIG4ssbB8=
 cosmossdk.io/errors/v2 v2.0.0-20240731132947-df72853b3ca5 h1:IQNdY2kB+k+1OM2DvqFG1+UgeU1JzZrWtwuWzI3ZfwA=

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,8 @@ func ProvideEnvironment[T transaction.Tx](
 	config *runtimev2.Module,
 	key depinject.ModuleKey,
 	appBuilder *AppBuilder[T],
+	kvFactory store.KVStoreServiceFactory,
+	headerService header.Service,
 ) (
 	appmodulev2.Environment,
 	store.KVStoreService,
@@ -197,7 +199,7 @@ 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)
@@ -209,7 +211,7 @@ func ProvideEnvironment[T transaction.Tx](
 		BranchService:      stf.BranchService{},
 		EventService:       stf.NewEventService(),
 		GasService:         stf.NewGasMeterService(),
-		HeaderService:      stf.HeaderService{},
+		HeaderService:      headerService,
 		QueryRouterService: stf.NewQueryRouterService(),
 		MsgRouterService:   stf.NewMsgRouterService([]byte(key.Name())),
 		TransactionService: services.NewContextAwareTransactionService(),
@@ -220,8 +222,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 +236,20 @@ 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 services.NewGenesisKVService(
+				actor,
+				stf.NewKVStoreService(actor),
+			)
+		}
+		headerService header.Service = services.NewGenesisHeaderService(stf.HeaderService{})
+		cometService  comet.Service  = &services.ContextAwareCometInfoService{}
+	)
+	return depinject.Supply(
+		kvServiceFactory,
+		headerService,
+		cometService,
+	)
 }

runtime/v2/services/genesis.go

@@ -0,0 +1,90 @@
+package services
+
+import (
+	"context"
+	"fmt"
+
+	"cosmossdk.io/core/header"
+	"cosmossdk.io/core/store"
+)
+
+var _ store.KVStoreService = (*GenesisKVStoreService)(nil)
+
+type genesisContextKeyType struct{}
+
+var genesisContextKey = genesisContextKeyType{}
+
+type genesisContext struct {
+	state store.WriterMap
+}
+
+func NewGenesisContext(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
+	}
+	return g.state, nil
+}
+
+type GenesisKVStoreService struct {
+	actor            []byte
+	executionService store.KVStoreService
+}
+
+func NewGenesisKVService(
+	actor []byte,
+	execution store.KVStoreService,
+) *GenesisKVStoreService {
+	return &GenesisKVStoreService{
+		actor:            actor,
+		executionService: execution,
+	}
+}
+
+// OpenKVStore implements store.KVStoreService.
+func (g *GenesisKVStoreService) OpenKVStore(ctx context.Context) store.KVStore {
+	v := ctx.Value(genesisContextKey)
+	if v == nil {
+		return g.executionService.OpenKVStore(ctx)
+	}
+	genCtx, ok := v.(*genesisContext)
+	if !ok {
+		panic(fmt.Errorf("unexpected genesis context type: %T", v))
+	}
+	state, err := genCtx.state.GetWriter(g.actor)
+	if err != nil {
+		panic(err)
+	}
+	return state
+}
+
+type GenesisHeaderService struct {
+	executionService header.Service
+}
+
+// HeaderInfo implements header.Service.
+func (g *GenesisHeaderService) HeaderInfo(ctx context.Context) header.Info {
+	v := ctx.Value(genesisContextKey)
+	if v == nil {
+		return g.executionService.HeaderInfo(ctx)
+	}
+	return header.Info{}
+}
+
+func NewGenesisHeaderService(executionService header.Service) *GenesisHeaderService {
+	return &GenesisHeaderService{
+		executionService: executionService,
+	}
+}
+
+var _ header.Service = (*GenesisHeaderService)(nil)

server/v2/appmanager/appmanager.go

@@ -43,25 +43,19 @@ func (a AppManager[T]) InitGenesis(
 	initGenesisJSON []byte,
 	txDecoder transaction.Codec[T],
 ) (*server.BlockResponse, corestore.WriterMap, error) {
-	v, zeroState, err := a.db.StateLatest()
-	if err != nil {
-		return nil, 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, nil, errors.New("cannot init genesis on non-zero state")
-	}
-
 	var genTxs []T
-	genesisState, err := a.stf.RunWithCtx(ctx, zeroState, func(ctx context.Context) error {
-		return a.initGenesis(ctx, bytes.NewBuffer(initGenesisJSON), func(jsonTx json.RawMessage) error {
+	genesisState, err := a.initGenesis(
+		ctx,
+		bytes.NewBuffer(initGenesisJSON),
+		func(jsonTx json.RawMessage) error {
 			genTx, err := txDecoder.DecodeJSON(jsonTx)
 			if err != nil {
 				return fmt.Errorf("failed to decode genesis transaction: %w", err)
 			}
 			genTxs = append(genTxs, genTx)
 			return nil
-		})
-	})
+		},
+	)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to import genesis state: %w", err)
 	}
@@ -89,29 +83,11 @@ 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) {
-	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)
+	if a.exportGenesis == nil {
+		return nil, errors.New("export genesis function not set")
 	}
 
-	return bz, nil
+	return a.exportGenesis(ctx, version)
 }
 
 func (a AppManager[T]) DeliverBlock(
@@ -180,10 +156,6 @@ func (a AppManager[T]) Query(ctx context.Context, version uint64, request transa
 // QueryWithState executes a query with the provided state. This allows to process a query
 // independently of the db state. For example, it can be used to process a query with temporary
 // and uncommitted state
-func (a AppManager[T]) QueryWithState(
-	ctx context.Context,
-	state corestore.ReaderMap,
-	request transaction.Msg,
-) (transaction.Msg, error) {
+func (a AppManager[T]) QueryWithState(ctx context.Context, state corestore.ReaderMap, request transaction.Msg) (transaction.Msg, error) {
 	return a.stf.Query(ctx, state, a.config.QueryGasLimit, request)
 }

server/v2/appmanager/genesis.go

@@ -4,11 +4,25 @@ import (
 	"context"
 	"encoding/json"
 	"io"
+
+	"cosmossdk.io/core/store"
 )
 
 type (
 	// ExportGenesis is a function type that represents the export of the genesis state.
 	ExportGenesis func(ctx context.Context, version uint64) ([]byte, error)
-	// InitGenesis is a function type that represents the initialization of the genesis state.
-	InitGenesis func(ctx context.Context, src io.Reader, txHandler func(json.RawMessage) error) error
+
+	// InitGenesis is a function that will run at application genesis, it will be called with
+	// the following arguments:
+	// - ctx: the context of the genesis operation
+	// - src: the source containing the raw genesis state
+	// - txHandler: a function capable of decoding a json tx, will be run for each genesis
+	//   transaction
+	//
+	// It must return a map of the dirty state after the genesis operation.
+	InitGenesis func(
+		ctx context.Context,
+		src io.Reader,
+		txHandler func(json.RawMessage) error,
+	) (store.WriterMap, error)
 )