Merge branch 'main' into chore01

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

libp2p/relay/circuit_v2/transport.py

@@ -219,11 +219,25 @@ class CircuitV2Transport(ITransport):
             # Get a relay from the list of discovered relays
             relays = self.discovery.get_relays()
             if relays:
-                # TODO: Implement more sophisticated relay selection
-                # For now, just return the first available relay
-                return relays[0]
+                # Prioritize relays with active reservations
+                relays_with_reservations = []
+                other_relays = []
+
+                for relay_id in relays:
+                    relay_info = self.discovery.get_relay_info(relay_id)
+                    if relay_info and relay_info.has_reservation:
+                        relays_with_reservations.append(relay_id)
+                    else:
+                        other_relays.append(relay_id)
+
+                # Return first available relay with reservation, or fallback to others
+                if relays_with_reservations:
+                    return relays_with_reservations[
+                        attempts % len(relays_with_reservations)
+                    ]
+                elif other_relays:
+                    return other_relays[attempts % len(other_relays)]
 
-            # Wait and try discovery
             await trio.sleep(1)
             attempts += 1
 

libp2p/stream_muxer/mplex/mplex_stream.py

@@ -1,6 +1,5 @@
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-import time
 from types import (
     TracebackType,
 )
@@ -101,12 +100,6 @@ 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
@@ -118,8 +111,8 @@ class MplexStream(IMuxedStream):
     # class of IMuxedConn. Ignoring this type assignment should not pose
     # any risk.
     muxed_conn: "Mplex"  # type: ignore[assignment]
-    read_deadline: float | None
-    write_deadline: float | None
+    read_deadline: int | None
+    write_deadline: int | None
 
     rw_lock: ReadWriteLock
     close_lock: trio.Lock
@@ -163,30 +156,6 @@ 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)
@@ -213,9 +182,6 @@ 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(
@@ -226,13 +192,8 @@ 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:
@@ -246,20 +207,6 @@ 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
@@ -279,43 +226,15 @@ 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
@@ -396,9 +315,8 @@ class MplexStream(IMuxedStream):
 
         :return: True if successful
         """
-        deadline = time.time() + ttl
-        self.read_deadline = deadline
-        self.write_deadline = deadline
+        self.read_deadline = ttl
+        self.write_deadline = ttl
         return True
 
     def set_read_deadline(self, ttl: int) -> bool:
@@ -407,7 +325,7 @@ class MplexStream(IMuxedStream):
 
         :return: True if successful
         """
-        self.read_deadline = time.time() + ttl
+        self.read_deadline = ttl
         return True
 
     def set_write_deadline(self, ttl: int) -> bool:
@@ -416,7 +334,7 @@ class MplexStream(IMuxedStream):
 
         :return: True if successful
         """
-        self.write_deadline = ttl + time.time()
+        self.write_deadline = ttl
         return True
 
     def get_remote_address(self) -> tuple[str, int] | None: