Conform stream_id to go-mplex

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"<peer.ID {self.to_string()[2:8]}>"
         else:
             return self.to_string()
 

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

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]"]
+    buffers: Dict[StreamID, "asyncio.Queue[bytes]"]
-    stream_queue: "asyncio.Queue[int]"
+    stream_queue: "asyncio.Queue[StreamID]"
-    next_stream_id: int
+    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_channel_id = 0
-            self.next_stream_id = 0
-        else:
-            self.next_stream_id = 1
 
         # 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
+        next_id = self.next_channel_id
-        self.next_stream_id += 2
+        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()
+        channel_id = self._get_next_channel_id()
-        name = str(stream_id)
+        stream_id = StreamID(channel_id=channel_id, is_initiator=True)
-        stream = MplexStream(name, stream_id, True, self)
+        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

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
+    stream_id: StreamID
-    initiator: bool
     mplex_conn: IMuxedConn
     read_deadline: int
     write_deadline: int
@@ -22,18 +22,14 @@ class MplexStream(IMuxedStream):
 
     _buf: bytearray
 
-    def __init__(
+    def __init__(self, name: str, stream_id: StreamID, mplex_conn: IMuxedConn) -> None:
-        self, name: str, stream_id: int, initiator: bool, mplex_conn: IMuxedConn
-    ) -> None:
         """
         create new MuxedStream in muxer
-        :param stream_id: stream stream id
+        :param stream_id: stream id of this stream
-        :param initiator: boolean if this is an initiator
         :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)

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
+    try:
-    # Clean up
+        yield _hosts
-    listeners = []
+    finally:
-    for _host in _hosts:
+        # Clean up
-        for listener in _host.get_network().listeners.values():
+        listeners = []
-            listener.server.close()
+        for _host in _hosts:
-            listeners.append(listener)
+            for listener in _host.get_network().listeners.values():
-    await asyncio.gather(*[listener.server.wait_closed() for listener in listeners])
+                listener.server.close()
+                listeners.append(listener)
+        await asyncio.gather(*[listener.server.wait_closed() for listener in listeners])
 
 
 @pytest.fixture