try better parallelization

Author: omerfirmak

cmd/migration-checker/main.go

@@ -24,7 +24,8 @@ import (
 )
 
 var accountsDone atomic.Uint64
-var trieCheckers chan struct{}
+var numLoaders int
+var trieLoaders chan struct{}
 
 type dbs struct {
 	zkDb  *leveldb.Database
@@ -54,33 +55,19 @@ 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{}{}
+	numLoaders = runtime.GOMAXPROCS(0) * (*parallelismMultipler)
+	trieLoaders = make(chan struct{}, numLoaders)
+	for i := 0; i < numLoaders; i++ {
+		trieLoaders <- 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
+	for i := 0; i < numLoaders; i++ {
+		<-trieLoaders
 	}
 }
 
@@ -177,8 +164,9 @@ 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
+		<-trieLoaders
 		go func() {
+			trieLoaders <- struct{}{}
 			defer func() {
 				if p := recover(); p != nil {
 					fmt.Println(p)
@@ -189,7 +177,6 @@ func checkAccountEquality(label string, dbs *dbs, zkAccountBytes, mptAccountByte
 			checkTrieEquality(dbs, zkRoot, mptRoot, label, checkStorageEquality, false, paranoid)
 			accountsDone.Add(1)
 			fmt.Println("Accounts done:", accountsDone.Load())
-			trieCheckers <- struct{}{}
 		}()
 	} else {
 		accountsDone.Add(1)
@@ -207,42 +194,33 @@ func checkStorageEquality(label string, _ *dbs, zkStorageBytes, mptStorageBytes
 	}
 }
 
-func loadMPT(mptTrie *trie.SecureTrie, parallel bool) chan map[string][]byte {
+func loadMPT(mptTrie *trie.SecureTrie, top bool) chan map[string][]byte {
 	startKey := make([]byte, 32)
-	workers := 1 << 5
-	if !parallel {
-		workers = 1
-	}
-	step := byte(0xFF) / byte(workers)
-
 	mptLeafMap := make(map[string][]byte, 1000)
 	var mptLeafMutex sync.Mutex
 
 	var mptWg sync.WaitGroup
-	for i := 0; i < workers; i++ {
-		startKey[0] = byte(i) * step
+	for i := 0; i < 255; i++ {
+		startKey[0] = byte(i)
 		trieIt := trie.NewIterator(mptTrie.NodeIterator(startKey))
 
 		mptWg.Add(1)
+		<-trieLoaders
 		go func() {
-			defer mptWg.Done()
+			defer func() {
+				mptWg.Done()
+				trieLoaders <- struct{}{}
+			}()
 			for trieIt.Next() {
-				if parallel {
-					mptLeafMutex.Lock()
-				}
-
+				mptLeafMutex.Lock()
 				if _, ok := mptLeafMap[string(trieIt.Key)]; ok {
 					mptLeafMutex.Unlock()
 					break
 				}
-
 				mptLeafMap[string(dup(trieIt.Key))] = dup(trieIt.Value)
+				mptLeafMutex.Unlock()
 
-				if parallel {
-					mptLeafMutex.Unlock()
-				}
-
-				if parallel && len(mptLeafMap)%10000 == 0 {
+				if top && len(mptLeafMap)%10000 == 0 {
 					fmt.Println("MPT Accounts Loaded:", len(mptLeafMap))
 				}
 			}
@@ -257,7 +235,8 @@ func loadMPT(mptTrie *trie.SecureTrie, parallel bool) chan map[string][]byte {
 	return respChan
 }
 
-func loadZkTrie(zkTrie *trie.ZkTrie, parallel, paranoid bool) chan map[string][]byte {
+func loadZkTrie(zkTrie *trie.ZkTrie, top, paranoid bool) chan map[string][]byte {
+	parallelismCutoffDepth := 8
 	zkLeafMap := make(map[string][]byte, 1000)
 	var zkLeafMutex sync.Mutex
 	zkDone := make(chan map[string][]byte)
@@ -268,20 +247,20 @@ func loadZkTrie(zkTrie *trie.ZkTrie, parallel, paranoid bool) chan map[string][]
 				panic(fmt.Sprintf("preimage not found zk trie %s", hex.EncodeToString(key)))
 			}
 
-			if parallel {
-				zkLeafMutex.Lock()
-			}
-
+			zkLeafMutex.Lock()
 			zkLeafMap[string(dup(preimageKey))] = value
+			zkLeafMutex.Unlock()
 
-			if parallel {
-				zkLeafMutex.Unlock()
-			}
-
-			if parallel && len(zkLeafMap)%10000 == 0 {
+			if top && len(zkLeafMap)%10000 == 0 {
 				fmt.Println("ZK Accounts Loaded:", len(zkLeafMap))
 			}
-		}, parallel, paranoid)
+		}, func(f func()) {
+			<-trieLoaders
+			go func() {
+				f()
+				trieLoaders <- struct{}{}
+			}()
+		}, parallelismCutoffDepth, paranoid)
 		zkDone <- zkLeafMap
 	}()
 	return zkDone

trie/zk_trie.go

@@ -239,15 +239,15 @@ func (t *ZkTrie) Witness() map[string]struct{} {
 	panic("not implemented")
 }
 
-func (t *ZkTrie) CountLeaves(cb func(key, value []byte), parallel, verifyNodeHashes bool) uint64 {
+func (t *ZkTrie) CountLeaves(cb func(key, value []byte), spawnWorker func(func()), parallelismCutoffDepth int, verifyNodeHashes bool) uint64 {
 	root, err := t.ZkTrie.Tree().Root()
 	if err != nil {
 		panic("CountLeaves cannot get root")
 	}
-	return t.countLeaves(root, cb, 0, parallel, verifyNodeHashes)
+	return t.countLeaves(0, root, cb, spawnWorker, parallelismCutoffDepth, verifyNodeHashes)
 }
 
-func (t *ZkTrie) countLeaves(root *zkt.Hash, cb func(key, value []byte), depth int, parallel, verifyNodeHashes bool) uint64 {
+func (t *ZkTrie) countLeaves(depth int, root *zkt.Hash, cb func(key, value []byte), spawnWorker func(func()), parallelismCutoffDepth int, verifyNodeHashes bool) uint64 {
 	if root == nil {
 		return 0
 	}
@@ -271,19 +271,19 @@ func (t *ZkTrie) countLeaves(root *zkt.Hash, cb func(key, value []byte), depth i
 		cb(append([]byte{}, rootNode.NodeKey.Bytes()...), append([]byte{}, rootNode.Data()...))
 		return 1
 	} else {
-		if parallel && depth < 5 {
+		if depth < parallelismCutoffDepth {
 			count := make(chan uint64)
 			leftT := t.Copy()
 			rightT := t.Copy()
 			go func() {
-				count <- leftT.countLeaves(rootNode.ChildL, cb, depth+1, parallel, verifyNodeHashes)
+				spawnWorker(func() {
+					count <- leftT.countLeaves(depth+1, rootNode.ChildL, cb, spawnWorker, parallelismCutoffDepth, verifyNodeHashes)
+				})
 			}()
-			go func() {
-				count <- rightT.countLeaves(rootNode.ChildR, cb, depth+1, parallel, verifyNodeHashes)
-			}()
-			return <-count + <-count
+			return rightT.countLeaves(depth+1, rootNode.ChildR, cb, spawnWorker, parallelismCutoffDepth, verifyNodeHashes) + <-count
 		} else {
-			return t.countLeaves(rootNode.ChildL, cb, depth+1, parallel, verifyNodeHashes) + t.countLeaves(rootNode.ChildR, cb, depth+1, parallel, verifyNodeHashes)
+			return t.countLeaves(depth+1, rootNode.ChildL, cb, spawnWorker, parallelismCutoffDepth, verifyNodeHashes) +
+				t.countLeaves(depth+1, rootNode.ChildR, cb, spawnWorker, parallelismCutoffDepth, verifyNodeHashes)
 		}
 	}
 }