feat(saexec): Configure state

blocks/blockstest/blocks.go

@@ -130,6 +130,11 @@ func NewGenesis(tb testing.TB, db ethdb.Database, xdb saetypes.ExecutionResults,
 	}
 
 	tdb := state.NewDatabaseWithConfig(db, conf.tdbConfig).TrieDB()
+	defer func() {
+		if err := tdb.Close(); err != nil {
+			tb.Errorf("triedb.Database.Close(): %v", err)
+		}
+	}()
 	_, hash, err := core.SetupGenesisBlock(db, tdb, gen)
 	require.NoError(tb, err, "core.SetupGenesisBlock()")
 	require.NoErrorf(tb, tdb.Commit(hash, true), "%T.Commit(core.SetupGenesisBlock(...))", tdb)

sae/always.go

@@ -6,12 +6,16 @@ package sae
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 
 	"github.com/ava-labs/avalanchego/database"
 	"github.com/ava-labs/avalanchego/snow"
 	snowcommon "github.com/ava-labs/avalanchego/snow/engine/common"
 	"github.com/ava-labs/libevm/core"
+	"github.com/ava-labs/libevm/core/types"
+	"github.com/ava-labs/libevm/ethdb"
+	"github.com/ava-labs/libevm/params"
 	"github.com/ava-labs/libevm/triedb"
 
 	"github.com/ava-labs/strevm/adaptor"
@@ -50,25 +54,39 @@ func (vm *SinceGenesis[_]) Initialize(
 	appSender snowcommon.AppSender,
 ) error {
 	db := newEthDB(avaDB)
-	tdb := triedb.NewDatabase(db, vm.config.DBConfig.TrieDBConfig)
 
-	genesis := new(core.Genesis)
-	if err := json.Unmarshal(genesisBytes, genesis); err != nil {
-		return fmt.Errorf("json.Unmarshal(%T): %v", genesis, err)
-	}
-	config, _, err := core.SetupGenesisBlock(db, tdb, genesis)
+	g, config, err := createGenesisBlock(db, vm.config.DBConfig.TrieDBConfig(), genesisBytes)
 	if err != nil {
-		return fmt.Errorf("core.SetupGenesisBlock(...): %v", err)
+		return err
 	}
 
-	inner, err := NewVM(ctx, vm.hooks, vm.config, snowCtx, config, db, genesis.ToBlock(), appSender)
+	inner, err := NewVM(ctx, vm.hooks, vm.config, snowCtx, config, db, g, appSender)
 	if err != nil {
 		return err
 	}
 	vm.VM = inner
 	return nil
 }
 
