separate dbs

Author: omerfirmak

cmd/migration-checker/main.go

@@ -26,7 +26,7 @@ import (
 
 var accountsDone atomic.Uint64
 var numLoaders int
-var trieLoaders chan struct{}
+var trieLoaders chan dbs
 
 type dbs struct {
 	zkDb  *leveldb.Database
@@ -48,40 +48,37 @@ func main() {
 		log.Println(http.ListenAndServe("0.0.0.0:6060", nil))
 	}()
 
-	zkDb, err := leveldb.NewCustom(*zkDbPath, "zkDb", func(options *opt.Options) {
-		// Set default options
-		options.ReadOnly = true
-		options.BlockCacher = opt.NoCacher
-		options.BlockCacheCapacity = 0
-		options.OpenFilesCacher = opt.NoCacher
-		options.OpenFilesCacheCapacity = 0
-		options.DisableBlockCache = true
-	})
-	panicOnError(err, "", "failed to open zk db")
-	mptDb, err := leveldb.NewCustom(*mptDbPath, "mptDb", func(options *opt.Options) {
-		// Set default options
-		options.ReadOnly = true
-		options.BlockCacher = opt.NoCacher
-		options.BlockCacheCapacity = 0
-		options.OpenFilesCacher = opt.NoCacher
-		options.OpenFilesCacheCapacity = 0
-		options.DisableBlockCache = true
-	})
-	panicOnError(err, "", "failed to open mpt db")
-
 	zkRootHash := common.HexToHash(*zkRoot)
 	mptRootHash := common.HexToHash(*mptRoot)
 
 	numLoaders = runtime.GOMAXPROCS(0) * (*parallelismMultipler)
-	trieLoaders = make(chan struct{}, numLoaders)
+	trieLoaders = make(chan dbs, numLoaders)
 	for i := 0; i < numLoaders; i++ {
-		trieLoaders <- struct{}{}
+		zkDb, err := leveldb.NewCustom(*zkDbPath, "zkDb", func(options *opt.Options) {
+			// Set default options
+			options.ReadOnly = true
+			options.BlockCacher = opt.NoCacher
+			options.BlockCacheCapacity = 0
+			options.OpenFilesCacheCapacity = 128
+			options.DisableBlockCache = true
+		})
+		panicOnError(err, "", "failed to open zk db")
+		mptDb, err := leveldb.NewCustom(*mptDbPath, "mptDb", func(options *opt.Options) {
+			// Set default options
+			options.ReadOnly = true
+			options.BlockCacher = opt.NoCacher
+			options.BlockCacheCapacity = 0
+			options.OpenFilesCacheCapacity = 128
+			options.DisableBlockCache = true
+		})
+		panicOnError(err, "", "failed to open mpt db")
+		trieLoaders <- dbs{
+			zkDb:  zkDb,
+			mptDb: mptDb,
+		}
 	}
 
-	checkTrieEquality(&dbs{
-		zkDb:  zkDb,
-		mptDb: mptDb,
-	}, zkRootHash, mptRootHash, "", checkAccountEquality, true, *paranoid)
+	checkTrieEquality(zkRootHash, mptRootHash, "", checkAccountEquality, true, *paranoid)
 
 	for i := 0; i < numLoaders; i++ {
 		<-trieLoaders
@@ -97,7 +94,7 @@ func panicOnError(err error, label, msg string) {
 func dup(s []byte) []byte {
 	return append([]byte{}, s...)
 }
-func checkTrieEquality(dbs *dbs, zkRoot, mptRoot common.Hash, label string, leafChecker func(string, *dbs, []byte, []byte, bool), top, paranoid bool) {
+func checkTrieEquality(zkRoot, mptRoot common.Hash, label string, leafChecker func(string, []byte, []byte, bool), top, paranoid bool) {
 	done := make(chan struct{})
 	start := time.Now()
 	if !top {
@@ -114,13 +111,8 @@ func checkTrieEquality(dbs *dbs, zkRoot, mptRoot common.Hash, label string, leaf
 	}
 	defer close(done)
 
-	zkTrie, err := trie.NewZkTrie(zkRoot, trie.NewZktrieDatabaseFromTriedb(trie.NewDatabaseWithConfig(dbs.zkDb, &trie.Config{Preimages: true})))
-	panicOnError(err, label, "failed to create zk trie")
-	mptTrie, err := trie.NewSecureNoTracer(mptRoot, trie.NewDatabaseWithConfig(dbs.mptDb, &trie.Config{Preimages: true}))
-	panicOnError(err, label, "failed to create mpt trie")
-
-	mptLeafCh := loadMPT(mptTrie, top)
-	zkLeafCh := loadZkTrie(zkTrie, top, paranoid)
+	mptLeafCh := loadMPT(mptRoot, top)
+	zkLeafCh := loadZkTrie(zkRoot, top, paranoid)
 
 	mptLeafMap := <-mptLeafCh
 	zkLeafMap := <-zkLeafCh
@@ -154,11 +146,11 @@ func checkTrieEquality(dbs *dbs, zkRoot, mptRoot common.Hash, label string, leaf
 			panic(fmt.Sprintf("%s key %s (preimage %s) not found in mpt", label, hex.EncodeToString(mptKey), hex.EncodeToString([]byte(preimageKey))))
 		}
 
-		leafChecker(fmt.Sprintf("%s key: %s", label, hex.EncodeToString([]byte(preimageKey))), dbs, zkValue, mptVal, paranoid)
+		leafChecker(fmt.Sprintf("%s key: %s", label, hex.EncodeToString([]byte(preimageKey))), zkValue, mptVal, paranoid)
 	}
 }
 
-func checkAccountEquality(label string, dbs *dbs, zkAccountBytes, mptAccountBytes []byte, paranoid bool) {
+func checkAccountEquality(label string, zkAccountBytes, mptAccountBytes []byte, paranoid bool) {
 	mptAccount := &types.StateAccount{}
 	panicOnError(rlp.DecodeBytes(mptAccountBytes, mptAccount), label, "failed to decode mpt account")
 	zkAccount, err := types.UnmarshalStateAccount(zkAccountBytes)
@@ -181,17 +173,15 @@ func checkAccountEquality(label string, dbs *dbs, zkAccountBytes, mptAccountByte
 	} else if zkAccount.Root != (common.Hash{}) {
 		zkRoot := common.BytesToHash(zkAccount.Root[:])
 		mptRoot := common.BytesToHash(mptAccount.Root[:])
-		<-trieLoaders
 		go func() {
-			trieLoaders <- struct{}{}
 			defer func() {
 				if p := recover(); p != nil {
 					fmt.Println(p)
 					os.Exit(1)
 				}
 			}()
 
-			checkTrieEquality(dbs, zkRoot, mptRoot, label, checkStorageEquality, false, paranoid)
+			checkTrieEquality(zkRoot, mptRoot, label, checkStorageEquality, false, paranoid)
 			accountsDone.Add(1)
 			fmt.Println("Accounts done:", accountsDone.Load())
 		}()
@@ -201,7 +191,7 @@ func checkAccountEquality(label string, dbs *dbs, zkAccountBytes, mptAccountByte
 	}
 }
 
-func checkStorageEquality(label string, _ *dbs, zkStorageBytes, mptStorageBytes []byte, _ bool) {
+func checkStorageEquality(label string, zkStorageBytes, mptStorageBytes []byte, _ bool) {
 	zkValue := common.BytesToHash(zkStorageBytes)
 	_, content, _, err := rlp.Split(mptStorageBytes)
 	panicOnError(err, label, "failed to decode mpt storage")
@@ -211,22 +201,25 @@ func checkStorageEquality(label string, _ *dbs, zkStorageBytes, mptStorageBytes
 	}
 }
 
-func loadMPT(mptTrie *trie.SecureTrie, top bool) chan map[string][]byte {
+func loadMPT(mptRoot common.Hash, top bool) chan map[string][]byte {
 	startKey := make([]byte, 32)
 	mptLeafMap := make(map[string][]byte, 1000)
 	var mptLeafMutex sync.Mutex
 
 	var mptWg sync.WaitGroup
 	for i := 0; i < 255; i++ {
 		startKey[0] = byte(i)
+
+		loaderDbs := <-trieLoaders
+		mptTrie, err := trie.NewSecureNoTracer(mptRoot, trie.NewDatabaseWithConfig(loaderDbs.mptDb, &trie.Config{Preimages: true}))
+		panicOnError(err, mptRoot.Hex(), "failed to create mpt trie")
 		trieIt := trie.NewIterator(mptTrie.NodeIterator(startKey))
 
 		mptWg.Add(1)
-		<-trieLoaders
 		go func() {
 			defer func() {
 				mptWg.Done()
-				trieLoaders <- struct{}{}
+				trieLoaders <- loaderDbs
 			}()
 			for trieIt.Next() {
 				mptLeafMutex.Lock()
@@ -252,12 +245,20 @@ func loadMPT(mptTrie *trie.SecureTrie, top bool) chan map[string][]byte {
 	return respChan
 }
 
-func loadZkTrie(zkTrie *trie.ZkTrie, top, paranoid bool) chan map[string][]byte {
+func loadZkTrie(zkRoot common.Hash, 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)
 	go func() {
+		loaderDbs := <-trieLoaders
+		defer func() {
+			trieLoaders <- loaderDbs
+		}()
+
+		zkTrie, err := trie.NewZkTrie(zkRoot, trie.NewZktrieDatabaseFromTriedb(trie.NewDatabaseWithConfig(loaderDbs.zkDb, &trie.Config{Preimages: true})))
+		panicOnError(err, zkRoot.Hex(), "failed to create zk trie")
+
 		zkTrie.CountLeaves(func(key, value []byte) {
 			preimageKey := zkTrie.GetKey(key)
 			if len(preimageKey) == 0 {
@@ -272,10 +273,10 @@ func loadZkTrie(zkTrie *trie.ZkTrie, top, paranoid bool) chan map[string][]byte
 				fmt.Println("ZK Accounts Loaded:", len(zkLeafMap))
 			}
 		}, func(f func()) {
-			<-trieLoaders
+			tempDb := <-trieLoaders
 			go func() {
 				defer func() {
-					trieLoaders <- struct{}{}
+					trieLoaders <- tempDb
 				}()
 				f()
 			}()