perf: lock free account/slot reads & add pre-block BAL prefetch

core/blockchain.go

@@ -1902,18 +1902,18 @@ func (bc *BlockChain) processBlock(parentRoot common.Hash, block *types.Block, s
 
 	bc.cacheConfig.TrieCleanNoPrefetch = true
 	if bc.cacheConfig.TrieCleanNoPrefetch {
-		// statedb, err = state.New(parentRoot, bc.statedb)
+		statedb, err = state.New(parentRoot, bc.statedb)
+		if err != nil {
+			return nil, err
+		}
+		// reader, err := bc.statedb.ReaderWithCache(parentRoot)
+		// if err != nil {
+		// 	return nil, err
+		// }
+		// statedb, err = state.NewWithReader(parentRoot, bc.statedb, reader)
 		// if err != nil {
 		// 	return nil, err
 		// }
-		reader, err := bc.statedb.ReaderWithCache(parentRoot)
-		if err != nil {
-			return nil, err
-		}
-		statedb, err = state.NewWithReader(parentRoot, bc.statedb, reader)
-		if err != nil {
-			return nil, err
-		}
 	} else {
 		fmt.Println("prefetch enabled===========================")
 		// If prefetching is enabled, run that against the current state to pre-cache

core/parallel_state_processor.go

@@ -48,13 +48,23 @@ func (p *ParallelStateProcessor) Process(block *types.Block, statedb *state.Stat
 	var (
 		context   vm.BlockContext
 		signer    = types.MakeSigner(p.config, header.Number, header.Time)
+		initialdb *state.StateDB
 	)
 
 	if preStateType == BALPreState {
 		start := time.Now()
-		// statedb.PrefetchStateBAL(block.NumberU64())
 		statedb.PreComputePostState(block.NumberU64())
+		PrefetchMergeBALTime += time.Since(start)
+
+		// copy initialdb before PrefetchStateBAL to avoid redundant copy
+		initialdb = statedb.Copy()
+
+		start = time.Now()
+		// Must prefetch bal before syscall or merkle root might mismatch
+		statedb.PrefetchStateBAL(block.NumberU64())
 		PrefetchBALTime += time.Since(start)
+	} else {
+		initialdb = statedb.Copy()
 	}
 
 	// Apply pre-execution system calls.
@@ -72,10 +82,10 @@ func (p *ParallelStateProcessor) Process(block *types.Block, statedb *state.Stat
 		ProcessParentBlockHash(block.ParentHash(), evm)
 	}
 
-	return p.executeParallel(block, statedb, cfg, gp, signer, context)
+	return p.executeParallel(block, statedb, cfg, gp, signer, context, initialdb)
 }
 
-func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *state.StateDB, cfg vm.Config, gp *GasPool, signer types.Signer, context vm.BlockContext) (*ProcessResult, error) {
+func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *state.StateDB, cfg vm.Config, gp *GasPool, signer types.Signer, context vm.BlockContext, initialdb *state.StateDB) (*ProcessResult, error) {
 	var (
 		receipts    = make(types.Receipts, len(block.Transactions()))
 		header      = block.Header()
@@ -85,17 +95,19 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 
 		preStateProvider PreStateProvider
 		workers          errgroup.Group
+		// statedbbal       = statedb.Copy()
 	)
-	workers.SetLimit(runtime.NumCPU() - 6)
+	workers.SetLimit(runtime.NumCPU() - 4)
+
 	// Fetch prestate for each tx
 
 	// todo: handle gp with RW lock
 	switch preStateType {
 	case BALPreState:
 		{
-			start := time.Now()
-			// statedb.MergePostBal()
-			PrefetchMergeBALTime += time.Since(start)
+			// start := time.Now()
+			// statedb.PrefetchStateBAL(block.NumberU64())
+			// PrefetchBALTime += time.Since(start)
 		}
 
 	case SeqPreState:
@@ -117,10 +129,8 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	exeStart := time.Now()
 	postEntries := make([][]state.JournalEntry, len(block.Transactions()))
 
-	initialdb := statedb.Copy()
 	for i, tx := range block.Transactions() {
 		i := i
-
 		workers.Go(func() error {
 			var (
 				cleanStatedb *state.StateDB
@@ -131,7 +141,7 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 				{
 					cleanStatedb = initialdb.Copy()
 					cleanStatedb.SetTxContext(tx.Hash(), i)
-					err = cleanStatedb.SetTxBALReader()
+					err = cleanStatedb.SetTxBALReader(statedb)
 				}
 			case SeqPreState:
 				{
@@ -175,6 +185,9 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	// - Ommit preimages for now
 	usedGas := uint64(0)
 
+	// set it to avoid read bal post state, -2 is a magic tx number
+	statedb.SetTxContext(common.Hash{}, -2)
+
 	start := time.Now()
 	for i, receipt := range receipts {
 		if receipt == nil {

core/state/reader.go

@@ -71,6 +71,9 @@ type StateReader interface {
 	// - Returns an error only if an unexpected issue occurs
 	// - The returned storage slot is safe to modify after the call
 	Storage(addr common.Address, slot common.Hash) (common.Hash, error)
+
+	AccountBAL(addr common.Address) (*types.StateAccount, error)
+	StorageBAL(addr common.Address, slot common.Hash, tr *trie.StateTrie) (common.Hash, error)
 }
 
 // Reader defines the interface for accessing accounts, storage slots and contract
@@ -172,6 +175,10 @@ func (r *flatReader) Account(addr common.Address) (*types.StateAccount, error) {
 	return acct, nil
 }
 
+func (r *flatReader) AccountBAL(addr common.Address) (*types.StateAccount, error) {
+	return r.Account(addr)
+}
+
 // Storage implements StateReader, retrieving the storage slot specified by the
 // address and slot key.
 //
@@ -200,6 +207,10 @@ func (r *flatReader) Storage(addr common.Address, key common.Hash) (common.Hash,
 	return value, nil
 }
 
+func (r *flatReader) StorageBAL(addr common.Address, key common.Hash, tr *trie.StateTrie) (common.Hash, error) {
+	return r.Storage(addr, key)
+}
+
 // trieReader implements the StateReader interface, providing functions to access
 // state from the referenced trie.
 //
@@ -268,6 +279,14 @@ func (r *trieReader) Account(addr common.Address) (*types.StateAccount, error) {
 	return r.account(addr)
 }
 
+func (r *trieReader) AccountBAL(addr common.Address) (*types.StateAccount, error) {
+	account, err := r.mainTrie.GetAccount(addr)
+	if err != nil {
+		return nil, err
+	}
+	return account, nil
+}
+
 // Storage implements StateReader, retrieving the storage slot specified by the
 // address and slot key.
 //
@@ -314,6 +333,15 @@ func (r *trieReader) Storage(addr common.Address, key common.Hash) (common.Hash,
 	return value, nil
 }
 
+func (r *trieReader) StorageBAL(addr common.Address, key common.Hash, tr *trie.StateTrie) (common.Hash, error) {
+	ret, err := tr.GetStorage(addr, key.Bytes())
+	if err != nil {
+		return common.Hash{}, err
+	}
+
+	return common.BytesToHash(ret), nil
+}
+
 // multiStateReader is the aggregation of a list of StateReader interface,
 // providing state access by leveraging all readers. The checking priority
 // is determined by the position in the reader list.
@@ -354,6 +382,18 @@ func (r *multiStateReader) Account(addr common.Address) (*types.StateAccount, er
 	return nil, errors.Join(errs...)
 }
 
+func (r *multiStateReader) AccountBAL(addr common.Address) (*types.StateAccount, error) {
+	var errs []error
+	for _, reader := range r.readers {
+		acct, err := reader.AccountBAL(addr)
+		if err == nil {
+			return acct, nil
+		}
+		errs = append(errs, err)
+	}
+	return nil, errors.Join(errs...)
+}
+
 // Storage implementing StateReader interface, retrieving the storage slot
 // associated with a particular account address and slot key.
 //
@@ -372,6 +412,18 @@ func (r *multiStateReader) Storage(addr common.Address, slot common.Hash) (commo
 	return common.Hash{}, errors.Join(errs...)
 }
 
+func (r *multiStateReader) StorageBAL(addr common.Address, slot common.Hash, tr *trie.StateTrie) (common.Hash, error) {
+	var errs []error
+	for _, reader := range r.readers {
+		slot, err := reader.StorageBAL(addr, slot, tr)
+		if err == nil {
+			return slot, nil
+		}
+		errs = append(errs, err)
+	}
+	return common.Hash{}, errors.Join(errs...)
+}
+
 // reader is the wrapper of ContractCodeReader and StateReader interface.
 type reader struct {
 	ContractCodeReader
@@ -574,16 +626,16 @@ func (r *readerWithBAL) Storage(addr common.Address, slot common.Hash) (common.H
 	}
 
 	val, err := r.Reader.Storage(addr, slot)
-	// Cache storage
-	if postVals != nil && postVals[addr] != nil {
-		if postVals[addr].StorageKV != nil {
-			postVals[addr].StorageKV[slot] = val
-		} else {
-			postVals[addr].StorageKV = map[common.Hash]common.Hash{
-				slot: val,
-			}
-		}
-	}
+	// Don't Cache storage due to concurrent map writes because postVals might share the same map.
+	// if postVals != nil && postVals[addr] != nil {
+	// 	if postVals[addr].StorageKV != nil {
+	// 		postVals[addr].StorageKV[slot] = val
+	// 	} else {
+	// 		postVals[addr].StorageKV = map[common.Hash]common.Hash{
+	// 			slot: val,
+	// 		}
+	// 	}
+	// }
 	if err != nil {
 		return common.Hash{}, err
 	}

core/state/statedb.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"maps"
 	"os"
+	"runtime"
 	"slices"
 	"sync"
 	"sync/atomic"
@@ -304,14 +305,18 @@ func (s *StateDB) PreComputePostState(blockNumber uint64) {
 	s.postBAL = postBal
 }
 
-func (s *StateDB) SetTxBALReader() error {
+func (s *StateDB) SetTxBALReader(preblockReader Reader) error {
 	if s.postBAL == nil {
 		return errors.New("cannot set bal reader without postBAL")
 	}
-	s.reader = newReaderWithBAL(s.reader, s.txIndex, s.postBAL)
+	s.reader = newReaderWithBAL(preblockReader, s.txIndex, s.postBAL)
 	return nil
 }
 
+func (s *StateDB) SetOriginalReader(reader Reader) {
+	s.reader = reader
+}
+
 func (s *StateDB) PrefetchStateBAL(blockNumber uint64) {
 	s.blockNumber = blockNumber
 	switch balType {
@@ -331,19 +336,20 @@ func (s *StateDB) prefetchBalPreblockKeys() {
 	}
 	var (
 		preBal      = AllBlockAccessLists[s.blockNumber]
+		lenMaxSlots = 0
 		lenSlots    = 0
 		workers     errgroup.Group
 	)
 
-	workers.SetLimit(25)
+	workers.SetLimit(runtime.NumCPU() - 2)
 	// Fetch pre-block acccount state
 	accounts := make(chan *stateObject, len(preBal))
 	for _, acl := range preBal {
 		addr := acl.Address
-		lenSlots += len(acl.StorageKeys)
+		lenMaxSlots += len(acl.StorageKeys)
 
 		workers.Go(func() error {
-			acct, err := s.reader.Account(addr)
+			acct, err := s.reader.AccountBAL(addr)
 			obj := newObject(s, addr, acct)
 			accounts <- obj
 			if err != nil {
@@ -361,14 +367,31 @@ func (s *StateDB) prefetchBalPreblockKeys() {
 		s.setStateObject(obj)
 	}
 
+	close(accounts)
+
 	// Fetch pre-block storage state
-	storages := make(chan *StorageKV, lenSlots)
+	storages := make(chan *StorageKV, lenMaxSlots)
 	for _, acl := range preBal {
 		addr := acl.Address
 		keys := acl.StorageKeys
+
+		obj := s.stateObjects[addr]
+
+		if obj.origin == nil {
+			continue
+		}
+
+		lenSlots += len(acl.StorageKeys)
+
+		tr, err := trie.NewStateTrie(trie.StorageTrieID(s.originalRoot, obj.addrHash, obj.origin.Root), s.db.TrieDB())
+		if err != nil {
+			log.Error("fail to create trie for:", addr)
+			panic("fail to create trie for:")
+		}
+
 		for _, key := range keys {
 			workers.Go(func() error {
-				val, err := s.reader.Storage(addr, key)
+				val, err := s.reader.StorageBAL(addr, key, tr)
 				kv := &StorageKV{&addr, &key, &val}
 				storages <- kv
 				if err != nil {
@@ -389,6 +412,8 @@ func (s *StateDB) prefetchBalPreblockKeys() {
 			s.setStateObject(account)
 		}
 	}
+
+	close(storages)
 }
 
 var (
@@ -635,7 +660,8 @@ func (s *StateDB) MergeState(entries []JournalEntry) {
 		case codeChange:
 			s.SetCode(v.account, v.prevCode)
 		case refundChange:
-			s.setRefund(v.prev)
+			// refundchange doesn't affect stateroot
+			// s.setRefund(v.prev)
 		default:
 
 		}
@@ -1015,7 +1041,7 @@ func (s *StateDB) CreateAccount(addr common.Address) {
 // correctly handle EIP-6780 'delete-in-same-transaction' logic.
 func (s *StateDB) CreateContract(addr common.Address) {
 	obj := s.getStateObject(addr)
-	if !obj.newContract {
+	if obj != nil && !obj.newContract {
 		obj.newContract = true
 		s.journal.createContract(addr)
 		s.updateJournal.createContract(addr)

core/state/statedb_readerbal.go

@@ -0,0 +1,37 @@
+package state
+
+import (
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/trie"
+)
+
+func (s *StateDB) Account(addr common.Address) (*types.StateAccount, error) {
+	// s.getStateObject() shouldn't be used, cause when acct is nil, it'll load account from DB and result concurrent map writes for accts
+	obj := s.stateObjects[addr]
+	return &obj.data, nil
+}
+
+func (s *StateDB) AccountBAL(addr common.Address) (*types.StateAccount, error) {
+	panic("AccountBAL not implemented for statedb")
+}
+
+func (s *StateDB) Code(addr common.Address, codeHash common.Hash) ([]byte, error) {
+	return s.reader.Code(addr, codeHash)
+}
+
+func (s *StateDB) CodeSize(addr common.Address, codeHash common.Hash) (int, error) {
+	return s.reader.CodeSize(addr, codeHash)
+}
+
+func (s *StateDB) Storage(addr common.Address, slot common.Hash) (common.Hash, error) {
+	stateObject := s.stateObjects[addr]
+	if stateObject != nil {
+		return stateObject.GetState(slot), nil
+	}
+	return common.Hash{}, nil
+}
+
+func (s *StateDB) StorageBAL(addr common.Address, slot common.Hash, tr *trie.StateTrie) (common.Hash, error) {
+	panic("StorageBAL not implemented for statedb")
+}

trie/secure_trie.go

@@ -19,6 +19,7 @@ package trie
 import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/trie/trienode"
 	"github.com/ethereum/go-ethereum/triedb/database"
@@ -105,7 +106,8 @@ func (t *StateTrie) MustGet(key []byte) []byte {
 // If the specified storage slot is not in the trie, nil will be returned.
 // If a trie node is not found in the database, a MissingNodeError is returned.
 func (t *StateTrie) GetStorage(_ common.Address, key []byte) ([]byte, error) {
-	enc, err := t.trie.Get(t.hashKey(key))
+	// enc, err := t.trie.Get(t.hashKey(key))
+	enc, err := t.trie.Get(crypto.Keccak256(key))
 	if err != nil || len(enc) == 0 {
 		return nil, err
 	}
@@ -117,7 +119,8 @@ func (t *StateTrie) GetStorage(_ common.Address, key []byte) ([]byte, error) {
 // If the specified account is not in the trie, nil will be returned.
 // If a trie node is not found in the database, a MissingNodeError is returned.
 func (t *StateTrie) GetAccount(address common.Address) (*types.StateAccount, error) {
-	res, err := t.trie.Get(t.hashKey(address.Bytes()))
+	// res, err := t.trie.Get(t.hashKey(address.Bytes()))
+	res, err := t.trie.Get(crypto.Keccak256(address.Bytes()))
 	if res == nil || err != nil {
 		return nil, err
 	}

trie/trie.go

@@ -616,7 +616,7 @@ func (t *Trie) resolveAndTrack(n hashNode, prefix []byte) (node, error) {
 	if err != nil {
 		return nil, err
 	}
-	t.tracer.onRead(prefix, blob)
+	// t.tracer.onRead(prefix, blob)
 
 	// The returned node blob won't be changed afterward. No need to
 	// deep-copy the slice.