perf: parallel fetching trie

cmd/geth/chaincmd.go

@@ -47,6 +47,7 @@ import (
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/metrics"
 	"github.com/ethereum/go-ethereum/params"
+	"github.com/ethereum/go-ethereum/trie"
 	"github.com/urfave/cli/v2"
 )
 
@@ -450,6 +451,7 @@ func showMetrics() {
 	// execution
 	fmt.Println("prefetchTimer", blockPrefetchExecuteTimer.Total())
 	fmt.Println("accountReadTimer", accountReadTimer.Total())
+	fmt.Println("ResolveTimer", trie.ResolveTime)
 	fmt.Println("storageReadTimer", storageReadTimer.Total())
 	fmt.Println("blockExecutionTimer", blockExecutionTimer.Total())
 
@@ -476,6 +478,7 @@ func showMetrics() {
 	fmt.Println("PrefetchMergeTime:", core.PrefetchMergeBALTime)
 	fmt.Println("ParallelExeTime: ", core.ParallelExeTime)
 	fmt.Println("PostMergeTime:    ", core.PostMergeTime)
+	fmt.Println("PrefetchTrieTime: ", core.PrefetchTrieTimer)
 
 	// total
 	fmt.Println("blockInsertTimer", blockInsertTimer.Total())

core/parallel_state_processor.go

@@ -21,6 +21,7 @@ var (
 	PrefetchMergeBALTime = time.Duration(0)
 	ParallelExeTime      = time.Duration(0)
 	PostMergeTime        = time.Duration(0)
+	PrefetchTrieTimer    = time.Duration(0)
 )
 
 type ParallelStateProcessor struct {
@@ -84,11 +85,17 @@ func (p *ParallelStateProcessor) Process(block *types.Block, statedb *state.Stat
 		defer wg.Done()
 
 		start := time.Now()
-		postState = statedb.Copy()
+		postState = statedb.CopyState()
+		// Stop prefetcher cause we'll directly fetch tries in parallel
+		postState.StopPrefetcher()
+
 		postState.MergePostBalStates()
-		// prewarm the updating trie
-		postState.IntermediateRoot(true)
 		PostMergeTime += time.Since(start)
+
+		// Prewarm the updating trie
+		start = time.Now()
+		postState.PrefetchTrie()
+		PrefetchTrieTimer += time.Since(start)
 	}()
 
 	go func() {

core/state/state_object.go

@@ -144,6 +144,9 @@ func (s *stateObject) getTrie() (Trie, error) {
 func (s *stateObject) getPrefetchedTrie() Trie {
 	// If there's nothing to meaningfully return, let the user figure it out by
 	// pulling the trie from disk.
+	if s.trie != nil {
+		return s.trie
+	}
 	if (s.data.Root == types.EmptyRootHash && !s.db.db.TrieDB().IsVerkle()) || s.db.prefetcher == nil {
 		return nil
 	}
@@ -512,9 +515,9 @@ func (s *stateObject) simpleCopy(db *StateDB) *stateObject {
 		data:               s.data,
 		code:               s.code,
 		originStorage:      s.originStorage.Copy(),
-		pendingStorage:     make(Storage),
-		dirtyStorage:       make(Storage),
-		uncommittedStorage: make(Storage),
+		pendingStorage:     s.pendingStorage.Copy(),
+		dirtyStorage:       s.dirtyStorage.Copy(),
+		uncommittedStorage: s.uncommittedStorage.Copy(),
 		dirtyCode:          s.dirtyCode,
 		selfDestructed:     s.selfDestructed,
 		newContract:        s.newContract,

core/state/statedb.go

@@ -732,6 +732,54 @@ func (s *StateDB) MergePostBalStates() {
 	}
 }
 
+func (s *StateDB) PrefetchTrie() {
+	var workers errgroup.Group
+
+	for addr := range s.journal.dirties {
+		addr := addr
+		// s.trie.GetAccount(addr)
+		workers.Go(func() error {
+			_, err := s.trie.GetAccount(addr)
+			if err != nil {
+				fmt.Println("error fetching account:", err)
+			}
+			return err
+		})
+
+		workers.Go(func() error {
+			obj, exist := s.stateObjects[addr]
+			if exist {
+				tr, err := s.db.OpenStorageTrie(s.originalRoot, addr, obj.data.Root, nil)
+				if err != nil {
+					log.Warn("Trie prefetcher failed opening storage trie", "root", obj.data.Root, "err", err)
+					return err
+				}
+				obj.trie = tr
+
+				if tr == nil {
+					return nil
+				}
+				for slot := range obj.dirtyStorage {
+					slot := slot
+					// tr.GetStorage(addr, slot[:])
+					workers.Go(func() error {
+						_, err := tr.GetStorage(addr, slot[:])
+						return err
+					})
+				}
+			}
+			return nil
+		})
+	}
+	workers.Wait()
+
+	obj, exist := s.stateObjects[params.BeaconRootsAddress]
+	if exist {
+		tr, _ := s.db.OpenStorageTrie(s.originalRoot, params.BeaconRootsAddress, obj.data.Root, nil)
+		obj.trie = tr
+	}
+}
+
 func (s *StateDB) setRefund(gas uint64) {
 	s.refund = gas
 }
@@ -1179,31 +1227,45 @@ func (s *StateDB) Copy() *StateDB {
 
 func (s *StateDB) CopyState() *StateDB {
 	// Copy all the basic fields, initialize the memory ones
-	start := time.Now()
 	state := &StateDB{
 		db:                   s.db,
+		prefetcher:           s.prefetcher,
 		trie:                 s.trie,
 		reader:               s.reader,
 		originalRoot:         s.originalRoot,
 		stateObjects:         make(map[common.Address]*stateObject, len(s.stateObjects)),
 		stateObjectsDestruct: make(map[common.Address]*stateObject, len(s.stateObjectsDestruct)),
-		mutations:            make(map[common.Address]*mutation),
-		logs:                 make(map[common.Hash][]*types.Log),
-		preimages:            make(map[common.Hash][]byte),
-		journal:              newJournal(),
+		mutations:            make(map[common.Address]*mutation, len(s.mutations)),
+		dbErr:                s.dbErr,
+		refund:               s.refund,
+		thash:                s.thash,
+		txIndex:              s.txIndex,
+		logs:                 make(map[common.Hash][]*types.Log, len(s.logs)),
+		logSize:              s.logSize,
+		preimages:            maps.Clone(s.preimages),
+
+		// Do we need to copy the access list and transient storage?
+		// In practice: No. At the start of a transaction, these two lists are empty.
+		// In practice, we only ever copy state _between_ transactions/blocks, never
+		// in the middle of a transaction. However, it doesn't cost us much to copy
+		// empty lists, so we do it anyway to not blow up if we ever decide copy them
+		// in the middle of a transaction.
+		accessList:       s.accessList.Copy(),
+		transientStorage: s.transientStorage.Copy(),
+		journal:          s.journal.copy(),
+		// The update journal is not copies to avoid duplicated updates in later transactions.
+		blockNumber:   s.blockNumber,
 		updateJournal: newJournal(),
-		accessList:           newAccessList(),
-		transientStorage:     newTransientStorage(),
+		postSnapshot:  s.postSnapshot,
+		postVals:      s.postVals,
+	}
+	if s.witness != nil {
+		state.witness = s.witness.Copy()
+	}
+	if s.accessEvents != nil {
+		state.accessEvents = s.accessEvents.Copy()
 	}
-	StateNewTime += time.Since(start)
-	// if s.witness != nil {
-	// 	state.witness = s.witness.Copy()
-	// }
-	// if s.accessEvents != nil {
-	// 	state.accessEvents = s.accessEvents.Copy()
-	// }
 	// Deep copy cached state objects.
-	start = time.Now()
 	for addr, obj := range s.stateObjects {
 		state.stateObjects[addr] = obj.simpleCopy(state)
 	}
@@ -1211,20 +1273,19 @@ func (s *StateDB) CopyState() *StateDB {
 	for addr, obj := range s.stateObjectsDestruct {
 		state.stateObjectsDestruct[addr] = obj.simpleCopy(state)
 	}
-	StateDeepCpTime += time.Since(start)
 	// Deep copy the object state markers.
-	// for addr, op := range s.mutations {
-	// 	state.mutations[addr] = op.copy()
-	// }
+	for addr, op := range s.mutations {
+		state.mutations[addr] = op.copy()
+	}
 	// Deep copy the logs occurred in the scope of block
-	// for hash, logs := range s.logs {
-	// 	cpy := make([]*types.Log, len(logs))
-	// 	for i, l := range logs {
-	// 		cpy[i] = new(types.Log)
-	// 		*cpy[i] = *l
-	// 	}
-	// 	state.logs[hash] = cpy
-	// }
+	for hash, logs := range s.logs {
+		cpy := make([]*types.Log, len(logs))
+		for i, l := range logs {
+			cpy[i] = new(types.Log)
+			*cpy[i] = *l
+		}
+		state.logs[hash] = cpy
+	}
 	return state
 }
 

trie/trie.go

@@ -21,6 +21,7 @@ import (
 	"bytes"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -324,6 +325,8 @@ func (t *Trie) update(key, value []byte) error {
 	return nil
 }
 
+var ResolveTime = time.Duration(0)
+
 func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error) {
 	if len(key) == 0 {
 		if v, ok := n.(valueNode); ok {
@@ -387,7 +390,9 @@ func (t *Trie) insert(n node, prefix, key []byte, value node) (bool, node, error
 		// We've hit a part of the trie that isn't loaded yet. Load
 		// the node and insert into it. This leaves all child nodes on
 		// the path to the value in the trie.
+		start := time.Now()
 		rn, err := t.resolveAndTrack(n, prefix)
+		ResolveTime += time.Since(start)
 		if err != nil {
 			return false, nil, err
 		}
@@ -517,7 +522,10 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 				// shortNode{..., shortNode{...}}.  Since the entry
 				// might not be loaded yet, resolve it just for this
 				// check.
+				start := time.Now()
 				cnode, err := t.resolve(n.Children[pos], append(prefix, byte(pos)))
+				ResolveTime += time.Since(start)
+
 				if err != nil {
 					return false, nil, err
 				}
@@ -548,7 +556,9 @@ func (t *Trie) delete(n node, prefix, key []byte) (bool, node, error) {
 		// We've hit a part of the trie that isn't loaded yet. Load
 		// the node and delete from it. This leaves all child nodes on
 		// the path to the value in the trie.
+		start := time.Now()
 		rn, err := t.resolveAndTrack(n, prefix)
+		ResolveTime += time.Since(start)
 		if err != nil {
 			return false, nil, err
 		}