PR Feedback

Author: pipermerriam

p2p/peer.py

@@ -38,6 +38,7 @@
 
 from eth_utils import (
     decode_hex,
+    to_tuple,
 )
 
 from eth_typing import BlockNumber, Hash32
@@ -199,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['MsgCollector']:
+        """
+        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_collector = MsgCollector()
+
+        with msg_collector.subscribe_peer(self):
+            yield msg_collector
+
     @property
     def received_msgs_count(self) -> int:
         return sum(self.received_msgs.values())
@@ -622,6 +635,17 @@ def subscribe_peer(self, peer: BasePeer) -> Iterator[None]:
             peer.remove_subscriber(self)
 
 
+class MsgCollector(PeerSubscriber):
+    logger = logging.getLogger('p2p.peer.MsgCollector')
+    msg_queue_maxsize = 200
+    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]):
     """
     PeerPool maintains connections to up-to max_peers on a given network.
@@ -676,21 +700,25 @@ def unsubscribe(self, subscriber: PeerSubscriber) -> None:
             peer.remove_subscriber(subscriber)
 
     async def start_peer(self, peer: BasePeer) -> None:
+        asyncio.ensure_future(peer.run(finished_callback=self._peer_finished))
+        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 collector:
+                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 collector.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
@@ -768,7 +796,7 @@ 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
@@ -785,25 +813,14 @@ 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.
-            try:
-                class MsgBuffer(PeerSubscriber):
-                    logger = logging.getLogger('p2p.peer.MsgBuffer')
-                    msg_queue_maxsize = 200
-                    subscription_msg_types = {protocol.Command}
-
-                msg_buffer = MsgBuffer()
-
-                with msg_buffer.subscribe_peer(peer):
-                    headers = await peer.handler.get_block_headers(  # type: ignore
-                        start_block,
-                        max_headers=2,
-                        reverse=False,
-                        timeout=CHAIN_SPLIT_CHECK_TIMEOUT,
-                    )
 
-                msgs = [msg_buffer.msg_queue.get_nowait()[1:] for _ in range(msg_buffer.queue_size)]
+            try:
+                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(
@@ -820,15 +837,18 @@ class MsgBuffer(PeerSubscriber):
                     "Invalid header response during DAO fork check: {}".format(err)
                 ) from err
 
-            parent, header = headers
+            if len(headers) != 2:
+                raise DAOForkCheckFailure(
+                    "Peer %s failed to return DAO fork check headers".format(peer)
+                )
+            else:
+                parent, header = headers
 
             try:
                 vm_class.validate_header(header, parent, check_seal=True)
             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.

p2p/service.py

@@ -17,6 +17,15 @@
 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']
