Updated Yamux impl.,added tests for yamux and mplex

Author: kaneki003

libp2p/stream_muxer/mplex/mplex_stream.py

@@ -46,8 +46,9 @@ 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]"
@@ -80,6 +81,8 @@ 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()
 
@@ -113,63 +116,65 @@ async def read(self, n: int | None = None) -> bytes:
         :param n: number of bytes to read
         :return: bytes actually read
         """
-        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.
+        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.
                 try:
-                    data = await self.incoming_data_channel.receive()
+                    data = self.incoming_data_channel.receive_nowait()
                     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)
+                    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)
 
     async def write(self, data: bytes) -> None:
         """
         Write to stream.
 
         :return: number of bytes written
         """
-        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 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)
 
     async def close(self) -> None:
         """

libp2p/stream_muxer/yamux/yamux.py

@@ -77,6 +77,8 @@ 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."""
@@ -98,16 +100,32 @@ async def write(self, data: bytes) -> None:
         # Flow control: Check if we have enough send window
         total_len = len(data)
         sent = 0
-
+        logging.debug(f"Stream {self.stream_id}: Starts writing {total_len} bytes ")
         while sent < total_len:
+            # Wait for available window with timeout
+            timeout = False
             async with self.window_lock:
-                # Wait for available window
-                while self.send_window == 0 and not self.closed:
-                    # Release lock while waiting
+                if self.send_window == 0:
+                    logging.debug(
+                        f"Stream {self.stream_id}: Window is zero, waiting for update"
+                    )
+                    # Release lock and wait with timeout
                     self.window_lock.release()
-                    await trio.sleep(0.01)
+                    # To avoid re-acquiring the lock immediately,
+                    with trio.move_on_after(5.0) as cancel_scope:
+                        while self.send_window == 0 and not self.closed:
+                            await trio.sleep(0.01)
+                        # If we timed out, cancel the scope
+                        timeout = cancel_scope.cancelled_caught
+                    # Re-acquire lock
                     await self.window_lock.acquire()
 
+                # If we timed out waiting for window update, raise an error
+                if timeout:
+                    raise MuxedStreamError(
+                        "Timed out waiting for window update after 5 seconds."
+                    )
+
                 if self.closed:
                     raise MuxedStreamError("Stream is closed")
 
@@ -123,25 +141,53 @@ async def write(self, data: bytes) -> None:
                 await self.conn.secured_conn.write(header + chunk)
                 sent += to_send
 
-        # If window is getting low, consider updating
-        if self.send_window < DEFAULT_WINDOW_SIZE // 2:
-            await self.send_window_update()
-
-    async def send_window_update(self, increment: int | None = None) -> None:
-        """Send a window update to peer."""
+    async def send_window_update(
+        self, increment: int | None, skip_lock: bool = False
+    ) -> None:
+        """
+        Send a window update to peer.
+
+        param:increment: The amount to increment the window size by.
+        If None, uses the difference between DEFAULT_WINDOW_SIZE
+        and current receive window.
+        param:skip_lock (bool): If True, skips acquiring window_lock.
+        This should only be used when calling from a context
+        that already holds the lock.
+        """
+        increment_value = 0
         if increment is None:
-            increment = DEFAULT_WINDOW_SIZE - self.recv_window
-
-        if increment <= 0:
+            increment_value = DEFAULT_WINDOW_SIZE - self.recv_window
+        else:
+            increment_value = increment
+        if increment_value <= 0:
+            # If increment is zero or negative, skip sending update
+            logging.debug(
+                f"Stream {self.stream_id}: Skipping window update"
+                f"(increment={increment})"
+            )
             return
+        logging.debug(
+            f"Stream {self.stream_id}: Sending window update with increment={increment}"
+        )
 
-        async with self.window_lock:
-            self.recv_window += increment
+        async def _do_window_update() -> None:
+            self.recv_window += increment_value
             header = struct.pack(
-                YAMUX_HEADER_FORMAT, 0, TYPE_WINDOW_UPDATE, 0, self.stream_id, increment
+                YAMUX_HEADER_FORMAT,
+                0,
+                TYPE_WINDOW_UPDATE,
+                0,
+                self.stream_id,
+                increment_value,
             )
             await self.conn.secured_conn.write(header)
 
+        if skip_lock:
+            await _do_window_update()
+        else:
+            async with self.window_lock:
+                await _do_window_update()
+
     async def read(self, n: int | None = -1) -> bytes:
         # Handle None value for n by converting it to -1
         if n is None:
@@ -198,11 +244,19 @@ async def read(self, n: int | None = -1) -> bytes:
                 # Return all buffered data
                 data = bytes(buffer)
                 buffer.clear()
-                logging.debug(f"Stream {self.stream_id}: Returning {len(data)} bytes")
                 return data
 
-        # For specific size read (n > 0), return available data immediately
-        return await self.conn.read_stream(self.stream_id, n)
+        data = await self.conn.read_stream(self.stream_id, n)
+        async with self.window_lock:
+            self.recv_window -= len(data)
+            # Automatically send a window update if recv_window is low
+            if self.recv_window <= DEFAULT_WINDOW_SIZE // 2:
+                logging.debug(
+                    f"Stream {self.stream_id}: "
+                    f"Low recv_window ({self.recv_window}), sending update"
+                )
+                await self.send_window_update(None, skip_lock=True)
+        return data
 
     async def close(self) -> None:
         if not self.send_closed:

newsfragments/639.feature.rst

@@ -0,0 +1 @@
+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.

tests/core/stream_muxer/test_mplex_read_write_lock.py

@@ -0,0 +1,124 @@
+import pytest
+import trio
+
+from libp2p.abc import ISecureConn
+from libp2p.crypto.keys import PrivateKey, PublicKey
+from libp2p.peer.id import ID
+from libp2p.stream_muxer.mplex.constants import (
+    HeaderTags,
+)
+from libp2p.stream_muxer.mplex.datastructures import (
+    StreamID,
+)
+from libp2p.stream_muxer.mplex.mplex import (
+    Mplex,
+)
+from libp2p.stream_muxer.mplex.mplex_stream import (
+    MplexStream,
+)
+
+
+class DummySecureConn(ISecureConn):
+    """A minimal implementation of ISecureConn for testing."""
+
+    async def write(self, data: bytes) -> None:
+        pass
+
+    async def read(self, n: int | None = -1) -> bytes:
+        return b""
+
+    async def close(self) -> None:
+        pass
+
+    def get_remote_address(self) -> tuple[str, int] | None:
+        return None
+
+    def get_local_peer(self) -> ID:
+        return ID(b"local")
+
+    def get_local_private_key(self) -> PrivateKey:
+        return PrivateKey()  # Dummy key for testing
+
+    def get_remote_peer(self) -> ID:
+        return ID(b"remote")
+
+    def get_remote_public_key(self) -> PublicKey:
+        return PublicKey()  # Dummy key for testing
+
+
+class DummyMuxedConn(Mplex):
+    """A minimal mock of Mplex for testing read/write locks."""
+
+    def __init__(self) -> None:
+        self.secured_conn = DummySecureConn()
+        self.peer_id = ID(b"dummy")
+        self.streams = {}
+        self.streams_lock = trio.Lock()
+        self.event_shutting_down = trio.Event()
+        self.event_closed = trio.Event()
+        self.event_started = trio.Event()
+        self.stream_backlog_limit = 256
+        self.stream_backlog_semaphore = trio.Semaphore(256)
+        channels = trio.open_memory_channel[MplexStream](0)
+        self.new_stream_send_channel, self.new_stream_receive_channel = channels
+
+    async def send_message(
+        self, flag: HeaderTags, data: bytes, stream_id: StreamID
+    ) -> None:
+        await trio.sleep(0.01)
+
+
+@pytest.mark.trio
+async def test_concurrent_writes_are_serialized():
+    stream_id = StreamID(1, True)
+    send_log = []
+
+    class LoggingMuxedConn(DummyMuxedConn):
+        async def send_message(
+            self, flag: HeaderTags, data: bytes, stream_id: StreamID
+        ) -> None:
+            send_log.append(data)
+            await trio.sleep(0.01)
+
+    memory_send, memory_recv = trio.open_memory_channel(8)
+    stream = MplexStream(
+        name="test",
+        stream_id=stream_id,
+        muxed_conn=LoggingMuxedConn(),
+        incoming_data_channel=memory_recv,
+    )
+
+    async def writer(data):
+        await stream.write(data)
+
+    async with trio.open_nursery() as nursery:
+        for i in range(5):
+            nursery.start_soon(writer, f"msg-{i}".encode())
+    # Order doesn't matter due to concurrent execution
+    assert sorted(send_log) == sorted([f"msg-{i}".encode() for i in range(5)])
+
+
+@pytest.mark.trio
+async def test_concurrent_reads_are_serialized():
+    stream_id = StreamID(2, True)
+    muxed_conn = DummyMuxedConn()
+    memory_send, memory_recv = trio.open_memory_channel(8)
+    results = []
+    stream = MplexStream(
+        name="test",
+        stream_id=stream_id,
+        muxed_conn=muxed_conn,
+        incoming_data_channel=memory_recv,
+    )
+    for i in range(5):
+        await memory_send.send(f"data-{i}".encode())
+    await memory_send.aclose()
+
+    async def reader():
+        data = await stream.read(6)
+        results.append(data)
+
+    async with trio.open_nursery() as nursery:
+        for _ in range(5):
+            nursery.start_soon(reader)
+    assert sorted(results) == [f"data-{i}".encode() for i in range(5)]