Fixed variable imports

Author: parth-soni07

libp2p/relay/circuit_v2/config.py

@@ -46,6 +46,35 @@
 RESERVATION_REFRESH_THRESHOLD = 0.8  # Refresh at 80% of TTL
 MAX_CONCURRENT_RESERVATIONS = 2
 
+# Timeout constants for different components
+DEFAULT_DISCOVERY_STREAM_TIMEOUT = 10  # seconds
+DEFAULT_PEER_PROTOCOL_TIMEOUT = 5  # seconds
+DEFAULT_PROTOCOL_READ_TIMEOUT = 15  # seconds
+DEFAULT_PROTOCOL_WRITE_TIMEOUT = 15  # seconds
+DEFAULT_PROTOCOL_CLOSE_TIMEOUT = 10  # seconds
+DEFAULT_DCUTR_READ_TIMEOUT = 30  # seconds
+DEFAULT_DCUTR_WRITE_TIMEOUT = 30  # seconds
+DEFAULT_DIAL_TIMEOUT = 10  # seconds
+
+
+@dataclass
+class TimeoutConfig:
+    """Timeout configuration for different Circuit Relay v2 components."""
+
+    # Discovery timeouts
+    discovery_stream_timeout: int = DEFAULT_DISCOVERY_STREAM_TIMEOUT
+    peer_protocol_timeout: int = DEFAULT_PEER_PROTOCOL_TIMEOUT
+
+    # Core protocol timeouts
+    protocol_read_timeout: int = DEFAULT_PROTOCOL_READ_TIMEOUT
+    protocol_write_timeout: int = DEFAULT_PROTOCOL_WRITE_TIMEOUT
+    protocol_close_timeout: int = DEFAULT_PROTOCOL_CLOSE_TIMEOUT
+
+    # DCUtR timeouts
+    dcutr_read_timeout: int = DEFAULT_DCUTR_READ_TIMEOUT
+    dcutr_write_timeout: int = DEFAULT_DCUTR_WRITE_TIMEOUT
+    dial_timeout: int = DEFAULT_DIAL_TIMEOUT
+
 
 # Relay roles enum
 class RelayRole(Flag):
@@ -83,6 +112,9 @@ class RelayConfig:
     max_circuit_duration: int = DEFAULT_MAX_CIRCUIT_DURATION
     max_circuit_bytes: int = DEFAULT_MAX_CIRCUIT_BYTES
 
+    # Timeout configuration
+    timeouts: TimeoutConfig = field(default_factory=TimeoutConfig)
+
     # ---------------------------------------------------------------------
     # Backwards-compat boolean helpers.  Existing code that still accesses
     # ``cfg.enable_hop, cfg.enable_stop, cfg.enable_client`` will continue to work.

libp2p/relay/circuit_v2/dcutr.py

@@ -29,6 +29,11 @@
 from libp2p.peer.peerinfo import (
     PeerInfo,
 )
+from libp2p.relay.circuit_v2.config import (
+    DEFAULT_DCUTR_READ_TIMEOUT,
+    DEFAULT_DCUTR_WRITE_TIMEOUT,
+    DEFAULT_DIAL_TIMEOUT,
+)
 from libp2p.relay.circuit_v2.nat import (
     ReachabilityChecker,
 )
@@ -47,11 +52,7 @@
 # Maximum message size for DCUtR (4KiB as per spec)
 MAX_MESSAGE_SIZE = 4 * 1024
 
-# Timeouts
-STREAM_READ_TIMEOUT = 30  # seconds
-STREAM_WRITE_TIMEOUT = 30  # seconds
-DIAL_TIMEOUT = 10  # seconds
-
+# DCUtR protocol constants
 # Maximum number of hole punch attempts per peer
 MAX_HOLE_PUNCH_ATTEMPTS = 5
 
@@ -70,18 +71,33 @@ class DCUtRProtocol(Service):
     hole punching, after they have established an initial connection through a relay.
     """
 
-    def __init__(self, host: IHost):
+    def __init__(
+        self,
+        host: IHost,
+        read_timeout: int = DEFAULT_DCUTR_READ_TIMEOUT,
+        write_timeout: int = DEFAULT_DCUTR_WRITE_TIMEOUT,
+        dial_timeout: int = DEFAULT_DIAL_TIMEOUT,
+    ):
         """
         Initialize the DCUtR protocol.
 
         Parameters
         ----------
         host : IHost
             The libp2p host this protocol is running on
