varun-r-mallya/py-libp2p
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/py-libp2p.git synced 2026-02-12 16:10:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed variable imports

Browse Source
This commit is contained in:
parth-soni07
2025-09-21 01:36:06 +05:30
parent 35a4bf2d42
commit 721da9364e
6 changed files with 121 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
libp2p/relay/circuit_v2/discovery.py
View File

@ -31,6 +31,11 @@ from libp2p.tools.async_service import (
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 @@ from .protocol_buffer import (
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 @@ class RelayDiscovery(Service):
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,6 +105,10 @@ class RelayDiscovery(Service):
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__()
@ -108,6 +116,8 @@ class RelayDiscovery(Service):
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 @@ class RelayDiscovery(Service):
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 @@ class RelayDiscovery(Service):
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 @@ class RelayDiscovery(Service):
# 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 @@ class RelayDiscovery(Service):
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 @@ class RelayDiscovery(Service):
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