fix: an attempt at making witness generation more stable (#1166)

Author: omerfirmak

core/blockchain.go

@@ -1699,7 +1699,7 @@ func (bc *BlockChain) insertChain(chain types.Blocks, verifySeals bool) (int, er
 		}
 
 		// Enable prefetching to pull in trie node paths while processing transactions
-		statedb.StartPrefetcher("chain", nil)
+		statedb.StartPrefetcher("chain")
 		activeState = statedb
 
 		// If we have a followup block, run that against the current state to pre-cache
@@ -1836,7 +1836,7 @@ func (bc *BlockChain) BuildAndWriteBlock(parentBlock *types.Block, header *types
 		return nil, NonStatTy, err
 	}
 
-	statedb.StartPrefetcher("l1sync", nil)
+	statedb.StartPrefetcher("l1sync")
 	defer statedb.StopPrefetcher()
 
 	header.ParentHash = parentBlock.Hash()

core/state/state_object.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"io"
 	"math/big"
+	"slices"
 	"time"
 
 	"github.com/scroll-tech/go-ethereum/common"
@@ -344,8 +345,17 @@ func (s *stateObject) updateTrie(db Database) Trie {
 	tr := s.getTrie(db)
 	hasher := s.db.hasher
 
+	sortedPendingStorages := make([]common.Hash, 0, len(s.pendingStorage))
+	for key := range s.pendingStorage {
+		sortedPendingStorages = append(sortedPendingStorages, key)
+	}
+	slices.SortFunc(sortedPendingStorages, func(a, b common.Hash) int {
+		return bytes.Compare(a[:], b[:])
+	})
+
 	usedStorage := make([][]byte, 0, len(s.pendingStorage))
-	for key, value := range s.pendingStorage {
+	for _, key := range sortedPendingStorages {
+		value := s.pendingStorage[key]
 		// Skip noop changes, persist actual changes
 		if value == s.originStorage[key] {
 			continue

core/state/statedb.go

@@ -18,9 +18,11 @@
 package state
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"math/big"
+	"slices"
 	"sort"
 	"time"
 
@@ -160,18 +162,20 @@ func New(root common.Hash, db Database, snaps *snapshot.Tree) (*StateDB, error)
 	return sdb, nil
 }
 
+// WithWitness sets the witness for the state database.
+func (s *StateDB) WithWitness(witness *stateless.Witness) {
+	s.witness = witness
+}
+
 // StartPrefetcher initializes a new trie prefetcher to pull in nodes from the
 // state trie concurrently while the state is mutated so that when we reach the
 // commit phase, most of the needed data is already hot.
-func (s *StateDB) StartPrefetcher(namespace string, witness *stateless.Witness) {
+func (s *StateDB) StartPrefetcher(namespace string) {
 	if s.prefetcher != nil {
 		s.prefetcher.close()
 		s.prefetcher = nil
 	}
 
-	// Enable witness collection if requested
-	s.witness = witness
-
 	if s.snap != nil {
 		s.prefetcher = newTriePrefetcher(s.db, s.originalRoot, namespace)
 	}
@@ -967,8 +971,17 @@ func (s *StateDB) IntermediateRoot(deleteEmptyObjects bool) common.Hash {
 			s.trie = trie
 		}
 	}
-	usedAddrs := make([][]byte, 0, len(s.stateObjectsPending))
+
+	sortedPendingAddrs := make([]common.Address, 0, len(s.stateObjectsPending))
 	for addr := range s.stateObjectsPending {
+		sortedPendingAddrs = append(sortedPendingAddrs, addr)
+	}
+	slices.SortFunc(sortedPendingAddrs, func(a, b common.Address) int {
+		return bytes.Compare(a[:], b[:])
+	})
+
+	usedAddrs := make([][]byte, 0, len(s.stateObjectsPending))
+	for _, addr := range sortedPendingAddrs {
 		if obj := s.stateObjects[addr]; obj.deleted {
 			s.deleteStateObject(obj)
 			s.AccountDeleted += 1

eth/api.go

@@ -29,6 +29,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/syndtr/goleveldb/leveldb"
+
 	"github.com/scroll-tech/go-ethereum/common"
 	"github.com/scroll-tech/go-ethereum/common/hexutil"
 	"github.com/scroll-tech/go-ethereum/core"
@@ -44,7 +46,6 @@ import (
 	"github.com/scroll-tech/go-ethereum/rollup/rcfg"
 	"github.com/scroll-tech/go-ethereum/rpc"
 	"github.com/scroll-tech/go-ethereum/trie"
-	"github.com/syndtr/goleveldb/leveldb"
 )
 
 // PublicEthereumAPI provides an API to access Ethereum full node-related
@@ -354,13 +355,11 @@ func generateWitness(blockchain *core.BlockChain, block *types.Block) (*stateles
 	if err != nil {
 		return nil, fmt.Errorf("failed to retrieve parent state: %w", err)
 	}
+	statedb.WithWitness(witness)
 
 	// Collect storage locations that prover needs but sequencer might not touch necessarily
 	statedb.GetState(rcfg.L2MessageQueueAddress, rcfg.WithdrawTrieRootSlot)
 
-	statedb.StartPrefetcher("debug_execution_witness", witness)
-	defer statedb.StopPrefetcher()
-
 	receipts, _, usedGas, err := blockchain.Processor().Process(block, statedb, *blockchain.GetVMConfig())
 	if err != nil {
 		return nil, fmt.Errorf("failed to process block %d: %w", block.Number(), err)
@@ -370,16 +369,10 @@ func generateWitness(blockchain *core.BlockChain, block *types.Block) (*stateles
 		return nil, fmt.Errorf("failed to validate block %d: %w", block.Number(), err)
 	}
 
-	// FIXME: testWitness will fail from time to time, the problem is caused by occasional state root mismatch
-	// after processing the block based on witness. We need to investigate the root cause and fix it.
-	for retries := 0; retries < 5; retries++ {
-		if err = testWitness(blockchain, block, witness); err == nil {
-			return witness, nil
-		} else {
-			log.Warn("Failed to validate witness", "block", block.Number(), "error", err)
-		}
+	if err = testWitness(blockchain, block, witness); err != nil {
+		return nil, err
 	}
-	return witness, err
+	return witness, nil
 }
 
 func testWitness(blockchain *core.BlockChain, block *types.Block, witness *stateless.Witness) error {

params/version.go

@@ -24,7 +24,7 @@ import (
 const (
 	VersionMajor = 5         // Major version component of the current release
 	VersionMinor = 8         // Minor version component of the current release
-	VersionPatch = 35        // Patch version component of the current release
+	VersionPatch = 36        // Patch version component of the current release
 	VersionMeta  = "mainnet" // Version metadata to append to the version string
 )
 

rollup/pipeline/pipeline.go

@@ -228,7 +228,7 @@ func sendCancellable[T any, C comparable](resCh chan T, msg T, cancelCh <-chan C
 }
 
 func (p *Pipeline) traceAndApplyStage(txsIn <-chan *types.Transaction) (<-chan error, <-chan *BlockCandidate, error) {
-	p.state.StartPrefetcher("miner", nil)
+	p.state.StartPrefetcher("miner")
 	downstreamCh := make(chan *BlockCandidate, p.downstreamChCapacity())
 	resCh := make(chan error)
 	p.wg.Add(1)