Merge pull request #1138 from pipermerriam/piper/establish-round-trip-request-response-pattern

Establish pattern for generic round trip request response handling

Author: pipermerriam

p2p/peer.py

@@ -7,7 +7,6 @@
 import operator
 import random
 import struct
-import time
 from abc import (
     ABC,
     abstractmethod
@@ -39,6 +38,7 @@
 
 from eth_utils import (
     decode_hex,
+    to_tuple,
 )
 
 from eth_typing import BlockNumber, Hash32
@@ -155,14 +155,6 @@ class BasePeer(BaseService):
     head_td: int = None
     head_hash: Hash32 = None
 
-    # TODO: Instead of a fixed timeout, we should instead monitor response
-    # times for the peer and adjust our timeout accordingly
-    _response_timeout = 60
-    pending_requests: Dict[
-        Type[protocol.Command],
-        Tuple['BaseRequest', 'asyncio.Future[protocol._DecodedMsgType]'],
-    ]
-
     def __init__(self,
                  remote: Node,
                  privkey: datatypes.PrivateKey,
@@ -189,8 +181,6 @@ def __init__(self,
         self.start_time = datetime.datetime.now()
         self.received_msgs: Dict[protocol.Command, int] = collections.defaultdict(int)
 
-        self.pending_requests = {}
-
         self.egress_mac = egress_mac
         self.ingress_mac = ingress_mac
         # FIXME: Yes, the encryption is insecure, see: https://github.com/ethereum/devp2p/issues/32
@@ -210,6 +200,18 @@ async def process_sub_proto_handshake(
             self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
         raise NotImplementedError("Must be implemented by subclasses")
 
+    @contextlib.contextmanager
+    def collect_sub_proto_messages(self) -> Iterator['MsgBuffer']:
+        """
+        Can be used to gather up all messages that are sent to the peer.
+        """
+        if not self.is_running:
+            raise RuntimeError("Cannot collect messages if peer is not running")
+        msg_buffer = MsgBuffer()
+
+        with msg_buffer.subscribe_peer(self):
+            yield msg_buffer
+
     @property
     def received_msgs_count(self) -> int:
         return sum(self.received_msgs.values())
@@ -400,22 +402,6 @@ def handle_sub_proto_msg(self, cmd: protocol.Command, msg: protocol._DecodedMsgT
         else:
             self.logger.warn("Peer %s has no subscribers, discarding %s msg", self, cmd)
 
-        if cmd_type in self.pending_requests:
-            request, future = self.pending_requests[cmd_type]
-            try:
-                request.validate_response(msg)
-            except ValidationError as err:
-                self.logger.debug(
-                    "Response validation failure for pending %s request from peer %s: %s",
-                    cmd_type.__name__,
-                    self,
-                    err,
-                )
-                pass
-            else:
-                future.set_result(msg)
-                self.pending_requests.pop(cmd_type)
-
     def process_msg(self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
         if cmd.is_base_protocol:
             self.handle_p2p_msg(cmd, msg)
@@ -640,6 +626,25 @@ def subscribe(self, peer_pool: 'PeerPool') -> Iterator[None]:
         finally:
             peer_pool.unsubscribe(self)
 
+    @contextlib.contextmanager
+    def subscribe_peer(self, peer: BasePeer) -> Iterator[None]:
+        peer.add_subscriber(self)
+        try:
+            yield
+        finally:
+            peer.remove_subscriber(self)
+
+
+class MsgBuffer(PeerSubscriber):
+    logger = logging.getLogger('p2p.peer.MsgBuffer')
+    msg_queue_maxsize = 500
+    subscription_msg_types = {protocol.Command}
+
+    @to_tuple
+    def get_messages(self) -> Iterator['PEER_MSG_TYPE']:
+        while not self.msg_queue.empty():
+            yield self.msg_queue.get_nowait()
+
 
 class PeerPool(BaseService, AsyncIterable[BasePeer]):
     """
@@ -695,28 +700,33 @@ def unsubscribe(self, subscriber: PeerSubscriber) -> None:
             peer.remove_subscriber(subscriber)
 
     async def start_peer(self, peer: BasePeer) -> None:
+        asyncio.ensure_future(peer.run())
+        await self.wait(peer.events.started.wait(), timeout=1)
         try:
             # Although connect() may seem like a more appropriate place to perform the DAO fork
             # check, we do it here because we want to perform it for incoming peer connections as
             # well.
-            msgs = await self.ensure_same_side_on_dao_fork(peer)
+            with peer.collect_sub_proto_messages() as buffer:
+                await self.ensure_same_side_on_dao_fork(peer)
         except DAOForkCheckFailure as err:
             self.logger.debug("DAO fork check with %s failed: %s", peer, err)
             await peer.disconnect(DisconnectReason.useless_peer)
             return
-        asyncio.ensure_future(peer.run(finished_callback=self._peer_finished))
-        self._add_peer(peer, msgs)
+        else:
+            msgs = tuple((cmd, msg) for _, cmd, msg in buffer.get_messages())
+            self._add_peer(peer, msgs)
 
     def _add_peer(self,
                   peer: BasePeer,
-                  msgs: List[Tuple[protocol.Command, protocol._DecodedMsgType]]) -> None:
+                  msgs: Tuple[Tuple[protocol.Command, protocol._DecodedMsgType], ...]) -> None:
         """Add the given peer to the pool.
 
         Appart from adding it to our list of connected nodes and adding each of our subscriber's
         to the peer, we also add the given messages to our subscriber's queues.
         """
         self.logger.info('Adding %s to pool', peer)
         self.connected_nodes[peer.remote] = peer
+        peer.add_finished_callback(self._peer_finished)
         for subscriber in self._subscribers:
             subscriber.register_peer(peer)
             peer.add_subscriber(subscriber)
@@ -787,15 +797,13 @@ async def connect_to_nodes(self, nodes: Iterator[Node]) -> None:
                 await self.start_peer(peer)
 
     async def ensure_same_side_on_dao_fork(
-            self, peer: BasePeer) -> List[Tuple[protocol.Command, protocol._DecodedMsgType]]:
+            self, peer: BasePeer) -> None:
         """Ensure we're on the same side of the DAO fork as the given peer.
 
         In order to do that we have to request the DAO fork block and its parent, but while we
         wait for that we may receive other messages from the peer, which are returned so that they
         can be re-added to our subscribers' queues when the peer is finally added to the pool.
         """
-        from trinity.protocol.base_block_headers import BaseBlockHeaders
-        msgs = []
         for start_block, vm_class in self.vm_configuration:
             if not issubclass(vm_class, HomesteadVM):
                 continue
@@ -806,46 +814,33 @@ async def ensure_same_side_on_dao_fork(
                 break
 
             start_block = vm_class.dao_fork_block_number - 1
-            # TODO: This can be either an `ETHPeer` or an `LESPeer`.  Will be
-            # fixed once full awaitable request API is completed.
-            request = peer.request_block_headers(  # type: ignore
-                start_block,
-                max_headers=2,
-                reverse=False,
-            )
-            start = time.time()
+
             try:
-                while True:
-                    elapsed = int(time.time() - start)
-                    remaining_timeout = max(0, CHAIN_SPLIT_CHECK_TIMEOUT - elapsed)
-                    cmd, msg = await self.wait(
-                        peer.read_msg(), timeout=remaining_timeout)
-                    if isinstance(cmd, BaseBlockHeaders):
-                        headers = cmd.extract_headers(msg)
-                        break
-                    else:
-                        msgs.append((cmd, msg))
-                        continue
+                headers = await peer.requests.get_block_headers(  # type: ignore
+                    start_block,
+                    max_headers=2,
+                    reverse=False,
+                    timeout=CHAIN_SPLIT_CHECK_TIMEOUT,
+                )
+
             except (TimeoutError, PeerConnectionLost) as err:
                 raise DAOForkCheckFailure(
-                    "Timed out waiting for DAO fork header from {}: {}".format(peer, err))
+                    "Timed out waiting for DAO fork header from {}: {}".format(peer, err)
+                ) from err
             except MalformedMessage as err:
                 raise DAOForkCheckFailure(
                     "Malformed message while doing DAO fork check with {0}: {1}".format(
                         peer, err,
                     )
                 ) from err
-
-            try:
-                request.validate_headers(headers)
             except ValidationError as err:
                 raise DAOForkCheckFailure(
                     "Invalid header response during DAO fork check: {}".format(err)
-                )
+                ) from err
 
             if len(headers) != 2:
                 raise DAOForkCheckFailure(
-                    "Peer failed to return all requested headers for DAO fork check"
+                    "Peer %s failed to return DAO fork check headers".format(peer)
                 )
             else:
                 parent, header = headers
@@ -855,8 +850,6 @@ async def ensure_same_side_on_dao_fork(
             except EthValidationError as err:
                 raise DAOForkCheckFailure("Peer failed DAO fork check validation: {}".format(err))
 
-        return msgs
-
     def _peer_finished(self, peer: BaseService) -> None:
         """Remove the given peer from our list of connected nodes.
         This is passed as a callback to be called when a peer finishes.
@@ -1008,7 +1001,9 @@ def _test() -> None:
     from eth.chains.ropsten import RopstenChain, ROPSTEN_GENESIS_HEADER, ROPSTEN_VM_CONFIGURATION
     from eth.db.backends.memory import MemoryDB
     from trinity.protocol.eth.peer import ETHPeer
+    from trinity.protocol.eth.requests import HeaderRequest as ETHHeaderRequest
     from trinity.protocol.les.peer import LESPeer
+    from trinity.protocol.les.requests import HeaderRequest as LESHeaderRequest
     from tests.p2p.integration_test_helpers import FakeAsyncHeaderDB, connect_to_peers_loop
     logging.basicConfig(level=TRACE_LEVEL_NUM, format='%(asctime)s %(levelname)s: %(message)s')
 
@@ -1041,13 +1036,15 @@ async def request_stuff() -> None:
             '0x59af08ab31822c992bb3dad92ddb68d820aa4c69e9560f07081fa53f1009b152')
         if peer_class == ETHPeer:
             peer = cast(ETHPeer, peer)
-            peer.sub_proto.send_get_block_headers(block_hash, 1, 0, False)
+            peer.sub_proto.send_get_block_headers(ETHHeaderRequest(block_hash, 1, 0, False))
             peer.sub_proto.send_get_block_bodies([block_hash])
             peer.sub_proto.send_get_receipts([block_hash])
         else:
             peer = cast(LESPeer, peer)
             request_id = 1
-            peer.sub_proto.send_get_block_headers(block_hash, 1, 0, False, request_id)
+            peer.sub_proto.send_get_block_headers(
+                LESHeaderRequest(block_hash, 1, 0, False, request_id)
+            )
             peer.sub_proto.send_get_block_bodies([block_hash], request_id + 1)
             peer.sub_proto.send_get_receipts(block_hash, request_id + 2)
 

p2p/service.py

@@ -17,9 +17,19 @@
 from p2p.cancellable import CancellableMixin
 
 
+class ServiceEvents:
+    def __init__(self) -> None:
+        self.started = asyncio.Event()
+        self.stopped = asyncio.Event()
+        self.cleaned_up = asyncio.Event()
+        self.cancelled = asyncio.Event()
+        self.finished = asyncio.Event()
+
+
 class BaseService(ABC, CancellableMixin):
     logger: TraceLogger = None
     _child_services: List['BaseService']
+    _finished_callbacks: List[Callable[['BaseService'], None]]
     # Number of seconds cancel() will wait for run() to finish.
     _wait_until_finished_timeout = 5
 
@@ -32,8 +42,9 @@ def __init__(self, token: CancelToken=None, loop: asyncio.AbstractEventLoop = No
                 TraceLogger, logging.getLogger(self.__module__ + '.' + self.__class__.__name__))
 
         self._run_lock = asyncio.Lock()
-        self.cleaned_up = asyncio.Event()
+        self.events = ServiceEvents()
         self._child_services = []
+        self._finished_callbacks = []
 
         self.loop = loop
         base_token = CancelToken(type(self).__name__, loop=loop)
@@ -56,8 +67,12 @@ async def run(
         elif self.cancel_token.triggered:
             raise RuntimeError("Cannot restart a service that has already been cancelled")
 
+        if finished_callback:
+            self._finished_callbacks.append(finished_callback)
+
         try:
             async with self._run_lock:
+                self.events.started.set()
                 await self._run()
         except OperationCancelled as e:
             self.logger.info("%s finished: %s", self, e)
@@ -66,19 +81,20 @@ async def run(
         finally:
             # Trigger our cancel token to ensure all pending asyncio tasks and background
             # coroutines started by this service exit cleanly.
+            self.events.cancelled.set()
             self.cancel_token.trigger()
 
             await self.cleanup()
 
-            from p2p.peer import BasePeer  # type: ignore
-            if finished_callback is not None:
-                finished_callback(self)
-            elif isinstance(self, BasePeer):
-                # XXX: Only added to help debug https://github.com/ethereum/py-evm/issues/1023;
-                # should be removed eventually.
-                self.logger.warn("%s finished but had no finished_callback", self)
+            for callback in self._finished_callbacks:
+                callback(self)
+
+            self.events.finished.set()
             self.logger.debug("%s halted cleanly", self)
 
+    def add_finished_callback(self, finished_callback: Callable[['BaseService'], None]) -> None:
+        self._finished_callbacks.append(finished_callback)
+
     def run_child_service(self, child_service: 'BaseService') -> 'asyncio.Future[Any]':
         """
         Run a child service and keep a reference to it to be considered during the cleanup.
@@ -94,13 +110,12 @@ async def cleanup(self) -> None:
         The ``_cleanup()`` coroutine is invoked before the child services may have finished
         their cleanup.
         """
-
         await asyncio.gather(*[
-            child_service.cleaned_up.wait()
+            child_service.events.cleaned_up.wait()
             for child_service in self._child_services],
             self._cleanup()
         )
-        self.cleaned_up.set()
+        self.events.cleaned_up.set()
 
     async def cancel(self) -> None:
         """Trigger the CancelToken and wait for the cleaned_up event to be set."""
@@ -111,10 +126,11 @@ async def cancel(self) -> None:
             raise RuntimeError("Cannot cancel a service that has not been started")
 
         self.logger.debug("Cancelling %s", self)
+        self.events.cancelled.set()
         self.cancel_token.trigger()
         try:
             await asyncio.wait_for(
-                self.cleaned_up.wait(), timeout=self._wait_until_finished_timeout)
+                self.events.cleaned_up.wait(), timeout=self._wait_until_finished_timeout)
         except asyncio.futures.TimeoutError:
             self.logger.info("Timed out waiting for %s to finish its cleanup, exiting anyway", self)
         else:
@@ -131,7 +147,10 @@ async def threadsafe_cancel(self) -> None:
         :param poll_period: how many seconds to wait in between each check for service cleanup
         """
         asyncio.run_coroutine_threadsafe(self.cancel(), loop=self.loop)
-        await asyncio.wait_for(self.cleaned_up.wait(), timeout=self._wait_until_finished_timeout)
+        await asyncio.wait_for(
+            self.events.cleaned_up.wait(),
+            timeout=self._wait_until_finished_timeout,
+        )
 
     async def sleep(self, delay: float) -> None:
         """Coroutine that completes after a given time (in seconds)."""

tests/conftest.py

@@ -33,7 +33,7 @@ def vm_logger(namespaces=LOGGING_NAMESPACES):
 
         handler = logging.StreamHandler(sys.stdout)
 
-        # level = TRACE_LEVEL_NUM
+        # level = 5  # TRACE
         # level = logging.DEBUG
         # level = logging.INFO
         level = logging.ERROR