feat: Add Firewood as a state syncer

graft/coreth/core/blockchain_ext_test.go

@@ -6,8 +6,6 @@ package core
 import (
 	"fmt"
 	"math/big"
-	"os"
-	"path/filepath"
 	"slices"
 	"testing"
 
@@ -25,6 +23,7 @@ import (
 	"github.com/ava-labs/avalanchego/graft/coreth/params"
 	"github.com/ava-labs/avalanchego/graft/coreth/plugin/evm/customheader"
 	"github.com/ava-labs/avalanchego/graft/coreth/plugin/evm/upgrade/ap4"
+	"github.com/ava-labs/avalanchego/graft/evm/utils/utilstest"
 
 	ethparams "github.com/ava-labs/libevm/params"
 )
@@ -153,63 +152,6 @@ var reexecTests = []ReexecTest{
 	},
 }
 
-func copyMemDB(db ethdb.Database) (ethdb.Database, error) {
-	newDB := rawdb.NewMemoryDatabase()
-	iter := db.NewIterator(nil, nil)
-	defer iter.Release()
-	for iter.Next() {
-		if err := newDB.Put(iter.Key(), iter.Value()); err != nil {
-			return nil, err
-		}
-	}
-
-	return newDB, nil
-}
-
-// copyDir recursively copies all files and folders from a directory [src] to a
-// new temporary directory and returns the path to the new directory.
-func copyDir(t *testing.T, src string) string {
-	t.Helper()
-
-	if src == "" {
-		return ""
-	}
-
-	dst := t.TempDir()
-	err := filepath.Walk(src, func(path string, info os.FileInfo, err error) error {
-		if err != nil {
-			return err
-		}
-
-		// Calculate the relative path from src
-		relPath, err := filepath.Rel(src, path)
-		if err != nil {
-			return err
-		}
-
-		// Skip the root directory itself
-		if relPath == "." {
-			return nil
-		}
-
-		dstPath := filepath.Join(dst, relPath)
-
-		if info.IsDir() {
-			return os.MkdirAll(dstPath, info.Mode().Perm())
-		}
-
-		data, err := os.ReadFile(path)
-		if err != nil {
-			return err
-		}
-
-		return os.WriteFile(dstPath, data, info.Mode().Perm())
-	})
-
-	require.NoError(t, err)
-	return dst
-}
-
 // checkBlockChainState creates a new BlockChain instance and checks that exporting each block from
 // genesis to last accepted from the original instance yields the same last accepted block and state
 // root.