@@ -32,7 +41,7 @@ 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.loop = loop
@@ -58,6 +67,7 @@ async def run(
 
         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,6 +76,7 @@ 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()
@@ -77,6 +88,7 @@ async def run(
                 # 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)
+            self.events.finished.set()
             self.logger.debug("%s halted cleanly", self)
 
     def run_child_service(self, child_service: 'BaseService') -> 'asyncio.Future[Any]':
@@ -95,11 +107,11 @@ async def cleanup(self) -> None:
         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."""
@@ -110,10 +122,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:
@@ -130,7 +143,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/trinity/core/p2p-proto/test_peer_block_header_request_and_response_api.py

@@ -74,7 +74,7 @@ async def send_headers():
         remote.sub_proto.send_block_headers(headers)
 
     asyncio.ensure_future(send_headers())
-    response = await peer.handler.get_block_headers(*params)
+    response = await peer.requests.get_block_headers(*params)
 
     assert len(response) == len(headers)
     for expected, actual in zip(headers, response):
@@ -96,12 +96,12 @@ async def test_les_peer_get_headers_round_trip(les_peer_and_remote,
     peer, remote = les_peer_and_remote
 
     async def send_headers():
-        request_id = peer.handler.get_block_headers.pending_request[0].request_id
+        request_id = peer.requests.get_block_headers.pending_request[0].request_id
         remote.sub_proto.send_block_headers(headers, 0, request_id)
         await asyncio.sleep(0)
 
     asyncio.ensure_future(send_headers())
-    response = await peer.handler.get_block_headers(*params)
+    response = await peer.requests.get_block_headers(*params)
 
     assert len(response) == len(headers)
     for expected, actual in zip(headers, response):
@@ -121,7 +121,7 @@ async def send_responses():
         await asyncio.sleep(0)
 
     asyncio.ensure_future(send_responses())
-    response = await peer.handler.get_block_headers(0, 10, 0, False)
+    response = await peer.requests.get_block_headers(0, 10, 0, False)
 
     assert len(response) == len(headers)
     for expected, actual in zip(headers, response):
@@ -145,7 +145,7 @@ async def send_responses():
         await asyncio.sleep(0)
 
     asyncio.ensure_future(send_responses())
-    response = await peer.handler.get_block_headers(0, 5, 0, False)
+    response = await peer.requests.get_block_headers(0, 5, 0, False)
 
     assert len(response) == len(headers)
     for expected, actual in zip(headers, response):

trinity/nodes/base.py

@@ -151,7 +151,11 @@ async def _cleanup(self) -> None:
             # Stop the this IPCServer-specific event loop, so that the IPCServer thread will exit
             self._ipc_loop.stop()
 
-        await asyncio.gather(*[service.cleaned_up.wait() for service in self._auxiliary_services])
+        await asyncio.gather(*[
+            service.events.cleaned_up.wait()
+            for service
+            in self._auxiliary_services
+        ])
 
     def _make_new_loop_thread(self) -> asyncio.AbstractEventLoop:
         new_loop = asyncio.new_event_loop()

trinity/protocol/common/handlers.py

@@ -35,8 +35,7 @@ async def _run(self) -> None:
             manager = getattr(self, attr)
             self.run_child_service(manager)
 
-        while not self.cancel_token.triggered:
-            await self.sleep(0.01)
+        await self.cancel_token.wait()
 
     async def _cleanup(self) -> None:
         pass

trinity/protocol/common/managers.py

@@ -31,7 +31,7 @@
 ReturnType = TypeVar('ReturnType')
 
 
-class BaseRequestManager(PeerSubscriber, BaseService, Generic[PeerClass, RequestType, ResponseType, ReturnType]):
+class BaseRequestManager(PeerSubscriber, BaseService, Generic[PeerClass, RequestType, ResponseType, ReturnType]):  # noqa: E501
     #
     # PeerSubscriber
     #
@@ -55,10 +55,10 @@ def __init__(self, peer: PeerClass, token: CancelToken) -> None:
     # Service API
     #
     async def _run(self) -> None:
-        self.logger.debug("Running %s for peer %s", self.__class__.__name__, self._peer)
+        self.logger.debug("Launching %s for peer %s", self.__class__.__name__, self._peer)
 
         with self.subscribe_peer(self._peer):
-            while True:
+            while self.is_running:
                 peer, cmd, msg = await self.wait(
                     self.msg_queue.get(), token=self.cancel_token)
                 if peer != self._peer:
@@ -74,6 +74,9 @@ async def _cleanup(self) -> None:
 
     def _handle_msg(self, msg: ResponseType) -> None:
         if self.pending_request is None:
+            self.logger.debug(
+                "Got unexpected %s message from %", self.response_msg_name, self._peer
+            )
             return
 
         request, future = self.pending_request
@@ -83,7 +86,7 @@ def _handle_msg(self, msg: ResponseType) -> None:
         except ValidationError as err:
             self.logger.debug(
                 "Response validation failure for pending %s request from peer %s: %s",
-                self._response_msg_type.__name__,
+                self.response_msg_name,
                 self._peer,
                 err,
             )
@@ -109,6 +112,10 @@ def __call__(self) -> ReturnType:  # type: ignore
     def _response_msg_type(self) -> Type[Command]:
         pass
 
+    @property
+    def response_msg_name(self) -> str:
+        return self._response_msg_type.__name__
+
     @abstractmethod
     def _send_sub_proto_request(self, request: RequestType) -> None:
         pass
@@ -119,12 +126,12 @@ async def _wait_for_response(self,
         if self.pending_request is not None:
             self.logger.error(
                 "Already waiting for response to %s for peer: %s",
-                self._response_msg_type.__name__,
+                self.response_msg_name,
                 self._peer,
             )
             raise AlreadyWaiting(
                 "Already waiting for response to {0} for peer: {1}".format(
-                    self._response_msg_type.__name__,
+                    self.response_msg_name,
                     self._peer
                 )
             )

trinity/protocol/eth/managers.py

@@ -13,7 +13,7 @@
 )
 
 from trinity.protocol.common.managers import (
-    BaseRequestManager,
+    BaseRequestManager as _BaseRequestManager,
 )
 
 from .commands import BlockHeaders
@@ -23,7 +23,15 @@
     from .peer import ETHPeer  # noqa: F401
 
 
-class GetBlockHeadersRequestManager(BaseRequestManager['ETHPeer', HeaderRequest, Tuple[BlockHeader, ...], Tuple[BlockHeader, ...]]):
+BaseRequestManager = _BaseRequestManager[
+    'ETHPeer',
+    HeaderRequest,
+    Tuple[BlockHeader, ...],
+    Tuple[BlockHeader, ...],
+]
+
+
+class GetBlockHeadersRequestManager(BaseRequestManager):
     msg_queue_maxsize = 100
 
     _response_msg_type: Type[Command] = BlockHeaders