diff --git a/libp2p/peer/id.py b/libp2p/peer/id.py index c1b52f0..f75563c 100644 --- a/libp2p/peer/id.py +++ b/libp2p/peer/id.py @@ -41,7 +41,7 @@ class ID: def __str__(self) -> str: if FRIENDLY_IDS: - return self.to_string()[2:8] + return f"" else: return self.to_string() diff --git a/libp2p/stream_muxer/abc.py b/libp2p/stream_muxer/abc.py index f16b805..0600dee 100644 --- a/libp2p/stream_muxer/abc.py +++ b/libp2p/stream_muxer/abc.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Optional from libp2p.peer.id import ID from libp2p.security.secure_conn_interface import ISecureConn from libp2p.stream_muxer.mplex.constants import HeaderTags +from libp2p.stream_muxer.mplex.datastructures import StreamID if TYPE_CHECKING: # Prevent GenericProtocolHandlerFn introducing circular dependencies @@ -51,7 +52,7 @@ class IMuxedConn(ABC): """ @abstractmethod - async def read_buffer(self, stream_id: int) -> bytes: + async def read_buffer(self, stream_id: StreamID) -> bytes: """ Read a message from stream_id's buffer, check raw connection for new messages :param stream_id: stream id of stream to read from @@ -59,7 +60,7 @@ class IMuxedConn(ABC): """ @abstractmethod - async def read_buffer_nonblocking(self, stream_id: int) -> Optional[bytes]: + async def read_buffer_nonblocking(self, stream_id: StreamID) -> Optional[bytes]: """ Read a message from `stream_id`'s buffer, non-blockingly. """ @@ -78,7 +79,9 @@ class IMuxedConn(ABC): """ @abstractmethod - async def send_message(self, flag: HeaderTags, data: bytes, stream_id: int) -> int: + async def send_message( + self, flag: HeaderTags, data: bytes, stream_id: StreamID + ) -> int: """ sends a message over the connection :param header: header to use diff --git a/libp2p/stream_muxer/mplex/mplex.py b/libp2p/stream_muxer/mplex/mplex.py index 2dd2618..40fa67b 100644 --- a/libp2p/stream_muxer/mplex/mplex.py +++ b/libp2p/stream_muxer/mplex/mplex.py @@ -13,6 +13,7 @@ from libp2p.utils import ( ) from .constants import HeaderTags +from .datastructures import StreamID from .exceptions import StreamNotFound from .mplex_stream import MplexStream @@ -29,9 +30,9 @@ class Mplex(IMuxedConn): # TODO: `dataIn` in go implementation. Should be size of 8. # TODO: Also, `dataIn` is closed indicating EOF in Go. We don't have similar strategies # to let the `MplexStream`s know that EOF arrived (#235). - buffers: Dict[int, "asyncio.Queue[bytes]"] - stream_queue: "asyncio.Queue[int]" - next_stream_id: int + buffers: Dict[StreamID, "asyncio.Queue[bytes]"] + stream_queue: "asyncio.Queue[StreamID]" + next_channel_id: int # TODO: `generic_protocol_handler` should be refactored out of mplex conn. def __init__( @@ -49,10 +50,7 @@ class Mplex(IMuxedConn): """ self.secured_conn = secured_conn - if self.secured_conn.initiator: - self.next_stream_id = 0 - else: - self.next_stream_id = 1 + self.next_channel_id = 0 # Store generic protocol handler self.generic_protocol_handler = generic_protocol_handler @@ -85,7 +83,7 @@ class Mplex(IMuxedConn): """ raise NotImplementedError() - async def read_buffer(self, stream_id: int) -> bytes: + async def read_buffer(self, stream_id: StreamID) -> bytes: """ Read a message from buffer of the stream specified by `stream_id`, check secured connection for new messages. @@ -97,7 +95,7 @@ class Mplex(IMuxedConn): raise StreamNotFound(f"stream {stream_id} is not found") return await self.buffers[stream_id].get() - async def read_buffer_nonblocking(self, stream_id: int) -> Optional[bytes]: + async def read_buffer_nonblocking(self, stream_id: StreamID) -> Optional[bytes]: """ Read a message from buffer of the stream specified by `stream_id`, non-blockingly. `StreamNotFound` is raised when stream `stream_id` is not found in `Mplex`. @@ -108,13 +106,13 @@ class Mplex(IMuxedConn): return None return await self.buffers[stream_id].get() - def _get_next_stream_id(self) -> int: + def _get_next_channel_id(self) -> int: """ Get next available stream id :return: next available stream id for the connection """ - next_id = self.next_stream_id - self.next_stream_id += 2 + next_id = self.next_channel_id + self.next_channel_id += 1 return next_id async def open_stream(self) -> IMuxedStream: @@ -122,11 +120,12 @@ class Mplex(IMuxedConn): creates a new muxed_stream :return: a new ``MplexStream`` """ - stream_id = self._get_next_stream_id() - name = str(stream_id) - stream = MplexStream(name, stream_id, True, self) + channel_id = self._get_next_channel_id() + stream_id = StreamID(channel_id=channel_id, is_initiator=True) + name = str(channel_id) + stream = MplexStream(name, stream_id, self) self.buffers[stream_id] = asyncio.Queue() - # Default stream name is the `stream_id` + # Default stream name is the `channel_id` await self.send_message(HeaderTags.NewStream, name.encode(), stream_id) return stream @@ -135,10 +134,12 @@ class Mplex(IMuxedConn): accepts a muxed stream opened by the other end """ stream_id = await self.stream_queue.get() - stream = MplexStream(name, stream_id, False, self) + stream = MplexStream(name, stream_id, self) asyncio.ensure_future(self.generic_protocol_handler(stream)) - async def send_message(self, flag: HeaderTags, data: bytes, stream_id: int) -> int: + async def send_message( + self, flag: HeaderTags, data: bytes, stream_id: StreamID + ) -> int: """ sends a message over the connection :param header: header to use @@ -146,7 +147,7 @@ class Mplex(IMuxedConn): :param stream_id: stream the message is in """ # << by 3, then or with flag - header = (stream_id << 3) | flag.value + header = (stream_id.channel_id << 3) | flag.value header = encode_uvarint(header) if data is None: @@ -174,9 +175,10 @@ class Mplex(IMuxedConn): # TODO Deal with other types of messages using flag (currently _) while True: - stream_id, flag, message = await self.read_message() + channel_id, flag, message = await self.read_message() - if stream_id is not None and flag is not None and message is not None: + if channel_id is not None and flag is not None and message is not None: + stream_id = StreamID(channel_id=channel_id, is_initiator=bool(flag & 1)) if stream_id not in self.buffers: self.buffers[stream_id] = asyncio.Queue() await self.stream_queue.put(stream_id) @@ -214,6 +216,6 @@ class Mplex(IMuxedConn): return None, None, None flag = header & 0x07 - stream_id = header >> 3 + channel_id = header >> 3 - return stream_id, flag, message + return channel_id, flag, message diff --git a/libp2p/stream_muxer/mplex/mplex_stream.py b/libp2p/stream_muxer/mplex/mplex_stream.py index a6c2ce2..fe0261b 100644 --- a/libp2p/stream_muxer/mplex/mplex_stream.py +++ b/libp2p/stream_muxer/mplex/mplex_stream.py @@ -3,6 +3,7 @@ import asyncio from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream from .constants import HeaderTags +from .datastructures import StreamID class MplexStream(IMuxedStream): @@ -11,8 +12,7 @@ class MplexStream(IMuxedStream): """ name: str - stream_id: int - initiator: bool + stream_id: StreamID mplex_conn: IMuxedConn read_deadline: int write_deadline: int @@ -22,18 +22,14 @@ class MplexStream(IMuxedStream): _buf: bytearray - def __init__( - self, name: str, stream_id: int, initiator: bool, mplex_conn: IMuxedConn - ) -> None: + def __init__(self, name: str, stream_id: StreamID, mplex_conn: IMuxedConn) -> None: """ create new MuxedStream in muxer - :param stream_id: stream stream id - :param initiator: boolean if this is an initiator + :param stream_id: stream id of this stream :param mplex_conn: muxed connection of this muxed_stream """ self.name = name self.stream_id = stream_id - self.initiator = initiator self.mplex_conn = mplex_conn self.read_deadline = None self.write_deadline = None @@ -42,6 +38,10 @@ class MplexStream(IMuxedStream): self.stream_lock = asyncio.Lock() self._buf = bytearray() + @property + def is_initiator(self) -> bool: + return self.stream_id.is_initiator + async def read(self, n: int = -1) -> bytes: """ Read up to n bytes. Read possibly returns fewer than `n` bytes, @@ -85,7 +85,7 @@ class MplexStream(IMuxedStream): """ flag = ( HeaderTags.MessageInitiator - if self.initiator + if self.is_initiator else HeaderTags.MessageReceiver ) return await self.mplex_conn.send_message(flag, data, self.stream_id) @@ -98,7 +98,9 @@ class MplexStream(IMuxedStream): """ # TODO error handling with timeout # TODO understand better how mutexes are used from go repo - flag = HeaderTags.CloseInitiator if self.initiator else HeaderTags.CloseReceiver + flag = ( + HeaderTags.CloseInitiator if self.is_initiator else HeaderTags.CloseReceiver + ) await self.mplex_conn.send_message(flag, None, self.stream_id) remote_lock = False @@ -131,7 +133,7 @@ class MplexStream(IMuxedStream): if not self.remote_closed: flag = ( HeaderTags.ResetInitiator - if self.initiator + if self.is_initiator else HeaderTags.ResetReceiver ) await self.mplex_conn.send_message(flag, None, self.stream_id) diff --git a/tests/pubsub/conftest.py b/tests/pubsub/conftest.py index 82b7782..12faff3 100644 --- a/tests/pubsub/conftest.py +++ b/tests/pubsub/conftest.py @@ -23,14 +23,16 @@ async def hosts(num_hosts): await asyncio.gather( *[_host.get_network().listen(LISTEN_MADDR) for _host in _hosts] ) - yield _hosts - # Clean up - listeners = [] - for _host in _hosts: - for listener in _host.get_network().listeners.values(): - listener.server.close() - listeners.append(listener) - await asyncio.gather(*[listener.server.wait_closed() for listener in listeners]) + try: + yield _hosts + finally: + # Clean up + listeners = [] + for _host in _hosts: + for listener in _host.get_network().listeners.values(): + listener.server.close() + listeners.append(listener) + await asyncio.gather(*[listener.server.wait_closed() for listener in listeners]) @pytest.fixture