TraceDB: Snapshot-backed state for the trace baker

app/app.go

@@ -41,6 +41,7 @@ import (
 	"github.com/sei-protocol/sei-chain/sei-cosmos/server/config"
 	servertypes "github.com/sei-protocol/sei-chain/sei-cosmos/server/types"
 	storetypes "github.com/sei-protocol/sei-chain/sei-cosmos/store/types"
+	storev2_rootmulti "github.com/sei-protocol/sei-chain/sei-cosmos/storev2/rootmulti"
 	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
 	sdkerrors "github.com/sei-protocol/sei-chain/sei-cosmos/types/errors"
 	genesistypes "github.com/sei-protocol/sei-chain/sei-cosmos/types/genesis"
@@ -719,6 +720,15 @@ func New(
 			panic(fmt.Sprintf("failed to open trace db: %s", dbErr))
 		}
 		app.EvmKeeper.SetTraceDB(traceDB)
+
+		if app.evmRPCConfig.TraceBakeUseSnapshot {
+			if rs, ok := app.CommitMultiStore().(*storev2_rootmulti.Store); ok {
+				app.EvmKeeper.SetTraceSnapshotStore(evmkeeper.NewTraceSnapshotStore(app.evmRPCConfig.TraceBakeSnapshotWindow))
+				app.EvmKeeper.SetTraceSnapshotCapture(rs.SnapshotSCStore)
+			} else {
+				logger.Info("trace_bake_use_snapshot set but commit multistore is not storev2 rootmulti; falling back to SS-pebble")
+			}
+		}
 	}
 	app.adminConfig, err = admin.ReadConfig(appOpts)
 	if err != nil {
@@ -1085,6 +1095,9 @@ func (app *App) HandleClose() error {
 			errs = append(errs, fmt.Errorf("failed to close trace db: %w", err))
 		}
 	}
