concurrent merge post state while executing

core/parallel_state_processor.go

@@ -60,7 +60,7 @@ func (p *ParallelStateProcessor) Process(block *types.Block, statedb *state.Stat
 		initialdb = statedb.Copy()
 
 		start = time.Now()
-		// Must prefetch bal before syscall to avoid overring syscall's state, thus merkle root might mismatch
+		// Must prefetch bal before syscall to avoid overriding syscall's state, thus merkle root might mismatch
 		statedb.PrefetchStateBAL(block.NumberU64())
 		PrefetchBALTime += time.Since(start)
 	} else {
@@ -95,9 +95,12 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 
 		preStateProvider PreStateProvider
 		workers          errgroup.Group
+		postState        = statedb.Copy()
 		// statedbbal       = statedb.Copy()
+		// wg sync.WaitGroup
 	)
-	workers.SetLimit(runtime.NumCPU() - 4)
+	// leave some cpus for prefetching
+	workers.SetLimit(runtime.NumCPU() / 2)
 
 	// Fetch prestate for each tx
 
@@ -129,6 +132,16 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	exeStart := time.Now()
 	postEntries := make([][]state.JournalEntry, len(block.Transactions()))
 
+	workers.Go(func() error {
+		start := time.Now()
+		postState.MergePostBalStates()
+		// prewarm the updating trie
+		postState.IntermediateRoot(true)
+		PostMergeTime += time.Since(start)
+
+		return nil
+	})
+
 	for i, tx := range block.Transactions() {
 		i := i
 		workers.Go(func() error {
@@ -177,6 +190,7 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	if err != nil {
 		return nil, err
 	}
+
 	ParallelExeTime += time.Since(exeStart)
 	// Merge state changes
 	// - Append receipts
@@ -188,17 +202,16 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	// 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 {
+	for i := range receipts {
+		receipt := receipts[i]
 		if receipt == nil {
 			continue // Skip nil receipts
 		}
 		receipt.CumulativeGasUsed = usedGas + receipt.GasUsed
 		usedGas += receipt.GasUsed
 		allLogs = append(allLogs, receipt.Logs...)
-		statedb.MergeState(postEntries[i])
+		// statedb.MergeState(postEntries[i])
 	}
-	PostMergeTime += time.Since(start)
 
 	// Read requests if Prague is enabled.
 	evm := vm.NewEVM(context, statedb, p.config, cfg)
@@ -222,6 +235,10 @@ func (p *ParallelStateProcessor) executeParallel(block *types.Block, statedb *st
 	// Finalize the block, applying any consensus engine specific extras (e.g. block rewards)
 	p.chain.engine.Finalize(p.chain, header, statedb, block.Body())
 
+	// Last tx alreadly includes the state change in requests after all txs
+	// wg.Wait()
+	*statedb = *postState
+
 	return &ProcessResult{
 		Receipts: receipts,
 		Requests: requests,

core/state/statedb.go

@@ -275,7 +275,11 @@ func NewWithReader(root common.Hash, db Database, reader Reader) (*StateDB, erro
 
 func (s *StateDB) PreComputePostState(blockNumber uint64) {
 	s.blockNumber = blockNumber
-	postBal := AllBlockTxPostValues[blockNumber]
+	postBal := map[int]types.TxPostValues{}
+	for k, v := range AllBlockTxPostValues[blockNumber] {
+		// Clone the map to avoid race with MergePostBalStates
+		postBal[k] = maps.Clone(v)
+	}
 
 	for i := 0; i < len(postBal)-1; i++ {
 		prevVals := postBal[i]
@@ -668,6 +672,52 @@ func (s *StateDB) MergeState(entries []JournalEntry) {
 	}
 }
 
+func (s *StateDB) MergePostBalStates() {
+	if balType != BalPreblockKeysPostValues {
+		panic("MergePostBal is only supported with BalPreblockKeysPostValues")
+	}
+	var (
+		postBal = AllBlockTxPostValues[s.blockNumber]
+	)
+
+	for txIndex := range len(postBal) {
+		postVals := postBal[txIndex]
+
+		for addr, acct := range postVals {
+			account := s.getStateObject(addr)
+
+			// Contract creation
+			if account == nil {
+				newAcct := types.NewEmptyStateAccount()
+				s.setStateObject(newObject(s, addr, newAcct))
+				account = s.getStateObject(addr)
+			}
+
+			if acct.Destruct {
+				s.SelfDestruct(addr)
+				continue
+			}
+
+			// When acct.Nonce is 0, it's possible nonce is just not changed.
+			// So we should only update nonce when it's larger than original nonce.
+			if account.Nonce() < acct.Nonce {
+				account.SetNonce(acct.Nonce)
+			}
+
+			if acct.Balance != nil {
+				account.SetBalance(acct.Balance)
+			}
+
+			if acct.Code != nil {
+				account.SetCode(crypto.Keccak256Hash(acct.Code), acct.Code)
+			}
+			for k, v := range acct.StorageKV {
+				account.SetState(k, v)
+			}
+		}
+	}
+}
+
 func (s *StateDB) setRefund(gas uint64) {
 	s.refund = gas
 }

trie/trie.go

@@ -157,6 +157,10 @@ func (t *Trie) Get(key []byte) ([]byte, error) {
 }
 
 func (t *Trie) get(origNode node, key []byte, pos int) (value []byte, newnode node, didResolve bool, err error) {
+	// When concurrent executing, sometime the following switch will throw error
+	if origNode == nil {
+		return nil, nil, false, nil
+	}
 	switch n := (origNode).(type) {
 	case nil:
 		return nil, nil, false, nil