allow parallel storage trie loading

Author: omerfirmak

cmd/migration-checker/main.go

@@ -199,11 +199,11 @@ 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
+	if !top {
+		workers = 1 << 3
 	}
 	step := byte(0xFF) / byte(workers)
 
@@ -219,9 +219,7 @@ func loadMPT(mptTrie *trie.SecureTrie, parallel bool) chan map[string][]byte {
 		go func() {
 			defer mptWg.Done()
 			for trieIt.Next() {
-				if parallel {
-					mptLeafMutex.Lock()
-				}
+				mptLeafMutex.Lock()
 
 				if _, ok := mptLeafMap[string(trieIt.Key)]; ok {
 					mptLeafMutex.Unlock()
@@ -230,11 +228,9 @@ func loadMPT(mptTrie *trie.SecureTrie, parallel bool) chan map[string][]byte {
 
 				mptLeafMap[string(dup(trieIt.Key))] = dup(trieIt.Value)
 
-				if parallel {
-					mptLeafMutex.Unlock()
-				}
+				mptLeafMutex.Unlock()
 
-				if parallel && len(mptLeafMap)%10000 == 0 {
+				if top && len(mptLeafMap)%10000 == 0 {
 					fmt.Println("MPT Accounts Loaded:", len(mptLeafMap))
 				}
 			}
@@ -249,7 +245,7 @@ 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 {
 	zkLeafMap := make(map[string][]byte, 1000)
 	var zkLeafMutex sync.Mutex
 	zkDone := make(chan map[string][]byte)
@@ -260,20 +256,16 @@ 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
 
-			if parallel {
-				zkLeafMutex.Unlock()
-			}
+			zkLeafMutex.Unlock()
 
-			if parallel && len(zkLeafMap)%10000 == 0 {
+			if top && len(zkLeafMap)%10000 == 0 {
 				fmt.Println("ZK Accounts Loaded:", len(zkLeafMap))
 			}
-		}, parallel, paranoid)
+		}, top, 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), top, 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(root, cb, 0, top, verifyNodeHashes)
 }
 
-func (t *ZkTrie) countLeaves(root *zkt.Hash, cb func(key, value []byte), depth int, parallel, verifyNodeHashes bool) uint64 {
+func (t *ZkTrie) countLeaves(root *zkt.Hash, cb func(key, value []byte), depth int, top, verifyNodeHashes bool) uint64 {
 	if root == nil {
 		return 0
 	}
@@ -271,19 +271,24 @@ 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 {
+		parallelismDepth := 5
+		if !top {
+			parallelismDepth = 3
+		}
+
+		if depth < parallelismDepth {
 			count := make(chan uint64)
 			leftT := t.Copy()
 			rightT := t.Copy()
 			go func() {
-				count <- leftT.countLeaves(rootNode.ChildL, cb, depth+1, parallel, verifyNodeHashes)
+				count <- leftT.countLeaves(rootNode.ChildL, cb, depth+1, top, verifyNodeHashes)
 			}()
 			go func() {
-				count <- rightT.countLeaves(rootNode.ChildR, cb, depth+1, parallel, verifyNodeHashes)
+				count <- rightT.countLeaves(rootNode.ChildR, cb, depth+1, top, verifyNodeHashes)
 			}()
 			return <-count + <-count
 		} else {
-			return t.countLeaves(rootNode.ChildL, cb, depth+1, parallel, verifyNodeHashes) + t.countLeaves(rootNode.ChildR, cb, depth+1, parallel, verifyNodeHashes)
+			return t.countLeaves(rootNode.ChildL, cb, depth+1, top, verifyNodeHashes) + t.countLeaves(rootNode.ChildR, cb, depth+1, top, verifyNodeHashes)
 		}
 	}
 }