@@ -255,9 +197,8 @@ func checkBlockChainState(
 	require.NoErrorf(checkState(acceptedState), "Check state failed for newly generated blockchain")
 
 	// Copy the database over to prevent any issues when re-using [originalDB] after this call.
-	originalDB, err = copyMemDB(originalDB)
-	require.NoError(err)
-	newChainDataDir := copyDir(t, oldChainDataDir)
+	originalDB = utilstest.CopyEthDB(t, originalDB)
+	newChainDataDir := utilstest.CopyDir(t, oldChainDataDir)
 	restartedChain, err := create(originalDB, gspec, lastAcceptedBlock.Hash(), newChainDataDir)
 	require.NoError(err)
 	defer restartedChain.Stop()
@@ -1702,7 +1643,7 @@ func ReexecCorruptedStateTest(t *testing.T, create ReexecTestFunc) {
 	blockchain.Stop()
 
 	// Restart blockchain with existing state
-	newDir := copyDir(t, tempDir) // avoid file lock
+	newDir := utilstest.CopyDir(t, tempDir) // avoid file lock
 	restartedBlockchain, err := create(chainDB, gspec, chain[1].Hash(), newDir, 4096)
 	require.NoError(t, err)
 	defer restartedBlockchain.Stop()

graft/coreth/core/chain_makers.go

@@ -32,7 +32,6 @@ import (
 	"math/big"
 
 	"github.com/ava-labs/avalanchego/graft/coreth/consensus"
-	"github.com/ava-labs/avalanchego/graft/coreth/core/extstate"
 	"github.com/ava-labs/avalanchego/graft/coreth/params"
 	"github.com/ava-labs/avalanchego/graft/coreth/plugin/evm/customheader"
 	"github.com/ava-labs/avalanchego/graft/coreth/plugin/evm/customtypes"
@@ -43,6 +42,7 @@ import (
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/core/vm"
 	"github.com/ava-labs/libevm/ethdb"
+	"github.com/ava-labs/libevm/libevm/stateconf"
 	"github.com/ava-labs/libevm/triedb"
 	"github.com/holiman/uint256"
 )
@@ -273,6 +273,12 @@ func (b *BlockGen) SetOnBlockGenerated(onBlockGenerated func(*types.Block)) {
 // values. Inserting them into BlockChain requires use of FakePow or
 // a similar non-validating proof of work implementation.
 func GenerateChain(config *params.ChainConfig, parent *types.Block, engine consensus.Engine, db ethdb.Database, n int, gap uint64, gen func(int, *BlockGen)) ([]*types.Block, []types.Receipts, error) {
+	stateCache := state.NewDatabase(db)
+	defer stateCache.TrieDB().Close()
+	return GenerateChainFromStateCache(config, parent, engine, stateCache, n, gap, gen)
+}
+
+func GenerateChainFromStateCache(config *params.ChainConfig, parent *types.Block, engine consensus.Engine, stateCache state.Database, n int, gap uint64, gen func(int, *BlockGen)) ([]*types.Block, []types.Receipts, error) {
 	if config == nil {
 		config = params.TestChainConfig
 	}
@@ -301,7 +307,8 @@ func GenerateChain(config *params.ChainConfig, parent *types.Block, engine conse
 		}
 
 		// Write state changes to db
-		root, err := statedb.Commit(b.header.Number.Uint64(), config.IsEIP158(b.header.Number))
+		statedbOpts := stateconf.WithTrieDBUpdateOpts(stateconf.WithTrieDBUpdatePayload(block.ParentHash(), block.Hash()))
+		root, err := statedb.Commit(b.header.Number.Uint64(), config.IsEIP158(b.header.Number), statedbOpts)
 		if err != nil {
 			panic(fmt.Sprintf("state write error: %v", err))
 		}
@@ -314,16 +321,12 @@ func GenerateChain(config *params.ChainConfig, parent *types.Block, engine conse
 		return block, b.receipts, nil
 	}
 
-	// Forcibly use hash-based state scheme for retaining all nodes in disk.
-	triedb := triedb.NewDatabase(db, triedb.HashDefaults)
-	defer triedb.Close()
-
 	for i := 0; i < n; i++ {
-		statedb, err := state.New(parent.Root(), extstate.NewDatabaseWithNodeDB(db, triedb), nil)
+		statedb, err := state.New(parent.Root(), stateCache, nil)
 		if err != nil {
 			return nil, nil, err
 		}
-		block, receipts, err := genblock(i, parent, triedb, statedb)
+		block, receipts, err := genblock(i, parent, stateCache.TrieDB(), statedb)
 		if err != nil {
 			return nil, nil, err
 		}

graft/coreth/network/network.go

@@ -41,6 +41,7 @@ var (
 	_                     Network              = (*network)(nil)
 	_                     validators.Connector = (*network)(nil)
 	_                     common.AppHandler    = (*network)(nil)
+	_                     p2p.NodeSampler      = (*network)(nil)
 )
 
 // SyncedNetworkClient defines ability to send request / response through the Network
@@ -63,6 +64,7 @@ type SyncedNetworkClient interface {
 type Network interface {
 	validators.Connector
 	common.AppHandler
+	p2p.NodeSampler
 
 	SyncedNetworkClient
 
@@ -155,6 +157,24 @@ func NewNetwork(
 	}, nil
 }
 
+func (n *network) Sample(_ context.Context, limit int) []ids.NodeID {
+	if limit != 1 {
+		log.Warn("Sample called with limit > 1, but only 1 peer will be returned", "limit", limit)
+	}
+
+	n.lock.Lock()
+	defer n.lock.Unlock()
+	node, ok, err := n.peers.GetAnyPeer(nil)
+	if err != nil {
+		log.Error("error getting peer from peer tracker", "error", err)
+		return nil
+	}
+	if !ok {
+		return nil
+	}
+	return []ids.NodeID{node}
+}
+
 // SendAppRequestAny synchronously sends request to an arbitrary peer with a
 // node version greater than or equal to minVersion. If minVersion is nil,
 // the request will be sent to any peer regardless of their version.

graft/coreth/plugin/evm/atomic/sync/syncer_test.go

@@ -353,6 +353,7 @@ func setupTestInfrastructure(t *testing.T, serverTrieDB *triedb.Database) (conte
 		handlers.NewLeafsRequestHandler(serverTrieDB, state.TrieKeyLength, nil, message.CorethCodec, handlerstats.NewNoopHandlerStats()),
 		nil,
 		nil,
+		nil,
 	)
 
 	clientDB := versiondb.New(memdb.New())

graft/coreth/plugin/evm/config/config.md

@@ -221,7 +221,7 @@ Configuration is provided as a JSON object. All fields are optional unless other
 
 | Option | Type | Description | Default |
 |--------|------|-------------|---------|
-| `metrics-expensive-enabled` | bool | Enable expensive debug-level metrics; this includes Firewood metrics | `true` |
+| `metrics-expensive-enabled` | bool | Enable expensive debug-level metrics | `true` |
 
 ## Security and Access
 
@@ -259,7 +259,6 @@ Configuration is provided as a JSON object. All fields are optional unless other
 > **WARNING**: `firewood` and `path` schemes are untested in production. Using `path` is strongly discouraged. To use `firewood`, you must also set the following config options:
 >
 > - `populate-missing-tries: nil`
-> - `state-sync-enabled: false`
 > - `snapshot-cache: 0`
 
 Failing to set these options will result in errors on VM initialization. Additionally, not all APIs are available - see these portions of the config documentation for more details.

graft/coreth/plugin/evm/vm.go

@@ -41,6 +41,7 @@ import (
 	"github.com/ava-labs/avalanchego/cache/metercacher"
 	"github.com/ava-labs/avalanchego/codec"
 	"github.com/ava-labs/avalanchego/database"
+	"github.com/ava-labs/avalanchego/database/merkle/firewood/syncer"
 	"github.com/ava-labs/avalanchego/database/versiondb"
 	"github.com/ava-labs/avalanchego/graft/coreth/consensus/dummy"
 	"github.com/ava-labs/avalanchego/graft/coreth/core"
@@ -58,6 +59,7 @@ import (
 	"github.com/ava-labs/avalanchego/graft/coreth/precompile/precompileconfig"
 	"github.com/ava-labs/avalanchego/graft/coreth/warp"
 	"github.com/ava-labs/avalanchego/graft/evm/constants"
+	"github.com/ava-labs/avalanchego/graft/evm/firewood"
 	"github.com/ava-labs/avalanchego/graft/evm/message"
 	"github.com/ava-labs/avalanchego/graft/evm/rpc"
 	"github.com/ava-labs/avalanchego/graft/evm/sync/client"
@@ -142,7 +144,6 @@ var (
 	errShuttingDownVM                             = errors.New("shutting down VM")
 	errFirewoodSnapshotCacheDisabled              = errors.New("snapshot cache must be disabled for Firewood")
 	errFirewoodOfflinePruningUnsupported          = errors.New("offline pruning is not supported for Firewood")
-	errFirewoodStateSyncUnsupported               = errors.New("state sync is not yet supported for Firewood")
 	errFirewoodMissingTrieRepopulationUnsupported = errors.New("missing trie repopulation is not supported for Firewood")
 )
 
@@ -400,9 +401,6 @@ func (vm *VM) Initialize(
 		if vm.config.OfflinePruning {
 			return errFirewoodOfflinePruningUnsupported
 		}
-		if vm.config.StateSyncEnabled == nil || *vm.config.StateSyncEnabled {
-			return errFirewoodStateSyncUnsupported
-		}
 		if vm.config.PopulateMissingTries != nil {
 			return errFirewoodMissingTrieRepopulationUnsupported
 		}
@@ -582,36 +580,51 @@ func (vm *VM) initializeChain(lastAcceptedHash common.Hash) error {
 // If state sync is disabled, this function will wipe any ongoing summary from
 // disk to ensure that we do not continue syncing from an invalid snapshot.
 func (vm *VM) initializeStateSync(lastAcceptedHeight uint64) error {
-	// Create standalone EVM TrieDB (read only) for serving leafs requests.
-	// We create a standalone TrieDB here, so that it has a standalone cache from the one
-	// used by the node when processing blocks.
-	// However, Firewood does not support multiple TrieDBs, so we use the same one.
-	evmTrieDB := vm.eth.BlockChain().TrieDB()
-	if vm.ethConfig.StateScheme != customrawdb.FirewoodScheme {
-		evmTrieDB = triedb.NewDatabase(
+	leafHandlers := make(LeafHandlers)
+	leafMetricsNames := make(map[message.NodeType]string)
+	syncStats := handlerstats.GetOrRegisterHandlerStats(metrics.Enabled)
+
+	switch scheme := vm.ethConfig.StateScheme; scheme {
+	case rawdb.HashScheme, "":
+		// Create standalone EVM TrieDB (read only) for serving leafs requests.
+		// We create a standalone TrieDB here, so that it has a standalone cache from the one
+		// used by the node when processing blocks.
+		evmTrieDB := triedb.NewDatabase(
 			vm.chaindb,
 			&triedb.Config{
 				DBOverride: hashdb.Config{
 					CleanCacheSize: vm.config.StateSyncServerTrieCache * units.MiB,
 				}.BackendConstructor,
 			},
 		)
+		// register default leaf request handler for state trie
+		stateLeafRequestConfig := &extension.LeafRequestConfig{
+			LeafType:   message.StateTrieNode,
+			MetricName: "sync_state_trie_leaves",
+			Handler: handlers.NewLeafsRequestHandler(evmTrieDB,
+				message.StateTrieKeyLength,
+				vm.blockChain, vm.networkCodec,
+				syncStats,
+			),
+		}
+		leafHandlers[stateLeafRequestConfig.LeafType] = stateLeafRequestConfig.Handler
+		leafMetricsNames[stateLeafRequestConfig.LeafType] = stateLeafRequestConfig.MetricName
+	case customrawdb.FirewoodScheme:
+		tdb, ok := vm.eth.BlockChain().TrieDB().Backend().(*firewood.TrieDB)
+		if !ok {
+			return fmt.Errorf("expected a %T with %s scheme, got %T", tdb, customrawdb.FirewoodScheme, vm.eth.BlockChain().TrieDB().Backend())
+		}
+		n := vm.Network.P2PNetwork()
+		if err := n.AddHandler(p2p.FirewoodRangeProofHandlerID, syncer.NewGetRangeProofHandler(tdb.Firewood)); err != nil {
+			return fmt.Errorf("adding firewood range proof handler: %w", err)
+		}
+		if err := n.AddHandler(p2p.FirewoodChangeProofHandlerID, syncer.NewGetChangeProofHandler(tdb.Firewood)); err != nil {
+			return fmt.Errorf("adding firewood change proof handler: %w", err)
+		}
+	default:
+		log.Warn("state sync is not supported for this scheme, no leaf handlers will be registered", "scheme", scheme)
+		return nil
 	}
-	leafHandlers := make(LeafHandlers)
-	leafMetricsNames := make(map[message.NodeType]string)
-	// register default leaf request handler for state trie
-	syncStats := handlerstats.GetOrRegisterHandlerStats(metrics.Enabled)
-	stateLeafRequestConfig := &extension.LeafRequestConfig{
-		LeafType:   message.StateTrieNode,
-		MetricName: "sync_state_trie_leaves",
-		Handler: handlers.NewLeafsRequestHandler(evmTrieDB,
-			message.StateTrieKeyLength,
-			vm.blockChain, vm.networkCodec,
-			syncStats,
-		),
-	}
-	leafHandlers[stateLeafRequestConfig.LeafType] = stateLeafRequestConfig.Handler
-	leafMetricsNames[stateLeafRequestConfig.LeafType] = stateLeafRequestConfig.MetricName
 
 	extraLeafConfig := vm.extensionConfig.ExtraSyncLeafHandlerConfig
 	if extraLeafConfig != nil {
@@ -649,6 +662,7 @@ func (vm *VM) initializeStateSync(lastAcceptedHeight uint64) error {
 		StateSyncDone: vm.stateSyncDone,
 		Chain:         newChainContextAdapter(vm.eth),
 		State:         vm.State,
+		SnowCtx:       vm.ctx,
 		Client: client.New(
 			&client.Config{
 				Network:          vm.Network,