easier parallelism

Author: omerfirmak

cmd/migration-checker/main.go

@@ -5,11 +5,8 @@ import (
 	"encoding/hex"
 	"flag"
 	"fmt"
-	"os"
-	"runtime"
 	"sort"
 	"sync"
-	"sync/atomic"
 	"time"
 
 	"github.com/scroll-tech/go-ethereum/common"
@@ -19,22 +16,18 @@ import (
 	"github.com/scroll-tech/go-ethereum/trie"
 )
 
-var accountsDone atomic.Uint64
-var trieCheckers chan struct{}
-
 type dbs struct {
 	zkDb  *leveldb.Database
 	mptDb *leveldb.Database
 }
 
 func main() {
 	var (
-		mptDbPath            = flag.String("mpt-db", "", "path to the MPT node DB")
-		zkDbPath             = flag.String("zk-db", "", "path to the ZK node DB")
-		mptRoot              = flag.String("mpt-root", "", "root hash of the MPT node")
-		zkRoot               = flag.String("zk-root", "", "root hash of the ZK node")
-		paranoid             = flag.Bool("paranoid", false, "verifies all node contents against their expected hash")
-		parallelismMultipler = flag.Int("parallelism-multiplier", 4, "multiplier for the number of parallel workers")
+		mptDbPath = flag.String("mpt-db", "", "path to the MPT node DB")
+		zkDbPath  = flag.String("zk-db", "", "path to the ZK node DB")
+		mptRoot   = flag.String("mpt-root", "", "root hash of the MPT node")
+		zkRoot    = flag.String("zk-root", "", "root hash of the ZK node")
+		paranoid  = flag.Bool("paranoid", false, "verifies all node contents against their expected hash")
 	)
 	flag.Parse()
 
@@ -45,35 +38,10 @@ func main() {
 
 	zkRootHash := common.HexToHash(*zkRoot)
 	mptRootHash := common.HexToHash(*mptRoot)
-
-	numTrieCheckers := runtime.GOMAXPROCS(0) * (*parallelismMultipler)
-	trieCheckers = make(chan struct{}, numTrieCheckers)
-	for i := 0; i < numTrieCheckers; i++ {
-		trieCheckers <- struct{}{}
-	}
-
-	done := make(chan struct{})
-	totalCheckers := len(trieCheckers)
-	go func() {
-		for {
-			select {
-			case <-done:
-				return
-			case <-time.After(time.Minute):
-				fmt.Println("Active checkers:", totalCheckers-len(trieCheckers))
-			}
-		}
-	}()
-	defer close(done)
-
 	checkTrieEquality(&dbs{
 		zkDb:  zkDb,
 		mptDb: mptDb,
 	}, zkRootHash, mptRootHash, "", checkAccountEquality, true, *paranoid)
-
-	for i := 0; i < numTrieCheckers; i++ {
-		<-trieCheckers
-	}
 }
 
 func panicOnError(err error, label, msg string) {
@@ -120,6 +88,7 @@ func checkTrieEquality(dbs *dbs, zkRoot, mptRoot common.Hash, label string, leaf
 	for index, zkKv := range zkLeafs {
 		mptKv := mptLeafs[index]
 		leafChecker(fmt.Sprintf("%s key: %s", label, hex.EncodeToString([]byte(zkKv.key))), dbs, zkKv.value, mptKv.value, paranoid)
+		fmt.Println("Accounts done:", index+1, "/", len(zkLeafs))
 	}
 }
 
@@ -146,23 +115,7 @@ func checkAccountEquality(label string, dbs *dbs, zkAccountBytes, mptAccountByte
 	} else if zkAccount.Root != (common.Hash{}) {
 		zkRoot := common.BytesToHash(zkAccount.Root[:])
 		mptRoot := common.BytesToHash(mptAccount.Root[:])
-		<-trieCheckers
-		go func() {
-			defer func() {
-				if p := recover(); p != nil {
-					fmt.Println(p)
-					os.Exit(1)
-				}
-			}()
-
-			checkTrieEquality(dbs, zkRoot, mptRoot, label, checkStorageEquality, false, paranoid)
-			accountsDone.Add(1)
-			fmt.Println("Accounts done:", accountsDone.Load())
-			trieCheckers <- struct{}{}
-		}()
-	} else {
-		accountsDone.Add(1)
-		fmt.Println("Accounts done:", accountsDone.Load())
+		checkTrieEquality(dbs, zkRoot, mptRoot, label, checkStorageEquality, false, paranoid)
 	}
 }
 
@@ -183,9 +136,6 @@ type kv struct {
 func loadMPT(mptTrie *trie.SecureTrie, top bool) chan []kv {
 	startKey := make([]byte, 32)
 	workers := 1 << 5
-	if !top {
-		workers = 1 << 3
-	}
 	step := byte(256 / workers)
 
 	mptLeafs := make([]kv, 0, 1000)

trie/zk_trie.go

@@ -272,10 +272,6 @@ func (t *ZkTrie) countLeaves(root *zkt.Hash, cb func(key, value []byte), depth i
 		return 1
 	} else {
 		parallelismDepth := 5
-		if !top {
-			parallelismDepth = 3
-		}
-
 		if depth < parallelismDepth {
 			count := make(chan uint64)
 			leftT := t.Copy()