Handle in state downloader background task (#1167)

Author: pipermerriam

trinity/sync/full/state.py

@@ -62,7 +62,7 @@
     HexaryTrieSync,
     SyncRequest,
 )
-
+from trinity.utils.cancellation import async_handle_cancellation
 from trinity.utils.timer import Timer
 
 
@@ -234,9 +234,20 @@ async def request_nodes(self, node_keys: Iterable[Hash32]) -> None:
             self.request_tracker.active_requests[peer] = (time.time(), batch)
             asyncio.ensure_future(self._request_and_process_nodes(peer, batch))
 
+    @async_handle_cancellation
     async def _request_and_process_nodes(self, peer: ETHPeer, batch: Tuple[Hash32, ...]) -> None:
         self.logger.debug("Requesting %d trie nodes from %s", len(batch), peer)
-        node_data = await peer.requests.get_node_data(batch)
+        try:
+            node_data = await peer.requests.get_node_data(batch)
+        except TimeoutError as err:
+            self.logger.debug(
+                "Timed out waiting for %s trie nodes from %s: %s",
+                len(batch),
+                peer,
+                err,
+            )
+            node_data = tuple()
+
         try:
             self.request_tracker.active_requests.pop(peer)
         except KeyError:

trinity/utils/cancellation.py

@@ -0,0 +1,24 @@
+import functools
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    TypeVar,
+    Optional,
+)
+
+from cancel_token import OperationCancelled
+
+
+TReturn = TypeVar('TReturn')
+
+
+def async_handle_cancellation(awaitable_fn: Callable[..., Awaitable[TReturn]]
+                              ) -> Callable[..., Awaitable[Optional[TReturn]]]:
+    async def inner(*args: Any, **kwargs: Any) -> Optional[TReturn]:
+        try:
+            return await awaitable_fn(*args, **kwargs)
+        except OperationCancelled:
+            return None
+
+    return functools.update_wrapper(inner, awaitable_fn)