Use WriteTrieNode and DeleteTrieNode

blockchain/state/sync_test.go

@@ -388,8 +388,8 @@ func TestCheckStateConsistencyMissNode(t *testing.T) {
 				newState.DeleteCode(hash)
 			} else {
 				data, _ = srcDiskDB.ReadTrieNode(hash)
-				srcDiskDB.GetMemDB().Delete(hash[:])
-				newDiskDB.GetMemDB().Delete(hash[:])
+				srcDiskDB.DeleteTrieNode(hash)
+				newDiskDB.DeleteTrieNode(hash)
 			}
 			// Check consistency : errIterator
 			err = CheckStateConsistency(srcState, newState, srcRoot, 100, nil)
@@ -399,8 +399,13 @@ func TestCheckStateConsistencyMissNode(t *testing.T) {
 			}
 
 			// Recover nodes
-			srcDiskDB.GetMemDB().Put(hash[:], data)
-			newDiskDB.GetMemDB().Put(hash[:], data)
+			if code {
+				srcDiskDB.WriteCode(hash, data)
+				newDiskDB.WriteCode(hash, data)
+			} else {
+				srcDiskDB.WriteTrieNode(hash, data)
+				newDiskDB.WriteTrieNode(hash, data)
+			}
 		}
 	}
 
@@ -605,7 +610,7 @@ func TestIncompleteStateSync(t *testing.T) {
 		}
 		return false
 	}
-	checkTrieConsistency(srcState.TrieDB().DiskDB().(database.DBManager), srcRoot)
+	checkTrieConsistency(srcState.TrieDB().DiskDB(), srcRoot)
 
 	// Create a destination state and sync with the scheduler
 	dstDb := database.NewMemoryDBManager()
