Use a LevelDB instance as the node cache in StateDownloader

Closes: #1154

Author: gsalgado

tests/trinity/core/p2p-proto/test_state_sync.py

@@ -60,7 +60,8 @@ def test_trie_sync(random, event_loop):
     async def _test_trie_sync():
         src_trie, contents = make_random_trie(random)
         dest_db = FakeAsyncMemoryDB()
-        scheduler = HexaryTrieSync(src_trie.root_hash, dest_db, TraceLogger("test"))
+        nodes_cache = MemoryDB()
+        scheduler = HexaryTrieSync(src_trie.root_hash, dest_db, nodes_cache, TraceLogger("test"))
         requests = scheduler.next_batch()
         while len(requests) > 0:
             results = []
@@ -99,7 +100,8 @@ def make_random_state(n):
 async def test_state_sync():
     raw_db, state_root, contents = make_random_state(1000)
     dest_db = FakeAsyncMemoryDB()
-    scheduler = StateSync(state_root, dest_db, TraceLogger('test'))
+    nodes_cache = MemoryDB()
+    scheduler = StateSync(state_root, dest_db, nodes_cache, TraceLogger('test'))
     requests = scheduler.next_batch(10)
     while requests:
         results = []

trinity/sync/full/hexary_trie.py

@@ -4,19 +4,20 @@
     Callable,
     Dict,
     List,
-    Set,
     Tuple,
 )
 
 from eth_utils import (
     encode_hex,
 )
-from eth.utils.logging import TraceLogger
 
 from eth_typing import (
     Hash32
 )
 
+from eth.db.backends.base import BaseDB
+from eth.utils.logging import TraceLogger
+
 from trie.constants import (
     NODE_TYPE_BLANK,
     NODE_TYPE_BRANCH,
@@ -111,16 +112,24 @@ def _get_children(node: Hash32, depth: int
 
 class HexaryTrieSync:
 
-    def __init__(self, root_hash: Hash32, db: AsyncBaseDB, logger: TraceLogger) -> None:
+    def __init__(self,
+                 root_hash: Hash32,
+                 db: AsyncBaseDB,
+                 nodes_cache: BaseDB,
+                 logger: TraceLogger) -> None:
+        # Nodes that haven't been requested yet.
         self.queue: List[SyncRequest] = []
+        # Nodes that have been requested to a peer, but not yet committed to the DB, either
+        # because we haven't processed a reply containing them or because some of their children
+        # haven't been retrieved/committed yet.
         self.requests: Dict[Hash32, SyncRequest] = {}
         self.db = db
         self.root_hash = root_hash
         self.logger = logger
         # A cache of node hashes we know to exist in our DB, used to avoid querying the DB
         # unnecessarily as that's the main bottleneck when dealing with a large DB like for
         # ethereum's mainnet/ropsten.
-        self._existing_nodes: Set[Hash32] = set()
+        self.nodes_cache = nodes_cache
         self.committed_nodes = 0
         if root_hash in self.db:
             self.logger.info("Root node (%s) already exists in DB, nothing to do", root_hash)
@@ -150,11 +159,11 @@ async def schedule(self, node_key: Hash32, parent: SyncRequest, depth: int,
                        leaf_callback: Callable[[bytes, 'SyncRequest'], Awaitable[None]],
                        is_raw: bool = False) -> None:
         """Schedule a request for the node with the given key."""
-        if node_key in self._existing_nodes:
+        if node_key in self.nodes_cache:
             self.logger.trace("Node %s already exists in db", encode_hex(node_key))
             return
         if await self.db.coro_exists(node_key):
-            self._existing_nodes.add(node_key)
+            self.nodes_cache[node_key] = b''
             self.logger.trace("Node %s already exists in db", encode_hex(node_key))
             return
         self._schedule(node_key, parent, depth, leaf_callback, is_raw)
@@ -224,7 +233,7 @@ async def commit(self, request: SyncRequest) -> None:
         """
         self.committed_nodes += 1
         await self.db.coro_set(request.node_key, request.data)
-        self._existing_nodes.add(request.node_key)
+        self.nodes_cache[request.node_key] = b''
         self.requests.pop(request.node_key)
         for ancestor in request.parents:
             ancestor.dependencies -= 1

trinity/sync/full/state.py

@@ -2,6 +2,8 @@
 import collections
 import itertools
 import logging
+from pathlib import Path
+import tempfile
 import time
 from typing import (
     Any,
@@ -33,6 +35,7 @@
     BLANK_ROOT_HASH,
     EMPTY_SHA3,
 )
+from eth.db.backends.level import LevelDB
 from eth.rlp.accounts import Account
 from eth.utils.logging import TraceLogger
 
@@ -80,7 +83,11 @@ def __init__(self,
         self.chaindb = chaindb
         self.peer_pool = peer_pool
         self.root_hash = root_hash
-        self.scheduler = StateSync(root_hash, account_db, self.logger)
+        # We use a LevelDB instance for the nodes cache because a full state download, if run
+        # uninterrupted will visit more than 180M nodes, making an in-memory cache unfeasible.
+        self._nodes_cache_dir = tempfile.TemporaryDirectory(prefix="pyevm-state-sync-cache")
+        self.scheduler = StateSync(
+            root_hash, account_db, LevelDB(cast(Path, self._nodes_cache_dir.name)), self.logger)
         self._handler = PeerRequestHandler(self.chaindb, self.logger, self.cancel_token)
         self.request_tracker = TrieNodeRequestTracker(self._reply_timeout, self.logger)
         self._peer_missing_nodes: Dict[ETHPeer, Set[Hash32]] = collections.defaultdict(set)
@@ -193,6 +200,7 @@ async def _handle_get_block_headers(self, peer: ETHPeer, request: HeaderRequest)
         peer.sub_proto.send_block_headers(headers)
 
     async def _cleanup(self) -> None:
+        self._nodes_cache_dir.cleanup()
         # We don't need to cancel() anything, but we yield control just so that the coroutines we
         # run in the background notice the cancel token has been triggered and return.
         # Also, don't use self.sleep() here as the cancel token will be triggered and that will