+func createGenesisBlock(db ethdb.Database, tdbConfig *triedb.Config, genesisBytes []byte) (_ *types.Block, _ *params.ChainConfig, err error) {
+	tdb := triedb.NewDatabase(db, tdbConfig)
+	defer func() {
+		err = errors.Join(err, tdb.Close())
+	}()
+
+	genesis := new(core.Genesis)
+	if err := json.Unmarshal(genesisBytes, genesis); err != nil {
+		return nil, nil, fmt.Errorf("json.Unmarshal(%T): %v", genesis, err)
+	}
+
+	config, _, err := core.SetupGenesisBlock(db, tdb, genesis)
+	if err != nil {
+		return nil, nil, fmt.Errorf("core.SetupGenesisBlock(...): %v", err)
+	}
+
+	return genesis.ToBlock(), config, nil
+}
+
 // Shutdown gracefully closes the VM.
 func (vm *SinceGenesis[_]) Shutdown(ctx context.Context) error {
 	if vm.VM == nil {

saedb/tracker.go

@@ -14,22 +14,62 @@ import (
 	"github.com/ava-labs/libevm/core/state/snapshot"
 	"github.com/ava-labs/libevm/ethdb"
 	"github.com/ava-labs/libevm/triedb"
+	"github.com/ava-labs/libevm/triedb/hashdb"
 	"go.uber.org/zap"
 
 	"github.com/ava-labs/strevm/hook"
 )
 
+const (
+	// DefaultSnapshotCacheSizeMB is the snapshot cache size used by the executor.
+	DefaultSnapshotCacheSizeMB = 128
+	// DefaultTrieDBCacheSizeMB is the default cache size for trie nodes.
+	DefaultTrieDBCacheSizeMB = 512
+)
+
 // Config allows parameterization of the TrieDB and when
 // state is committed.
 type Config struct {
-	// TODO(alarso16): move minimal elements to config and construct in method.
-	TrieDBConfig *triedb.Config
-	Archival     bool // if true, will store every state on disk
+	Archival             bool // whether to store every state on disk
+	SnapshotCacheSizeMB  int  // Default: [DefaultSnapshotCacheSizeMB]. Set < 0 to disable snapshot.
+	TrieDBCacheSizeBytes int  // Default: [DefaultTrieDBCacheSizeMB]. Set < 0 to disable cache.
 }
 
-// SnapshotCacheSizeMB is the snapshot cache size used by a [Tracker].
-// TODO(alarso16): move to config
-const SnapshotCacheSizeMB = 128
+// TrieDBConfig returns a config that can be used to create a [triedb.Database] based on
+// the [Config] parameters provided.
+//
+// All [triedb.Database] must be closed, unless the TrieDB cache is disabled, as
+// this will result in a memory leak.
+func (c Config) TrieDBConfig() *triedb.Config {
+	cacheSize := DefaultTrieDBCacheSizeMB
+	switch {
+	case c.TrieDBCacheSizeBytes <= 0:
+		cacheSize = 0
+	case c.TrieDBCacheSizeBytes > 0:
+		cacheSize = c.TrieDBCacheSizeBytes
+	}
+
+	return &triedb.Config{
+		HashDB: &hashdb.Config{
+			CleanCacheSize: cacheSize,
+		},
+	}
+}
+
+func (c Config) snapConfig() *snapshot.Config {
+	size := DefaultSnapshotCacheSizeMB
+	switch {
+	case c.SnapshotCacheSizeMB < 0:
+		return nil
+	case c.SnapshotCacheSizeMB > 0:
+		size = c.SnapshotCacheSizeMB
+	}
+
+	return &snapshot.Config{
+		CacheSize:  size,
+		AsyncBuild: true,
+	}
+}
 
 var _ StateDBOpener = (*Tracker)(nil)
 
@@ -41,27 +81,28 @@ var _ StateDBOpener = (*Tracker)(nil)
 type Tracker struct {
 	snaps      *snapshot.Tree
 	cache      state.Database
-	isHashDB   bool
 	isArchival bool
 	log        logging.Logger
 }
 
 // NewTracker provides a new [Tracker] on the underlying database.
 func NewTracker(db ethdb.Database, c Config, lastExecuted common.Hash, log logging.Logger) (*Tracker, error) {
-	cache := state.NewDatabaseWithConfig(db, c.TrieDBConfig)
+	cache := state.NewDatabaseWithConfig(db, c.TrieDBConfig())
 	_, isHashDB := cache.TrieDB().Backend().(triedb.HashDB)
-	snapConf := snapshot.Config{
-		CacheSize:  SnapshotCacheSizeMB,
-		AsyncBuild: true,
+	if !isHashDB {
+		return nil, fmt.Errorf("unsupported DB: %T", cache.TrieDB().Backend())
 	}
-	snaps, err := snapshot.New(snapConf, db, cache.TrieDB(), lastExecuted)
-	if err != nil {
-		return nil, err
+	var snaps *snapshot.Tree
+	if snapConf := c.snapConfig(); snapConf != nil {
+		var err error
+		snaps, err = snapshot.New(*snapConf, db, cache.TrieDB(), lastExecuted)
+		if err != nil {
+			return nil, err
+		}
 	}
 	return &Tracker{
 		snaps:      snaps,
 		cache:      cache,
-		isHashDB:   isHashDB,
 		isArchival: c.Archival,
 		log:        log,
 	}, nil
@@ -74,11 +115,7 @@ func NewTracker(db ethdb.Database, c Config, lastExecuted common.Hash, log loggi
 // This state will be available in memory until [Tracker.Untrack] has been
 // called for the root as many times as [Tracker.Track] has been called.
 func (t *Tracker) Track(root common.Hash) {
-	if !t.isHashDB {
-		return
-	}
-
-	// Never returns an error because of the above check.
+	// Never returns an error because it's [triedb.HashDB].
 	if err := t.cache.TrieDB().Reference(root, common.Hash{}); err != nil {
 		t.log.Error("*triedb.Database.Reference()", zap.Error(err))
 	}
@@ -148,11 +185,7 @@ func LastHeightWithExecutionRootCommitted(db ethdb.Database, c Config, hooks hoo
 // This should be called on each block after its state is no longer
 // needed. If the state is already on disk, no operation is performed.
 func (t *Tracker) Untrack(root common.Hash) {
-	if !t.isHashDB {
-		return
-	}
-
-	// Never returns an error because of the above check.
+	// Never returns an error because it's [triedb.HashDB].
 	if err := t.cache.TrieDB().Dereference(root); err != nil {
 		t.log.Error("*triedb.Database.Dereference()", zap.Error(err))
 	}
@@ -173,21 +206,25 @@ func (t *Tracker) StateDB(root common.Hash) (*state.StateDB, error) {
 // recent state root.
 func (t *Tracker) Close(lastRoot common.Hash) (errs error) {
 	defer func() {
-		t.snaps.Release()
 		if err := t.cache.TrieDB().Close(); err != nil {
 			errs = errors.Join(errs, fmt.Errorf("triedb.Database.Close(): %v", err))
 		}
 	}()
 
+	if t.snaps == nil {
+		return nil
+	}
+
+	defer t.snaps.Release()
 	// We don't use [snapshot.Tree.Journal] because re-orgs are impossible under
 	// SAE so we don't mind flattening all snapshot layers to disk. Note that
 	// calling `Cap([disk root], 0)` returns an error when it's actually a
 	// no-op, so we ensure there are changes.
 	if lastRoot != t.snaps.DiskRoot() {
 		if err := t.snaps.Cap(lastRoot, 0); err != nil {
-			errs = errors.Join(errs, fmt.Errorf("snapshot.Tree.Cap([last post-execution state root], 0): %v", err))
+			return fmt.Errorf("snapshot.Tree.Cap([last post-execution state root], 0): %v", err)
 		}
 	}
 
-	return errs
+	return nil
 }

saexec/saexec_test.go

@@ -96,23 +96,21 @@ func newSUT(tb testing.TB, opts ...sutOption) (context.Context, *SUT) {
 	}, opts...)
 	config := saetest.ChainConfig()
 	db := rawdb.NewMemoryDatabase()
-	tdbConfig := &triedb.Config{}
 	xdb := saetest.NewExecutionResultsDB()
+	saedbConfig := saedb.Config{
+		Archival: sutCfg.archival,
+	}
 
 	wallet := saetest.NewUNSAFEWallet(tb, 1, types.LatestSigner(config))
 	alloc := saetest.MaxAllocFor(wallet.Addresses()...)
-	genesis := blockstest.NewGenesis(tb, db, xdb, config, alloc, blockstest.WithTrieDBConfig(tdbConfig), blockstest.WithGasTarget(sutCfg.hooks.Target))
+	genesis := blockstest.NewGenesis(tb, db, xdb, config, alloc, blockstest.WithTrieDBConfig(saedbConfig.TrieDBConfig()), blockstest.WithGasTarget(sutCfg.hooks.Target))
 
 	blockOpts := blockstest.WithBlockOptions(
 		blockstest.WithLogger(logger),
 	)
 	chain := blockstest.NewChainBuilder(config, genesis, blockOpts)
 	src := blocks.Source(chain.GetBlock)
 
-	saedbConfig := saedb.Config{
-		TrieDBConfig: tdbConfig,
-		Archival:     sutCfg.archival,
-	}
 	e, err := New(genesis, src.AsHeaderSource(), config, db, xdb, saedbConfig, sutCfg.hooks, logger)
 	require.NoError(tb, err, "New()")
 
@@ -915,7 +913,7 @@ func TestSnapshotPersistence(t *testing.T) {
 	// The crux of the test is whether we can recover the EOA nonce using only a
 	// new set of snapshots, recovered from the databases.
 	conf := snapshot.Config{
-		CacheSize: saedb.SnapshotCacheSizeMB,
+		CacheSize: saedb.DefaultSnapshotCacheSizeMB,
 		NoBuild:   true, // i.e. MUST be loaded from disk
 	}
 	snaps, err := snapshot.New(conf, sut.db, triedb.NewDatabase(e.db, nil), last.PostExecutionStateRoot())

worstcase/state_benchmark_test.go

@@ -34,7 +34,7 @@ func BenchmarkApplyTxWithSnapshot(b *testing.B) {
 			snaps, err := snapshot.New(
 				snapshot.Config{
 					AsyncBuild: false,
-					CacheSize:  saedb.SnapshotCacheSizeMB,
+					CacheSize:  saedb.DefaultSnapshotCacheSizeMB,
 				},
 				sut.db, sut.stateCache.TrieDB(), sut.genesis.PostExecutionStateRoot(),
 			)