+        read_timeout : int
+            Timeout for stream read operations, in seconds
+        write_timeout : int
+            Timeout for stream write operations, in seconds
+        dial_timeout : int
+            Timeout for dial operations, in seconds
 
         """
         super().__init__()
         self.host = host
+        self.read_timeout = read_timeout
+        self.write_timeout = write_timeout
+        self.dial_timeout = dial_timeout
         self.event_started = trio.Event()
         self._hole_punch_attempts: dict[ID, int] = {}
         self._direct_connections: set[ID] = set()
@@ -161,7 +177,7 @@ async def _handle_dcutr_stream(self, stream: INetStream) -> None:
 
             try:
                 # Read the CONNECT message
-                with trio.fail_after(STREAM_READ_TIMEOUT):
+                with trio.fail_after(self.read_timeout):
                     msg_bytes = await stream.read(MAX_MESSAGE_SIZE)
 
                 # Parse the message
@@ -196,7 +212,7 @@ async def _handle_dcutr_stream(self, stream: INetStream) -> None:
                 response.type = HolePunch.CONNECT
                 response.ObsAddrs.extend(our_addrs)
 
-                with trio.fail_after(STREAM_WRITE_TIMEOUT):
+                with trio.fail_after(self.write_timeout):
                     await stream.write(response.SerializeToString())
 
                 logger.debug(
@@ -206,7 +222,7 @@ async def _handle_dcutr_stream(self, stream: INetStream) -> None:
                 )
 
                 # Wait for SYNC message
-                with trio.fail_after(STREAM_READ_TIMEOUT):
+                with trio.fail_after(self.read_timeout):
                     sync_bytes = await stream.read(MAX_MESSAGE_SIZE)
 
                 # Parse the SYNC message
@@ -300,7 +316,7 @@ async def initiate_hole_punch(self, peer_id: ID) -> bool:
                 connect_msg.ObsAddrs.extend(our_addrs)
 
                 start_time = time.time()
-                with trio.fail_after(STREAM_WRITE_TIMEOUT):
+                with trio.fail_after(self.write_timeout):
                     await stream.write(connect_msg.SerializeToString())
 
                 logger.debug(
@@ -310,7 +326,7 @@ async def initiate_hole_punch(self, peer_id: ID) -> bool:
                 )
 
                 # Receive the peer's CONNECT message
-                with trio.fail_after(STREAM_READ_TIMEOUT):
+                with trio.fail_after(self.read_timeout):
                     resp_bytes = await stream.read(MAX_MESSAGE_SIZE)
 
                 # Calculate RTT
@@ -349,7 +365,7 @@ async def initiate_hole_punch(self, peer_id: ID) -> bool:
                 sync_msg = HolePunch()
                 sync_msg.type = HolePunch.SYNC
 
-                with trio.fail_after(STREAM_WRITE_TIMEOUT):
+                with trio.fail_after(self.write_timeout):
                     await stream.write(sync_msg.SerializeToString())
 
                 logger.debug("Sent SYNC message to %s", peer_id)
@@ -468,7 +484,7 @@ async def _dial_peer(self, peer_id: ID, addr: Multiaddr) -> None:
             peer_info = PeerInfo(peer_id, [addr])
 
             # Try to connect with timeout
-            with trio.fail_after(DIAL_TIMEOUT):
+            with trio.fail_after(self.dial_timeout):
                 await self.host.connect(peer_info)
 
             logger.info("Successfully connected to %s at %s", peer_id, addr)

libp2p/relay/circuit_v2/discovery.py

@@ -31,6 +31,11 @@
     Service,
 )
 
+from .config import (
+    DEFAULT_DISCOVERY_INTERVAL,
+    DEFAULT_DISCOVERY_STREAM_TIMEOUT,
+    DEFAULT_PEER_PROTOCOL_TIMEOUT,
+)
 from .pb.circuit_pb2 import (
     HopMessage,
 )
@@ -43,11 +48,8 @@
 
 logger = logging.getLogger("libp2p.relay.circuit_v2.discovery")
 
-# Constants
+# Discovery constants
 MAX_RELAYS_TO_TRACK = 10
-DEFAULT_DISCOVERY_INTERVAL = 60  # seconds
-STREAM_TIMEOUT = 10  # seconds
-PEER_PROTOCOL_TIMEOUT = 5  # seconds
 
 
 # Extended interfaces for type checking
@@ -87,6 +89,8 @@ def __init__(
         auto_reserve: bool = False,
         discovery_interval: int = DEFAULT_DISCOVERY_INTERVAL,
         max_relays: int = MAX_RELAYS_TO_TRACK,
+        stream_timeout: int = DEFAULT_DISCOVERY_STREAM_TIMEOUT,
+        peer_protocol_timeout: int = DEFAULT_PEER_PROTOCOL_TIMEOUT,
     ) -> None:
         """
         Initialize the discovery service.
