new service decorator to timeout get_contract_code

Author: carver

p2p/cancel_token.py

@@ -105,7 +105,8 @@ async def wait_first(self,
                          *awaitables: Awaitable[_TReturn],
                          token: CancelToken = None,
                          timeout: float = None) -> _TReturn:
-        """Wait for the first awaitable to complete, unless we timeout or the token chain is triggered.
+        """
+        Wait for the first awaitable to complete, unless we timeout or the token chain is triggered.
 
         The given token is chained with this service's token, so triggering either will cancel
         this.

p2p/constants.py

@@ -40,6 +40,9 @@
 # us.
 CHAIN_SPLIT_CHECK_TIMEOUT = 5 * REPLY_TIMEOUT
 
+# Default timeout before giving up on a caller-initiated interaction
+COMPLETION_TIMEOUT = 5
+
 # Types of LES Announce messages
 LES_ANNOUNCE_SIMPLE = 1
 LES_ANNOUNCE_SIGNED = 2

p2p/lightchain.py

@@ -38,7 +38,13 @@
 )
 from p2p.cancel_token import CancelToken
 from p2p import protocol
-from p2p.constants import REPLY_TIMEOUT
+from p2p.constants import (
+    COMPLETION_TIMEOUT,
+    REPLY_TIMEOUT,
+)
+from p2p.p2p_proto import (
+    DisconnectReason,
+)
 from p2p.peer import (
     LESPeer,
     PeerPool,
@@ -47,6 +53,7 @@
 from p2p.rlp import BlockBody
 from p2p.service import (
     BaseService,
+    service_timeout,
 )
 from p2p.utils import gen_request_id
 
@@ -144,6 +151,7 @@ async def get_account(self, block_hash: Hash32, address: Address) -> Account:
         return rlp.decode(rlp_account, sedes=Account)
 
     @alru_cache(maxsize=1024, cache_exceptions=False)
+    @service_timeout(COMPLETION_TIMEOUT)
     async def get_contract_code(self, block_hash: Hash32, address: Address) -> bytes:
         """
         :param block_hash: find code as of the block with block_hash

p2p/service.py

@@ -1,14 +1,17 @@
 from abc import ABC, abstractmethod
 import asyncio
+import functools
 import logging
-from typing import (  # noqa: #401
+from typing import (
     Any,
     Awaitable,
     Callable,
-    cast,
+    Dict,
     List,
     Optional,
-    TypeVar,
+    Tuple,
+    TypeVar
+    cast,
 )
 
 from evm.utils.logging import TraceLogger
@@ -129,6 +132,25 @@ async def _cleanup(self) -> None:
         raise NotImplementedError()
 
 
+def service_timeout(timeout):
+    """
+    Decorator to time out a method call.
+
+    :param timeout: seconds to wait before raising a timeout exception
+
+    :raise asyncio.futures.TimeoutError: if the call is not complete before timeout seconds
+    """
+    def decorator(func):
+        @functools.wraps(func)
+        async def wrapped(service: BaseService, *args: Tuple, **kwargs: Dict) -> Any:
+            return await service.wait(
+                func(service, *args, **kwargs),
+                timeout=timeout,
+            )
+        return wrapped
+    return decorator
+
+
 class EmptyService(BaseService):
     async def _run(self) -> None:
         pass