+	if ts := app.EvmKeeper.TraceSnapshotStore(); ts != nil {
+		ts.Close()
+	}
 
 	// Close receipt store
 	if app.receiptStore != nil {
@@ -2446,6 +2459,49 @@ func (app *App) RPCContextProvider(i int64) sdk.Context {
 	return ctx.WithIsEVM(true).WithTraceMode(true).WithIsCheckTx(false)
 }
 
+// SnapshotAwareRPCContextProvider builds SDK contexts from in-memory memiavl
+// snapshots; falls back to RPCContextProvider on miss or unsupported backend.
+func (app *App) SnapshotAwareRPCContextProvider() evmrpc.TraceContextProvider {
+	store := app.EvmKeeper.TraceSnapshotStore()
+	if store == nil {
+		return evmrpc.TraceContextProvider(func(i int64) (sdk.Context, func()) {
+			return app.RPCContextProvider(i), func() {}
+		})
+	}
+	rs, ok := app.CommitMultiStore().(*storev2_rootmulti.Store)
+	if !ok {
+		return evmrpc.TraceContextProvider(func(i int64) (sdk.Context, func()) {
+			return app.RPCContextProvider(i), func() {}
+		})
+	}
+	return evmrpc.TraceContextProvider(func(i int64) (sdk.Context, func()) {
+		if i <= 0 {
+			return app.RPCContextProvider(i), func() {}
+		}
+		snap, release := store.Lease(i)
+		if snap == nil {
+			return app.RPCContextProvider(i), func() {}
+		}
+		cms, err := rs.CacheMultiStoreFromCommitter(snap)
+		if err != nil {
+			release()
+			return app.RPCContextProvider(i), func() {}
+		}
+		checkCtx := app.GetCheckCtx()
+		closestUpgrade, upgradeHeight := app.UpgradeKeeper.GetClosestUpgrade(checkCtx, i)
+		if closestUpgrade == "" && upgradeHeight == 0 {
+			closestUpgrade = LatestUpgrade
+		}
+		ctx := sdk.NewContext(cms, checkCtx.BlockHeader(), true).
+			WithMinGasPrices(checkCtx.MinGasPrices()).
+			WithBlockHeight(i).
+			WithClosestUpgradeName(closestUpgrade).
+			WithIsEVM(true).WithTraceMode(true).WithIsCheckTx(false)
+		ctx = ctx.WithConsensusParams(app.GetConsensusParams(ctx))
+		return ctx, release
+	})
+}
+
 // RegisterTendermintService implements the Application.RegisterTendermintService method.
 func (app *App) RegisterTendermintService(clientCtx client.Context) {
 	tmservice.RegisterTendermintService(app.GRPCQueryRouter(), clientCtx, app.interfaceRegistry)
@@ -2460,8 +2516,10 @@ func (app *App) RegisterTendermintService(clientCtx client.Context) {
 		return app.legacyEncodingConfig.TxConfig
 	}
 
+	rpcCtxProvider := app.RPCContextProvider
+	traceCtxProvider := app.SnapshotAwareRPCContextProvider()
 	if app.evmRPCConfig.HTTPEnabled {
-		evmHTTPServer, err := evmrpc.NewEVMHTTPServer(app.evmRPCConfig, clientCtx.Client, &app.EvmKeeper, app.BeginBlockKeepers, app.BaseApp, app.TracerAnteHandler, app.RPCContextProvider, txConfigProvider, DefaultNodeHome, app.GetStateStore(), nil)
+		evmHTTPServer, err := evmrpc.NewEVMHTTPServer(app.evmRPCConfig, clientCtx.Client, &app.EvmKeeper, app.BeginBlockKeepers, app.BaseApp, app.TracerAnteHandler, rpcCtxProvider, txConfigProvider, DefaultNodeHome, app.GetStateStore(), nil, traceCtxProvider)
 		if err != nil {
 			panic(err)
 		}
@@ -2474,7 +2532,7 @@ func (app *App) RegisterTendermintService(clientCtx client.Context) {
 	}
 
 	if app.evmRPCConfig.WSEnabled {
-		evmWSServer, err := evmrpc.NewEVMWebSocketServer(app.evmRPCConfig, clientCtx.Client, &app.EvmKeeper, app.BeginBlockKeepers, app.BaseApp, app.TracerAnteHandler, app.RPCContextProvider, txConfigProvider, DefaultNodeHome, app.GetStateStore())
+		evmWSServer, err := evmrpc.NewEVMWebSocketServer(app.evmRPCConfig, clientCtx.Client, &app.EvmKeeper, app.BeginBlockKeepers, app.BaseApp, app.TracerAnteHandler, rpcCtxProvider, txConfigProvider, DefaultNodeHome, app.GetStateStore())
 		if err != nil {
 			panic(err)
 		}

app/app_test.go

@@ -27,6 +27,7 @@ import (
 	cosmosed25519 "github.com/sei-protocol/sei-chain/sei-cosmos/crypto/keys/ed25519"
 	"github.com/sei-protocol/sei-chain/sei-cosmos/crypto/keys/secp256k1"
 	cryptotypes "github.com/sei-protocol/sei-chain/sei-cosmos/crypto/types"
+	storev2_rootmulti "github.com/sei-protocol/sei-chain/sei-cosmos/storev2/rootmulti"
 	sdk "github.com/sei-protocol/sei-chain/sei-cosmos/types"
 	"github.com/sei-protocol/sei-chain/sei-cosmos/types/tx/signing"
 	xauthsigning "github.com/sei-protocol/sei-chain/sei-cosmos/x/auth/signing"
@@ -37,6 +38,7 @@ import (
 	tmtypes "github.com/sei-protocol/sei-chain/sei-tendermint/types"
 	testkeeper "github.com/sei-protocol/sei-chain/testutil/keeper"
 	"github.com/sei-protocol/sei-chain/x/evm/config"
+	evmkeeper "github.com/sei-protocol/sei-chain/x/evm/keeper"
 	evmtypes "github.com/sei-protocol/sei-chain/x/evm/types"
 	"github.com/sei-protocol/sei-chain/x/evm/types/ethtx"
 	oracletypes "github.com/sei-protocol/sei-chain/x/oracle/types"
@@ -1048,6 +1050,40 @@ func TestRPCContextProviderPopulatesConsensusParams(t *testing.T) {
 	})
 }
 
+func TestSnapshotAwareRPCContextProviderPopulatesConsensusParams(t *testing.T) {
+	valPub := cosmosed25519.GenPrivKey().PubKey()
+	accAddr := sdk.AccAddress(valPub.Address())
+	genAcc := authtypes.NewBaseAccount(accAddr, nil, 0, 0)
+	balance := banktypes.Balance{
+		Address: accAddr.String(),
+		Coins:   sdk.NewCoins(sdk.NewCoin(sdk.DefaultBondDenom, sdk.DefaultPowerReduction)),
+	}
+	tmPub, err := cryptocodec.ToTmPubKeyInterface(valPub)
+	require.NoError(t, err)
+	valSet := tmtypes.NewValidatorSet([]*tmtypes.Validator{tmtypes.NewValidator(tmPub, 1)})
+
+	testApp := app.SetupWithGenesisValSet(t, valSet, []authtypes.GenesisAccount{genAcc}, balance)
+	rs, ok := testApp.CommitMultiStore().(*storev2_rootmulti.Store)
+	require.True(t, ok)
+	snap := rs.SnapshotSCStore()
+	require.NotNil(t, snap)
+	testApp.EvmKeeper.SetTraceSnapshotStore(evmkeeper.NewTraceSnapshotStore(8))
+	testApp.EvmKeeper.TraceSnapshotStore().Put(snap.Version(), snap)
+
+	ctx, release := testApp.SnapshotAwareRPCContextProvider()(snap.Version())
+	defer release()
+
+	cp := ctx.ConsensusParams()
+	require.NotNil(t, cp, "ConsensusParams must be populated on the snapshot RPC ctx")
+	require.NotNil(t, cp.Block, "Block params must be populated")
+	require.Equal(t, app.DefaultConsensusParams.Block.MaxGas, cp.Block.MaxGas)
+	require.Equal(t, app.DefaultConsensusParams.Block.MaxBytes, cp.Block.MaxBytes)
+
+	rpcCtx := testApp.RPCContextProvider(snap.Version())
+	require.Equal(t, rpcCtx.ChainID(), ctx.ChainID())
+	require.Equal(t, rpcCtx.BlockTime(), ctx.BlockTime())
+}
+
 func finalizeToBlockProcessReq(req *abci.RequestFinalizeBlock) *app.BlockProcessRequest {
 	var height int64
 	var blockTime time.Time

evmrpc/config/config.go

@@ -150,6 +150,12 @@ type Config struct {
 	TraceBakeQueueSize    int      `mapstructure:"trace_bake_queue_size"`    // in-flight height queue (default 4096)
 	TraceBakeTracers      []string `mapstructure:"trace_bake_tracers"`       // tracers to bake (default ["callTracer"])
 	TraceBakeWindowBlocks int64    `mapstructure:"trace_bake_window_blocks"` // rolling prune window; 0 disables
+
+	// TraceBakeUseSnapshot captures an in-memory memiavl snapshot at
+	// EndBlock and uses it as the state backend for the baker, bypassing
+	// SS-pebble. Requires CosmosOnly write mode; falls back transparently.
+	TraceBakeUseSnapshot    bool  `mapstructure:"trace_bake_use_snapshot"`
+	TraceBakeSnapshotWindow int64 `mapstructure:"trace_bake_snapshot_window"` // recent snapshots to keep (default 64)
 }
 
 var DefaultConfig = Config{
@@ -187,11 +193,13 @@ var DefaultConfig = Config{
 		"sei_getEVMAddress",
 		"sei_getCosmosTx",
 	},
-	TraceBakeEnabled:      false,
-	TraceBakeWorkers:      1,
-	TraceBakeQueueSize:    4096,
-	TraceBakeTracers:      []string{"callTracer"},
-	TraceBakeWindowBlocks: 0,
+	TraceBakeEnabled:        false,
+	TraceBakeWorkers:        1,
+	TraceBakeQueueSize:      4096,
+	TraceBakeTracers:        []string{"callTracer"},
+	TraceBakeWindowBlocks:   0,
+	TraceBakeUseSnapshot:    false,
+	TraceBakeSnapshotWindow: 64,
 }
 
 const (
@@ -230,6 +238,8 @@ const (
 	flagTraceBakeQueueSize           = "evm.trace_bake_queue_size"
 	flagTraceBakeTracers             = "evm.trace_bake_tracers"
 	flagTraceBakeWindowBlocks        = "evm.trace_bake_window_blocks"
+	flagTraceBakeUseSnapshot         = "evm.trace_bake_use_snapshot"
+	flagTraceBakeSnapshotWindow      = "evm.trace_bake_snapshot_window"
 )
 
 func ReadConfig(opts servertypes.AppOptions) (Config, error) {
@@ -410,6 +420,16 @@ func ReadConfig(opts servertypes.AppOptions) (Config, error) {
 			return cfg, err
 		}
 	}
+	if v := opts.Get(flagTraceBakeUseSnapshot); v != nil {
+		if cfg.TraceBakeUseSnapshot, err = cast.ToBoolE(v); err != nil {
+			return cfg, err
+		}
+	}
+	if v := opts.Get(flagTraceBakeSnapshotWindow); v != nil {
+		if cfg.TraceBakeSnapshotWindow, err = cast.ToInt64E(v); err != nil {
+			return cfg, err
+		}
+	}
 
 	return cfg, nil
 }
@@ -585,4 +605,17 @@ trace_bake_tracers = [{{- range $i, $t := .EVM.TraceBakeTracers }}{{- if $i }},
 # Rolling cache window: prune blocks older than (latest - this).
 # 0 disables pruning (cache grows forever).
 trace_bake_window_blocks = {{ .EVM.TraceBakeWindowBlocks }}
+
+# TraceBakeUseSnapshot, when true, uses in-memory memiavl snapshots as the
+# state backend for trace baking when the store backend supports snapshots.
+# Watch these metrics when enabling on a high-throughput node:
+#   - memiavl_mem_node_total_size / memiavl_num_of_mem_node: rise if held
+#     snapshots are pinning too many COW nodes; lower the window or drop the
+#     memiavl snapshot interval.
+#   - trace baker dropped/baked counters: dropped > 0 or baked lagging chain
+#     tip means the baker is falling behind.
+trace_bake_use_snapshot = {{ .EVM.TraceBakeUseSnapshot }}
+
+# Number of recent memiavl snapshots to retain for trace baking.
+trace_bake_snapshot_window = {{ .EVM.TraceBakeSnapshotWindow }}
 `

evmrpc/config/config_test.go

@@ -139,7 +139,9 @@ func (o *opts) Get(k string) interface{} {
 		k == "evm.trace_bake_workers" ||
 		k == "evm.trace_bake_queue_size" ||
 		k == "evm.trace_bake_tracers" ||
-		k == "evm.trace_bake_window_blocks" {
+		k == "evm.trace_bake_window_blocks" ||
+		k == "evm.trace_bake_use_snapshot" ||
+		k == "evm.trace_bake_snapshot_window" {
 		return nil
 	}
 	panic("unknown key")

evmrpc/server.go

@@ -45,6 +45,7 @@ func NewEVMHTTPServer(
 	homeDir string,
 	stateStore types.StateStore,
 	isPanicOrSyntheticTxFunc func(ctx context.Context, hash common.Hash) (bool, error), // used in *ExcludeTraceFail endpoints
+	traceCtxProviders ...TraceContextProvider,
 ) (EVMServer, error) {
 
 	// Initialize global worker pool with configuration (metrics are embedded in pool)
@@ -89,8 +90,13 @@ func NewEVMHTTPServer(
 	sendAPI := NewSendAPI(tmClient, txConfigProvider, &SendConfig{slow: config.Slow}, k, beginBlockKeepers, ctxProvider, homeDir, simulateConfig, app, antehandler, ConnectionTypeHTTP, globalBlockCache, cacheCreationMutex, watermarks)
 
 	ctx := ctxProvider(LatestCtxHeight)
+	traceCtxProvider := defaultTraceContextProvider(ctxProvider)
+	if len(traceCtxProviders) > 0 && traceCtxProviders[0] != nil {
+		traceCtxProvider = traceCtxProviders[0]
+	}
 	txAPI := NewTransactionAPI(tmClient, k, ctxProvider, txConfigProvider, homeDir, ConnectionTypeHTTP, watermarks, globalBlockCache, cacheCreationMutex)
 	debugAPI := NewDebugAPI(tmClient, k, beginBlockKeepers, ctxProvider, txConfigProvider, simulateConfig, app, antehandler, ConnectionTypeHTTP, config, globalBlockCache, cacheCreationMutex, watermarks)
+	debugAPI.backend.SetTraceContextProvider(traceCtxProvider)
 	if config.TraceBakeEnabled {
 		StartTraceBakerForDebugAPI(debugAPI, TraceBakerConfig{
 			Workers:      config.TraceBakeWorkers,
@@ -109,6 +115,7 @@ func NewEVMHTTPServer(
 
 	seiTxAPI := NewSeiTransactionAPI(tmClient, k, ctxProvider, txConfigProvider, homeDir, ConnectionTypeHTTP, isPanicOrSyntheticTxFunc, watermarks, globalBlockCache, cacheCreationMutex)
 	seiDebugAPI := NewSeiDebugAPI(tmClient, k, beginBlockKeepers, ctxProvider, txConfigProvider, simulateConfig, app, antehandler, ConnectionTypeHTTP, config, globalBlockCache, cacheCreationMutex, watermarks)
+	seiDebugAPI.backend.SetTraceContextProvider(traceCtxProvider)
 
 	// DB semaphore aligned with worker count
 	dbReadSemaphore := make(chan struct{}, workerCount)