@@ -101,13 +105,19 @@ def __init__(
             How often to run discovery, in seconds
         max_relays : int
             Maximum number of relays to track
+        stream_timeout : int
+            Timeout for stream operations during discovery, in seconds
+        peer_protocol_timeout : int
+            Timeout for checking peer protocol support, in seconds
 
         """
         super().__init__()
         self.host = host
         self.auto_reserve = auto_reserve
         self.discovery_interval = discovery_interval
         self.max_relays = max_relays
+        self.stream_timeout = stream_timeout
+        self.peer_protocol_timeout = peer_protocol_timeout
         self._discovered_relays: dict[ID, RelayInfo] = {}
         self._protocol_cache: dict[
             ID, set[str]
@@ -167,19 +177,19 @@ async def discover_relays(self) -> None:
                     continue
 
                 # Don't wait too long for protocol info
-                with trio.move_on_after(PEER_PROTOCOL_TIMEOUT):
+                with trio.move_on_after(self.peer_protocol_timeout):
                     if await self._supports_relay_protocol(peer_id):
                         await self._add_relay(peer_id)
 
             # Limit number of relays we track
-            if len(self._discovered_relays) > MAX_RELAYS_TO_TRACK:
+            if len(self._discovered_relays) > self.max_relays:
                 # Sort by last seen time and keep only the most recent ones
                 sorted_relays = sorted(
                     self._discovered_relays.items(),
                     key=lambda x: x[1].last_seen,
                     reverse=True,
                 )
-                to_remove = sorted_relays[MAX_RELAYS_TO_TRACK:]
+                to_remove = sorted_relays[self.max_relays :]
                 for peer_id, _ in to_remove:
                     del self._discovered_relays[peer_id]
 
@@ -265,7 +275,7 @@ async def _check_via_peerstore(self, peer_id: ID) -> bool | None:
     async def _check_via_direct_connection(self, peer_id: ID) -> bool | None:
         """Check protocol support via direct connection."""
         try:
-            with trio.fail_after(STREAM_TIMEOUT):
+            with trio.fail_after(self.stream_timeout):
                 stream = await self.host.new_stream(peer_id, [PROTOCOL_ID])
                 if stream:
                     await stream.close()
@@ -371,7 +381,7 @@ async def make_reservation(self, peer_id: ID) -> bool:
 
             # Open a stream to the relay with timeout
             try:
-                with trio.fail_after(STREAM_TIMEOUT):
+                with trio.fail_after(self.stream_timeout):
                     stream = await self.host.new_stream(peer_id, [PROTOCOL_ID])
                     if not stream:
                         logger.error("Failed to open stream to relay %s", peer_id)
@@ -387,7 +397,7 @@ async def make_reservation(self, peer_id: ID) -> bool:
                     peer=self.host.get_id().to_bytes(),
                 )
 
-                with trio.fail_after(STREAM_TIMEOUT):
+                with trio.fail_after(self.stream_timeout):
                     await stream.write(request.SerializeToString())
 
                     # Wait for response
@@ -464,7 +474,7 @@ async def _cleanup_expired(self) -> None:
 
         for peer_id, relay_info in self._discovered_relays.items():
             # Check if relay hasn't been seen in a while (3x discovery interval)
-            if now - relay_info.last_seen > DEFAULT_DISCOVERY_INTERVAL * 3:
+            if now - relay_info.last_seen > self.discovery_interval * 3:
                 to_remove.append(peer_id)
                 continue