@@ -671,7 +676,7 @@ func TestIncompleteStateSync(t *testing.T) {
 			dstState.DeleteCode(node)
 		} else {
 			val, _ = dstDb.ReadTrieNode(node)
-			dstDb.GetMemDB().Delete(node[:])
+			dstDb.DeleteTrieNode(node)
 		}
 
 		if err := checkStateConsistency(dstDb, added[0]); err == nil {
@@ -683,11 +688,11 @@ func TestIncompleteStateSync(t *testing.T) {
 
 		err = CheckStateConsistencyParallel(srcState, dstState, srcRoot, nil)
 		assert.Error(t, err)
+
 		if code {
 			dstDb.WriteCode(node, val)
 		} else {
-			// insert a trie node to memory database
-			dstDb.GetMemDB().Put(node[:], val)
+			dstDb.WriteTrieNode(node, val)
 		}
 	}
 

datasync/downloader/downloader_test.go

@@ -149,7 +149,7 @@ func newTester() *downloadTester {
 		peerMissingStates: make(map[string]map[common.Hash]bool),
 	}
 	tester.stateDb = localdb
-	tester.stateDb.GetMemDB().Put(genesis.Root().Bytes(), []byte{0x00})
+	tester.stateDb.WriteTrieNode(genesis.Root(), []byte{0x00})
 
 	tester.downloader = New(FullSync, tester.stateDb, statedb.NewSyncBloom(1, tester.stateDb.GetMemDB()), new(event.TypeMux), tester, nil, tester.dropPeer, uint64(istanbul.WeightedRandom))
 
@@ -348,7 +348,7 @@ func (dl *downloadTester) CurrentBlock() *types.Block {
 
 	for i := len(dl.ownHashes) - 1; i >= 0; i-- {
 		if block := dl.ownBlocks[dl.ownHashes[i]]; block != nil {
-			if _, err := dl.stateDb.GetMemDB().Get(block.Root().Bytes()); err == nil {
+			if has, _ := dl.stateDb.HasTrieNode(block.Root()); has {
 				return block
 			}
 		}
@@ -430,15 +430,15 @@ func (dl *downloadTester) InsertChain(blocks types.Blocks) (int, error) {
 	for i, block := range blocks {
 		if parent, ok := dl.ownBlocks[block.ParentHash()]; !ok {
 			return i, fmt.Errorf("InsertChain: unknown parent at position %d / %d", i, len(blocks))
-		} else if _, err := dl.stateDb.GetMemDB().Get(parent.Root().Bytes()); err != nil {
-			return i, fmt.Errorf("InsertChain: unknown parent state %x: %v", parent.Root(), err)
+		} else if has, _ := dl.stateDb.HasTrieNode(parent.Root()); !has {
+			return i, fmt.Errorf("InsertChain: unknown parent state %x", parent.Root())
 		}
 		if _, ok := dl.ownHeaders[block.Hash()]; !ok {
 			dl.ownHashes = append(dl.ownHashes, block.Hash())
 			dl.ownHeaders[block.Hash()] = block.Header()
 		}
 		dl.ownBlocks[block.Hash()] = block
-		dl.stateDb.GetMemDB().Put(block.Root().Bytes(), []byte{0x00})
+		dl.stateDb.WriteTrieNode(block.Root(), []byte{0x00})
 		dl.ownChainTd[block.Hash()] = new(big.Int).Add(dl.ownChainTd[block.ParentHash()], block.BlockScore())
 	}
 	return len(blocks), nil
@@ -734,7 +734,7 @@ func (dlp *downloadTesterPeer) RequestNodeData(hashes []common.Hash) error {
 
 	results := make([][]byte, 0, len(hashes))
 	for _, hash := range hashes {
-		if data, err := dlp.dl.peerDb.GetMemDB().Get(hash.Bytes()); err == nil {
+		if data, err := dlp.dl.peerDb.ReadTrieNode(hash); err == nil {
 			if !dlp.dl.peerMissingStates[dlp.id][hash] {
 				results = append(results, data)
 			}

snapshot/generate_test.go

@@ -467,7 +467,7 @@ func TestGenerateCorruptAccountTrie(t *testing.T) {
 	// Delete an account trie leaf and ensure the generator chokes
 	// TODO-Klaytn-Snapshot put propoer block number
 	triedb.Commit(common.HexToHash("0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978"), false, 0)
-	diskdb.GetMemDB().Delete(common.HexToHash("0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7").Bytes())
+	diskdb.DeleteTrieNode(common.HexToHash("0x65145f923027566669a1ae5ccac66f945b55ff6eaeb17d2ea8e048b7d381f2d7"))
 
 	snap := generateSnapshot(diskdb, triedb, 16, common.HexToHash("0xa04693ea110a31037fb5ee814308a6f1d76bdab0b11676bdf4541d2de55ba978"))
 	select {
@@ -531,7 +531,7 @@ func TestGenerateMissingStorageTrie(t *testing.T) {
 	triedb.Commit(common.HexToHash("0xa2282b99de1fc11e32d26bee37707ef49a6978b2d375796a1b026a497193a2ef"), false, 0)
 
 	// Delete a storage trie root and ensure the generator chokes
-	diskdb.GetMemDB().Delete(common.HexToHash("0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67").Bytes())
+	diskdb.DeleteTrieNode(common.HexToHash("0xddefcd9376dd029653ef384bd2f0a126bb755fe84fdcc9e7cf421ba454f2bc67"))
 
 	snap := generateSnapshot(diskdb, triedb, 16, common.HexToHash("0xa2282b99de1fc11e32d26bee37707ef49a6978b2d375796a1b026a497193a2ef"))
 	select {
@@ -594,7 +594,7 @@ func TestGenerateCorruptStorageTrie(t *testing.T) {
 	triedb.Commit(common.HexToHash("0x4a651234bc4b8c7462b5ad4eb95bbb724eb636fed72bb5278d886f9ea4c345f8"), false, 0)
 
 	// Delete a storage trie leaf and ensure the generator chokes
-	diskdb.GetMemDB().Delete(common.HexToHash("0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371").Bytes())
+	diskdb.DeleteTrieNode(common.HexToHash("0x18a0f4d79cff4459642dd7604f303886ad9d77c30cf3d7d7cedb3a693ab6d371"))
 
 	snap := generateSnapshot(diskdb, triedb, 16, common.HexToHash("0x4a651234bc4b8c7462b5ad4eb95bbb724eb636fed72bb5278d886f9ea4c345f8"))
 	select {

storage/database/db_manager.go

@@ -162,7 +162,10 @@ type DBManager interface {
 	HasCodeWithPrefixFromOld(hash common.Hash) bool
 	ReadPreimageFromOld(hash common.Hash) []byte
 
+	// Write StateTrie
+	WriteTrieNode(hash common.Hash, node []byte) error
 	PutTrieNodeToBatch(batch Batch, hash common.Hash, node []byte) error
+	DeleteTrieNode(hash common.Hash) error
 	WritePreimages(number uint64, preimages map[common.Hash][]byte) error
 
 	// from accessors_indexes.go
@@ -1840,6 +1843,18 @@ func (dbm *databaseManager) ReadPreimageFromOld(hash common.Hash) []byte {
 	return data
 }
 
+func (dbm *databaseManager) WriteTrieNode(hash common.Hash, node []byte) error {
+	dbm.lockInMigration.RLock()
+	defer dbm.lockInMigration.RUnlock()
+
+	if dbm.inMigration {
+		if err := dbm.getDatabase(StateTrieMigrationDB).Put(hash[:], node); err != nil {
+			return err
+		}
+	}
+	return dbm.getDatabase(StateTrieDB).Put(hash[:], node)
+}
+
 func (dbm *databaseManager) PutTrieNodeToBatch(batch Batch, hash common.Hash, node []byte) error {
 	return batch.Put(hash[:], node)
 }
@@ -1864,6 +1879,10 @@ func (dbm *databaseManager) WritePreimages(number uint64, preimages map[common.H
 	return nil
 }
 
+func (dbm *databaseManager) DeleteTrieNode(hash common.Hash) error {
+	return dbm.getDatabase(StateTrieDB).Delete(hash[:])
+}
+
 // ReadTxLookupEntry retrieves the positional metadata associated with a transaction
 // hash to allow retrieving the transaction or receipt by hash.
 func (dbm *databaseManager) ReadTxLookupEntry(hash common.Hash) (common.Hash, uint64, uint64) {

storage/statedb/stacktrie.go

@@ -463,9 +463,7 @@ func (st *StackTrie) hash() {
 	h.sha.Write(h.tmp)
 	h.sha.Read(st.val)
 	if st.db != nil {
-		// TODO! Is it safe to Put the slice here?
-		// Do all db implementations copy the value provided?
-		st.db.GetStateTrieDB().Put(st.val, h.tmp)
+		st.db.WriteTrieNode(common.BytesToHash(st.val), h.tmp)
 	}
 }
 
@@ -509,7 +507,7 @@ func (st *StackTrie) Commit() (common.Hash, error) {
 		h.sha.Reset()
 		h.sha.Write(st.val)
 		h.sha.Read(ret)
-		st.db.GetStateTrieDB().Put(ret, st.val)
+		st.db.WriteTrieNode(common.BytesToHash(ret), st.val)
 		return common.BytesToHash(ret), nil
 	}
 	return common.BytesToHash(st.val), nil

storage/statedb/sync_test.go

@@ -508,7 +508,7 @@ func TestSyncOrdering(t *testing.T) {
 				t.Fatalf("failed to process result %v", err)
 			}
 		}
-		batch := diskdb.GetMemDB().NewBatch()
+		batch := diskdb.NewBatch(database.StateTrieDB)
 		if _, err := sched.Commit(batch); err != nil {
 			t.Fatalf("failed to commit data: %v", err)
 		}