Fix Mplex and Swarm

libp2p/stream_muxer/mplex/mplex_stream.py

@@ -3,7 +3,6 @@ from typing import TYPE_CHECKING
 import trio
 
 from libp2p.stream_muxer.abc import IMuxedStream
-from libp2p.utils import IQueue, TrioQueue
 
 from .constants import HeaderTags
 from .datastructures import StreamID
@@ -24,10 +23,11 @@ class MplexStream(IMuxedStream):
     read_deadline: int
     write_deadline: int
 
+    # TODO: Add lock for read/write to avoid interleaving receiving messages?
     close_lock: trio.Lock
 
     # NOTE: `dataIn` is size of 8 in Go implementation.
-    incoming_data: IQueue[bytes]
+    incoming_data_channel: "trio.MemoryReceiveChannel[bytes]"
 
     event_local_closed: trio.Event
     event_remote_closed: trio.Event
@@ -35,7 +35,13 @@ class MplexStream(IMuxedStream):
 
     _buf: bytearray
 
-    def __init__(self, name: str, stream_id: StreamID, muxed_conn: "Mplex") -> None:
+    def __init__(
+        self,
+        name: str,
+        stream_id: StreamID,
+        muxed_conn: "Mplex",
+        incoming_data_channel: "trio.MemoryReceiveChannel[bytes]",
+    ) -> None:
         """
         create new MuxedStream in muxer.
 
@@ -51,13 +57,30 @@ class MplexStream(IMuxedStream):
         self.event_remote_closed = trio.Event()
         self.event_reset = trio.Event()
         self.close_lock = trio.Lock()
-        self.incoming_data = TrioQueue()
+        self.incoming_data_channel = incoming_data_channel
         self._buf = bytearray()
 
     @property
     def is_initiator(self) -> bool:
         return self.stream_id.is_initiator
 
+    async def _read_until_eof(self) -> bytes:
+        async for data in self.incoming_data_channel:
+            self._buf.extend(data)
+        payload = self._buf
+        self._buf = self._buf[len(payload) :]
+        return bytes(payload)
+
+    def _read_return_when_blocked(self) -> bytes:
+        buf = bytearray()
+        while True:
+            try:
+                data = self.incoming_data_channel.receive_nowait()
+                buf.extend(data)
+            except (trio.WouldBlock, trio.EndOfChannel):
+                break
+        return buf
+
     async def read(self, n: int = -1) -> bytes:
         """
         Read up to n bytes. Read possibly returns fewer than `n` bytes, if
@@ -73,7 +96,40 @@ class MplexStream(IMuxedStream):
             )
         if self.event_reset.is_set():
             raise MplexStreamReset
-        return await self.incoming_data.get()
+        if n == -1:
+            return await self._read_until_eof()
+        if len(self._buf) == 0:
+            data: bytes
+            # Peek whether there is data available. If yes, we just read until there is no data,
+            # and then return.
+            try:
+                data = self.incoming_data_channel.receive_nowait()
+            except trio.EndOfChannel:
+                raise MplexStreamEOF
+            except trio.WouldBlock:
+                # We know `receive` will be blocked here. Wait for data here with `receive` and
+                # catch all kinds of errors here.
+                try:
+                    data = await self.incoming_data_channel.receive()
+                except trio.EndOfChannel:
+                    if self.event_reset.is_set():
+                        raise MplexStreamReset
+                    if self.event_remote_closed.is_set():
+                        raise MplexStreamEOF
+                except trio.ClosedResourceError as error:
+                    # Probably `incoming_data_channel` is closed in `reset` when we are waiting
+                    # for `receive`.
+                    if self.event_reset.is_set():
+                        raise MplexStreamReset
+                    raise Exception(
+                        "`incoming_data_channel` is closed but stream is not reset. "
+                        "This should never happen."
+                    ) from error
+            self._buf.extend(data)
+        self._buf.extend(self._read_return_when_blocked())
+        payload = self._buf[:n]
+        self._buf = self._buf[len(payload) :]
+        return bytes(payload)
 
     async def write(self, data: bytes) -> int:
         """
@@ -99,22 +155,26 @@ class MplexStream(IMuxedStream):
             if self.event_local_closed.is_set():
                 return
 
+        print(f"!@# stream.close: {self.muxed_conn._id}: step=0")
         flag = (
             HeaderTags.CloseInitiator if self.is_initiator else HeaderTags.CloseReceiver
         )
         # TODO: Raise when `muxed_conn.send_message` fails and `Mplex` isn't shutdown.
         await self.muxed_conn.send_message(flag, None, self.stream_id)
 
+        print(f"!@# stream.close: {self.muxed_conn._id}: step=1")
         _is_remote_closed: bool
         async with self.close_lock:
             self.event_local_closed.set()
             _is_remote_closed = self.event_remote_closed.is_set()
 
+        print(f"!@# stream.close: {self.muxed_conn._id}: step=2")
         if _is_remote_closed:
             # Both sides are closed, we can safely remove the buffer from the dict.
             async with self.muxed_conn.streams_lock:
                 if self.stream_id in self.muxed_conn.streams:
                     del self.muxed_conn.streams[self.stream_id]
+        print(f"!@# stream.close: {self.muxed_conn._id}: step=3")
 
     async def reset(self) -> None:
         """closes both ends of the stream tells this remote side to hang up."""
@@ -132,14 +192,15 @@ class MplexStream(IMuxedStream):
                     if self.is_initiator
                     else HeaderTags.ResetReceiver
                 )
-                async with trio.open_nursery() as nursery:
-                    nursery.start_soon(
-                        self.muxed_conn.send_message, flag, None, self.stream_id
-                    )
+                self.muxed_conn.manager.run_task(
+                    self.muxed_conn.send_message, flag, None, self.stream_id
+                )
 
             self.event_local_closed.set()
             self.event_remote_closed.set()
 
+            await self.incoming_data_channel.aclose()
+
         async with self.muxed_conn.streams_lock:
             if (
                 self.muxed_conn.streams is not None