Merge pull request #1184 from carver/request-manager-refactor

Reducing responsibility of BaseRequestManager

Author: carver

p2p/peer.py

@@ -20,6 +20,7 @@
     Dict,
     Iterator,
     List,
+    NamedTuple,
     Set,
     TYPE_CHECKING,
     Tuple,
@@ -209,7 +210,7 @@ async def send_sub_proto_handshake(self) -> None:
 
     @abstractmethod
     async def process_sub_proto_handshake(
-            self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
+            self, cmd: protocol.Command, msg: protocol.PayloadType) -> None:
         raise NotImplementedError("Must be implemented by subclasses")
 
     @contextlib.contextmanager
@@ -365,7 +366,7 @@ async def _run(self) -> None:
                 self.logger.debug("%s disconnected: %s", self, e)
                 return
 
-    async def read_msg(self) -> Tuple[protocol.Command, protocol._DecodedMsgType]:
+    async def read_msg(self) -> Tuple[protocol.Command, protocol.PayloadType]:
         header_data = await self.read(HEADER_LEN + MAC_LEN)
         header = self.decrypt_header(header_data)
         frame_size = self.get_frame_size(header)
@@ -392,7 +393,7 @@ async def read_msg(self) -> Tuple[protocol.Command, protocol._DecodedMsgType]:
             self.received_msgs[cmd] += 1
             return cmd, decoded_msg
 
-    def handle_p2p_msg(self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
+    def handle_p2p_msg(self, cmd: protocol.Command, msg: protocol.PayloadType) -> None:
         """Handle the base protocol (P2P) messages."""
         if isinstance(cmd, Disconnect):
             msg = cast(Dict[str, Any], msg)
@@ -406,12 +407,12 @@ def handle_p2p_msg(self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -
         else:
             raise UnexpectedMessage("Unexpected msg: {} ({})".format(cmd, msg))
 
-    def handle_sub_proto_msg(self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
+    def handle_sub_proto_msg(self, cmd: protocol.Command, msg: protocol.PayloadType) -> None:
         cmd_type = type(cmd)
 
         if self._subscribers:
             was_added = tuple(
-                subscriber.add_msg((self, cmd, msg))
+                subscriber.add_msg(PeerMessage(self, cmd, msg))
                 for subscriber
                 in self._subscribers
             )
@@ -424,14 +425,14 @@ 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)
 
-    def process_msg(self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
+    def process_msg(self, cmd: protocol.Command, msg: protocol.PayloadType) -> None:
         if cmd.is_base_protocol:
             self.handle_p2p_msg(cmd, msg)
         else:
             self.handle_sub_proto_msg(cmd, msg)
 
     async def process_p2p_handshake(
-            self, cmd: protocol.Command, msg: protocol._DecodedMsgType) -> None:
+            self, cmd: protocol.Command, msg: protocol.PayloadType) -> None:
         msg = cast(Dict[str, Any], msg)
         if not isinstance(cmd, Hello):
             await self.disconnect(DisconnectReason.bad_protocol)
@@ -478,7 +479,7 @@ def decrypt_header(self, data: bytes) -> bytes:
         self.ingress_mac.update(sxor(aes, header_ciphertext))
         expected_header_mac = self.ingress_mac.digest()[:HEADER_LEN]
         if not bytes_eq(expected_header_mac, header_mac):
-            raise DecryptionError('Invalid header mac: expected %s, got %s'.format(
+            raise DecryptionError('Invalid header mac: expected {}, got {}'.format(
                 expected_header_mac, header_mac))
         return self.aes_dec.update(header_ciphertext)
 
@@ -563,8 +564,14 @@ def __hash__(self) -> int:
         return hash(self.remote)
 
 
+class PeerMessage(NamedTuple):
+    peer: BasePeer
+    command: protocol.Command
+    payload: protocol.PayloadType
+
+
 class PeerSubscriber(ABC):
-    _msg_queue: 'asyncio.Queue[PEER_MSG_TYPE]' = None
+    _msg_queue: 'asyncio.Queue[PeerMessage]' = None
 
     @property
     @abstractmethod
@@ -609,7 +616,7 @@ def deregister_peer(self, peer: BasePeer) -> None:
         pass
 
     @property
-    def msg_queue(self) -> 'asyncio.Queue[PEER_MSG_TYPE]':
+    def msg_queue(self) -> 'asyncio.Queue[PeerMessage]':
         if self._msg_queue is None:
             self._msg_queue = asyncio.Queue(maxsize=self.msg_queue_maxsize)
         return self._msg_queue
@@ -618,26 +625,29 @@ def msg_queue(self) -> 'asyncio.Queue[PEER_MSG_TYPE]':
     def queue_size(self) -> int:
         return self.msg_queue.qsize()
 
-    def add_msg(self, msg: 'PEER_MSG_TYPE') -> bool:
+    def add_msg(self, msg: PeerMessage) -> bool:
         peer, cmd, _ = msg
 
         if not self.is_subscription_command(type(cmd)):
-            self.logger.trace(  # type: ignore
-                "Discarding %s msg from %s; not subscribed to msg type; "
-                "subscriptions: %s",
-                cmd, peer, self.subscription_msg_types,
-            )
+            if hasattr(self, 'logger'):
+                self.logger.trace(  # type: ignore
+                    "Discarding %s msg from %s; not subscribed to msg type; "
+                    "subscriptions: %s",
+                    cmd, peer, self.subscription_msg_types,
+                )
             return False
 
         try:
-            self.logger.trace(  # type: ignore
-                "Adding %s msg from %s to queue; queue_size=%d", cmd, peer, self.queue_size)
+            if hasattr(self, 'logger'):
+                self.logger.trace(  # type: ignore
+                    "Adding %s msg from %s to queue; queue_size=%d", cmd, peer, self.queue_size)
             self.msg_queue.put_nowait(msg)
             return True
         except asyncio.queues.QueueFull:
-            self.logger.warn(  # type: ignore
-                "%s msg queue is full; discarding %s msg from %s",
-                self.__class__.__name__, cmd, peer)
+            if hasattr(self, 'logger'):
+                self.logger.warn(  # type: ignore
+                    "%s msg queue is full; discarding %s msg from %s",
+                    self.__class__.__name__, cmd, peer)
             return False
 
     @contextlib.contextmanager
@@ -663,7 +673,7 @@ class MsgBuffer(PeerSubscriber):
     subscription_msg_types = {protocol.Command}
 
     @to_tuple
-    def get_messages(self) -> Iterator['PEER_MSG_TYPE']:
+    def get_messages(self) -> Iterator[PeerMessage]:
         while not self.msg_queue.empty():
             yield self.msg_queue.get_nowait()
 
@@ -740,7 +750,7 @@ async def start_peer(self, peer: BasePeer) -> None:
 
     def _add_peer(self,
                   peer: BasePeer,
-                  msgs: Tuple[Tuple[protocol.Command, protocol._DecodedMsgType], ...]) -> None:
+                  msgs: Tuple[Tuple[protocol.Command, protocol.PayloadType], ...]) -> 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
@@ -753,7 +763,7 @@ def _add_peer(self,
             subscriber.register_peer(peer)
             peer.add_subscriber(subscriber)
             for cmd, msg in msgs:
-                subscriber.add_msg((peer, cmd, msg))
+                subscriber.add_msg(PeerMessage(peer, cmd, msg))
 
     async def _run(self) -> None:
         # FIXME: PeerPool should probably no longer be a BaseService, but for now we're keeping it
@@ -1006,9 +1016,6 @@ def __init__(self,
         self.genesis_hash = genesis_hash
 
 
-PEER_MSG_TYPE = Tuple[BasePeer, protocol.Command, protocol._DecodedMsgType]
-
-
 def _test() -> None:
     """
     Create a Peer instance connected to a local geth instance and log messages exchanged with it.
@@ -1066,8 +1073,8 @@ async def request_stuff() -> None:
         hashes = tuple(header.hash for header in headers)
         if peer_class == ETHPeer:
             peer = cast(ETHPeer, peer)
-            peer.sub_proto._send_get_block_bodies(hashes)
-            peer.sub_proto._send_get_receipts(hashes)
+            peer.sub_proto.send_get_block_bodies(hashes)
+            peer.sub_proto.send_get_receipts(hashes)
         else:
             peer = cast(LESPeer, peer)
             request_id = 1

p2p/protocol.py

@@ -1,11 +1,14 @@
+from abc import ABC, abstractmethod
 import logging
 import struct
 from typing import (
     Any,
     Dict,
+    Generic,
     List,
     Tuple,
     Type,
+    TypeVar,
     TYPE_CHECKING,
     Union,
 )
@@ -20,19 +23,23 @@
 )
 from p2p.utils import get_devp2p_cmd_id
 
-
 # Workaround for import cycles caused by type annotations:
 # http://mypy.readthedocs.io/en/latest/common_issues.html#import-cycles
 if TYPE_CHECKING:
     from p2p.peer import ChainInfo, BasePeer  # noqa: F401
 
-
-_DecodedMsgType = Union[
+PayloadType = Union[
     Dict[str, Any],
     List[rlp.Serializable],
     Tuple[rlp.Serializable, ...],
 ]
 
+# A payload to be delivered with a request
+TRequestPayload = TypeVar('TRequestPayload', bound=PayloadType, covariant=True)
+
+# for backwards compatibility for internal references in p2p:
+_DecodedMsgType = PayloadType
+
 
 class Command:
     _cmd_id: int = None
@@ -60,7 +67,7 @@ def is_base_protocol(self) -> bool:
     def __str__(self) -> str:
         return "{} (cmd_id={})".format(self.__class__.__name__, self.cmd_id)
 
-    def encode_payload(self, data: Union[_DecodedMsgType, sedes.CountableList]) -> bytes:
+    def encode_payload(self, data: Union[PayloadType, sedes.CountableList]) -> bytes:
         if isinstance(data, dict):  # convert dict to ordered list
             if not isinstance(self.structure, list):
                 raise ValueError("Command.structure must be a list when data is a dict")
@@ -76,7 +83,7 @@ def encode_payload(self, data: Union[_DecodedMsgType, sedes.CountableList]) -> b
             encoder = sedes.List([type_ for _, type_ in self.structure])
         return rlp.encode(data, sedes=encoder)
 
-    def decode_payload(self, rlp_data: bytes) -> _DecodedMsgType:
+    def decode_payload(self, rlp_data: bytes) -> PayloadType:
         if isinstance(self.structure, sedes.CountableList):
             decoder = self.structure
         else:
@@ -97,13 +104,13 @@ def decode_payload(self, rlp_data: bytes) -> _DecodedMsgType:
             in zip(self.structure, data)
         }
 
-    def decode(self, data: bytes) -> _DecodedMsgType:
+    def decode(self, data: bytes) -> PayloadType:
         packet_type = get_devp2p_cmd_id(data)
         if packet_type != self.cmd_id:
             raise MalformedMessage("Wrong packet type: {}".format(packet_type))
         return self.decode_payload(data[1:])
 
-    def encode(self, data: _DecodedMsgType) -> Tuple[bytes, bytes]:
+    def encode(self, data: PayloadType) -> Tuple[bytes, bytes]:
         payload = self.encode_payload(data)
         enc_cmd_id = rlp.encode(self.cmd_id, sedes=rlp.sedes.big_endian_int)
         frame_size = len(enc_cmd_id) + len(payload)
@@ -123,6 +130,24 @@ def encode(self, data: _DecodedMsgType) -> Tuple[bytes, bytes]:
         return header, body
 
 
+class BaseRequest(ABC, Generic[TRequestPayload]):
+    """
+    Must define command_payload during init. This is the data that will
+    be sent to the peer with the request command.
+    """
+    command_payload: TRequestPayload
+
+    @property
+    @abstractmethod
+    def cmd_type(self) -> Type[Command]:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def response_type(self) -> Type[Command]:
+        raise NotImplementedError
+
+
 class Protocol:
     logger = logging.getLogger("p2p.protocol.Protocol")
     name: str = None
@@ -135,11 +160,20 @@ def __init__(self, peer: 'BasePeer', cmd_id_offset: int) -> None:
         self.peer = peer
         self.cmd_id_offset = cmd_id_offset
         self.commands = [cmd_class(cmd_id_offset) for cmd_class in self._commands]
+        self.cmd_by_type = {cmd_class: cmd_class(cmd_id_offset) for cmd_class in self._commands}
         self.cmd_by_id = dict((cmd.cmd_id, cmd) for cmd in self.commands)
 
     def send(self, header: bytes, body: bytes) -> None:
         self.peer.send(header, body)
 
+    def send_request(self, request: BaseRequest[PayloadType]) -> None:
+        command = self.cmd_by_type[request.cmd_type]
+        header, body = command.encode(request.command_payload)
+        self.send(header, body)
+
+    def supports_command(self, cmd_type: Type[Command]) -> bool:
+        return cmd_type in self.cmd_by_type
+
     def __repr__(self) -> str:
         return "(%s, %d)" % (self.name, self.version)
 

tests/p2p/test_peer_collect_sub_proto_msgs.py

@@ -5,10 +5,6 @@
 
 from trinity.protocol.eth.peer import ETHPeer
 from trinity.protocol.eth.commands import GetBlockHeaders, GetNodeData
-from trinity.protocol.eth.requests import (
-    HeaderRequest,
-    NodeDataRequest,
-)
 
 from tests.trinity.core.peer_helpers import (
     get_directly_linked_peers,
@@ -30,11 +26,11 @@ async def test_peer_subscriber_filters_messages(request, event_loop):
 
     with peer.collect_sub_proto_messages() as collector:
         assert collector in peer._subscribers
-        remote.sub_proto.send_get_node_data(NodeDataRequest([b'\x00' * 32]))
-        remote.sub_proto.send_get_block_headers(HeaderRequest(0, 1, 0, False))
-        remote.sub_proto.send_get_node_data(NodeDataRequest([b'\x00' * 32]))
-        remote.sub_proto.send_get_block_headers(HeaderRequest(1, 1, 0, False))
-        remote.sub_proto.send_get_node_data(NodeDataRequest([b'\x00' * 32]))
+        remote.sub_proto.send_get_node_data([b'\x00' * 32])
+        remote.sub_proto.send_get_block_headers(0, 1, 0, False)
+        remote.sub_proto.send_get_node_data([b'\x00' * 32])
+        remote.sub_proto.send_get_block_headers(1, 1, 0, False)
+        remote.sub_proto.send_get_node_data([b'\x00' * 32])
         await asyncio.sleep(0.01)
 
     assert collector not in peer._subscribers
@@ -51,7 +47,7 @@ async def test_peer_subscriber_filters_messages(request, event_loop):
     assert isinstance(all_messages[4][1], GetNodeData)
 
     # make sure it isn't still collecting
-    remote.sub_proto.send_get_block_headers(HeaderRequest(1, 1, 0, False))
+    remote.sub_proto.send_get_block_headers(1, 1, 0, False)
 
     await asyncio.sleep(0.01)
 

tests/p2p/test_peer_subscriber.py

@@ -8,10 +8,6 @@
 
 from trinity.protocol.eth.peer import ETHPeer
 from trinity.protocol.eth.commands import GetBlockHeaders
-from trinity.protocol.eth.requests import (
-    HeaderRequest,
-    NodeDataRequest,
-)
 
 from tests.trinity.core.peer_helpers import (
     get_directly_linked_peers,
@@ -48,11 +44,11 @@ async def test_peer_subscriber_filters_messages(request, event_loop):
     peer.add_subscriber(header_subscriber)
     peer.add_subscriber(all_subscriber)
 
-    remote.sub_proto.send_get_node_data(NodeDataRequest([b'\x00' * 32]))
-    remote.sub_proto.send_get_block_headers(HeaderRequest(0, 1, 0, False))
-    remote.sub_proto.send_get_node_data(NodeDataRequest([b'\x00' * 32]))
-    remote.sub_proto.send_get_block_headers(HeaderRequest(1, 1, 0, False))
-    remote.sub_proto.send_get_node_data(NodeDataRequest([b'\x00' * 32]))
+    remote.sub_proto.send_get_node_data(tuple([b'\x00' * 32]))
+    remote.sub_proto.send_get_block_headers(0, 1, 0, False)
+    remote.sub_proto.send_get_node_data(tuple([b'\x00' * 32]))
+    remote.sub_proto.send_get_block_headers(1, 1, 0, False)
+    remote.sub_proto.send_get_node_data(tuple([b'\x00' * 32]))
 
     # yeild to let remote and peer transmit.
     await asyncio.sleep(0.01)