updated doc-string and reverted mplex-changes

Author: kaneki003

libp2p/stream_muxer/mplex/mplex_stream.py

@@ -46,9 +46,8 @@ class MplexStream(IMuxedStream):
     read_deadline: int | None
     write_deadline: int | None
 
+    # TODO: Add lock for read/write to avoid interleaving receiving messages?
     close_lock: trio.Lock
-    read_lock: trio.Lock
-    write_lock: trio.Lock
 
     # NOTE: `dataIn` is size of 8 in Go implementation.
     incoming_data_channel: "trio.MemoryReceiveChannel[bytes]"
@@ -81,8 +80,6 @@ def __init__(
         self.event_remote_closed = trio.Event()
         self.event_reset = trio.Event()
         self.close_lock = trio.Lock()
-        self.read_lock = trio.Lock()
-        self.write_lock = trio.Lock()
         self.incoming_data_channel = incoming_data_channel
         self._buf = bytearray()
 
@@ -116,65 +113,63 @@ async def read(self, n: int | None = None) -> bytes:
         :param n: number of bytes to read
         :return: bytes actually read
         """
-        async with self.read_lock:
-            if n is not None and n < 0:
-                raise ValueError(
-                    "the number of bytes to read `n` must be non-negative or "
-                    f"`None` to indicate read until EOF, got n={n}"
-                )
-            if self.event_reset.is_set():
-                raise MplexStreamReset
-            if n is None:
-                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, then return.
+        if n is not None and n < 0:
+            raise ValueError(
+                "the number of bytes to read `n` must be non-negative or "
+                f"`None` to indicate read until EOF, got n={n}"
+            )
+        if self.event_reset.is_set():
+            raise MplexStreamReset
+        if n is None:
+            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, then return.
+            try:
+                data = self.incoming_data_channel.receive_nowait()
+                self._buf.extend(data)
+            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 = self.incoming_data_channel.receive_nowait()
+                    data = await self.incoming_data_channel.receive()
                     self._buf.extend(data)
                 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()
-                        self._buf.extend(data)
-                    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(self._read_return_when_blocked())
-            payload = self._buf[:n]
-            self._buf = self._buf[len(payload) :]
-            return bytes(payload)
+                    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(self._read_return_when_blocked())
+        payload = self._buf[:n]
+        self._buf = self._buf[len(payload) :]
+        return bytes(payload)
 
     async def write(self, data: bytes) -> None:
         """
         Write to stream.
 
         :return: number of bytes written
         """
-        async with self.write_lock:
-            if self.event_local_closed.is_set():
-                raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}")
-            flag = (
-                HeaderTags.MessageInitiator
-                if self.is_initiator
-                else HeaderTags.MessageReceiver
-            )
-            await self.muxed_conn.send_message(flag, data, self.stream_id)
+        if self.event_local_closed.is_set():
+            raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}")
+        flag = (
+            HeaderTags.MessageInitiator
+            if self.is_initiator
+            else HeaderTags.MessageReceiver
+        )
+        await self.muxed_conn.send_message(flag, data, self.stream_id)
 
     async def close(self) -> None:
         """

libp2p/stream_muxer/yamux/yamux.py

@@ -77,8 +77,6 @@ def __init__(self, stream_id: int, conn: "Yamux", is_initiator: bool) -> None:
         self.send_window = DEFAULT_WINDOW_SIZE
         self.recv_window = DEFAULT_WINDOW_SIZE
         self.window_lock = trio.Lock()
-        self.read_lock = trio.Lock()
-        self.write_lock = trio.Lock()
 
     async def __aenter__(self) -> "YamuxStream":
         """Enter the async context manager."""

newsfragments/639.feature.rst

@@ -1 +1,6 @@
-Added separate read and write locks to the `MplexStream` & `YamuxStream` class.This ensures thread-safe access and data integrity when multiple coroutines interact with the same MplexStream instance.
+Fixed several flow-control and concurrency issues in the `YamuxStream` class. Previously, stress-testing revealed that transferring data over `DEFAULT_WINDOW_SIZE` would break the stream due to inconsistent window update handling and lock management. The fixes include:
+
+- Removed sending of window updates during writes to maintain correct flow-control.
+- Added proper timeout handling when releasing and acquiring locks to prevent concurrency errors.
+- Corrected the `read` function to properly handle window updates for both `read_until_EOF` and `read_n_bytes`.
+- Added event logging at `send_window_updates` and `waiting_for_window_updates` for better observability.

tests/core/stream_muxer/test_mplex_read_write_lock.py

@@ -224,16 +224,14 @@ async def test_yamux_stream_reset(yamux_pair):
     await client_stream.reset()
     # After reset, reading should raise MuxedStreamReset or MuxedStreamEOF
     try:
-        while True:
-            await server_stream.read()
+        await server_stream.read()
     except (MuxedStreamEOF, MuxedStreamError):
         pass
     else:
         pytest.fail("Expected MuxedStreamEOF or MuxedStreamError")
     # Verify subsequent operations fail with StreamReset or EOF
     with pytest.raises(MuxedStreamError):
-        while True:
-            await server_stream.read()
+        await server_stream.read()
     with pytest.raises(MuxedStreamError):
         await server_stream.write(b"test")
     logging.debug("test_yamux_stream_reset complete")