feat(mplex): add timeout handling for read and write operation

Signed-off-by: varun-r-mallya <varunrmallya@gmail.com>

libp2p/stream_muxer/mplex/mplex_stream.py

@@ -1,5 +1,6 @@
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
+import time
 from types import (
     TracebackType,
 )
@@ -100,6 +101,12 @@ class ReadWriteLock:
                 self.release_write()
 
 
+class MplexStreamTimeout(Exception):
+    """Raised when a stream operation exceeds its deadline."""
+
+    pass
+
+
 class MplexStream(IMuxedStream):
     """
     reference: https://github.com/libp2p/go-mplex/blob/master/stream.go
@@ -111,8 +118,8 @@ class MplexStream(IMuxedStream):
     # class of IMuxedConn. Ignoring this type assignment should not pose
     # any risk.
     muxed_conn: "Mplex"  # type: ignore[assignment]
-    read_deadline: int | None
+    read_deadline: float | None
-    write_deadline: int | None
+    write_deadline: float | None
 
     rw_lock: ReadWriteLock
     close_lock: trio.Lock
@@ -156,6 +163,30 @@ class MplexStream(IMuxedStream):
     def is_initiator(self) -> bool:
         return self.stream_id.is_initiator
 
+    def _check_read_deadline(self) -> None:
+        """Check if read deadline has expired and raise timeout if needed."""
+        if self.read_deadline is not None and time.time() > self.read_deadline:
+            raise MplexStreamTimeout("Read operation exceeded deadline")
+
+    def _check_write_deadline(self) -> None:
+        """Check if write deadline has expired and raise timeout if needed."""
+        if self.write_deadline is not None and time.time() > self.write_deadline:
+            raise MplexStreamTimeout("Write operation exceeded deadline")
+
+    def _get_read_timeout(self) -> float | None:
+        """Calculate remaining time until read deadline."""
+        if self.read_deadline is None:
+            return None
+        remaining = self.read_deadline - time.time()
+        return max(0.0, remaining) if remaining > 0 else 0
+
+    def _get_write_timeout(self) -> float | None:
+        """Calculate remaining time until write deadline."""
+        if self.write_deadline is None:
+            return None
+        remaining = self.write_deadline - time.time()
+        return max(0.0, remaining) if remaining > 0 else 0
+
     async def _read_until_eof(self) -> bytes:
         async for data in self.incoming_data_channel:
             self._buf.extend(data)
@@ -182,6 +213,9 @@ class MplexStream(IMuxedStream):
         :param n: number of bytes to read
         :return: bytes actually read
         """
+        # check deadline before starting
+        self._check_read_deadline()
+
         async with self.rw_lock.read_lock():
             if n is not None and n < 0:
                 raise ValueError(
@@ -192,8 +226,13 @@ class MplexStream(IMuxedStream):
                 raise MplexStreamReset
             if n is None:
                 return await self._read_until_eof()
+
+            # check deadline again before potentially blocking operation
+            self._check_read_deadline()
+
             if len(self._buf) == 0:
                 data: bytes
+                timeout = self._get_read_timeout()
                 # Peek whether there is data available. If yes, we just read until
                 # there is no data, then return.
                 try:
@@ -207,6 +246,20 @@ class MplexStream(IMuxedStream):
                     try:
                         data = await self.incoming_data_channel.receive()
                         self._buf.extend(data)
+                        if timeout is not None and timeout <= 0:
+                            raise MplexStreamTimeout(
+                                "Read deadline exceeded while waiting for data"
+                            )
+
+                        if timeout is not None:
+                            with trio.fail_after(timeout):
+                                data = await self.incoming_data_channel.receive()
+                        else:
+                            data = await self.incoming_data_channel.receive()
+
+                        self._buf.extend(data)
+                    except trio.TooSlowError:
+                        raise MplexStreamTimeout("Read operation timed out")
                     except trio.EndOfChannel:
                         if self.event_reset.is_set():
                             raise MplexStreamReset
@@ -226,15 +279,43 @@ class MplexStream(IMuxedStream):
             self._buf = self._buf[len(payload) :]
             return bytes(payload)
 
+    async def _read_until_eof_with_timeout(self) -> bytes:
+        """Read until EOF with timeout support."""
+        timeout = self._get_read_timeout()
+
+        try:
+            if timeout is not None:
+                with trio.fail_after(timeout):
+                    async for data in self.incoming_data_channel:
+                        self._buf.extend(data)
+            else:
+                async for data in self.incoming_data_channel:
+                    self._buf.extend(data)
+        except trio.TooSlowError:
+            raise MplexStreamTimeout("Read until EOF operation timed out")
+
+        payload = self._buf
+        self._buf = self._buf[len(payload) :]
+        return bytes(payload)
+
     async def write(self, data: bytes) -> None:
         """
         Write to stream.
 
         :return: number of bytes written
         """
+        # Check deadline before starting
+        self._check_write_deadline()
+
         async with self.rw_lock.write_lock():
             if self.event_local_closed.is_set():
                 raise MplexStreamClosed(f"cannot write to closed stream: data={data!r}")
+
+            # Check deadline again after acquiring lock
+            timeout = self._get_write_timeout()
+            if timeout is not None and timeout <= 0:
+                raise MplexStreamTimeout("Write deadline exceeded")
+
             flag = (
                 HeaderTags.MessageInitiator
                 if self.is_initiator
@@ -315,8 +396,9 @@ class MplexStream(IMuxedStream):
 
         :return: True if successful
         """
-        self.read_deadline = ttl
+        deadline = time.time() + ttl
-        self.write_deadline = ttl
+        self.read_deadline = deadline
+        self.write_deadline = deadline
         return True
 
     def set_read_deadline(self, ttl: int) -> bool:
@@ -325,7 +407,7 @@ class MplexStream(IMuxedStream):
 
         :return: True if successful
         """
-        self.read_deadline = ttl
+        self.read_deadline = time.time() + ttl
         return True
 
     def set_write_deadline(self, ttl: int) -> bool:
@@ -334,7 +416,7 @@ class MplexStream(IMuxedStream):
 
         :return: True if successful
         """
-        self.write_deadline = ttl
+        self.write_deadline = ttl + time.time()
         return True
 
     def get_remote_address(self) -> tuple[str, int] | None: