← back to Exo
fix UDPDiscovery params, create a new transport every time we broadcast
3dd81a1e05810344798dc0ad65bfa3df5cf56103 · 2024-09-04 14:21:19 +0100 · Alex Cheema
Files touched
M exo/networking/udp_discovery.py
Diff
commit 3dd81a1e05810344798dc0ad65bfa3df5cf56103
Author: Alex Cheema <alexcheema123@gmail.com>
Date: Wed Sep 4 14:21:19 2024 +0100
fix UDPDiscovery params, create a new transport every time we broadcast
---
exo/networking/udp_discovery.py | 43 ++++++++++++++++++++---------------------
1 file changed, 21 insertions(+), 22 deletions(-)
diff --git a/exo/networking/udp_discovery.py b/exo/networking/udp_discovery.py
index aea8310e..6817ce8b 100644
--- a/exo/networking/udp_discovery.py
+++ b/exo/networking/udp_discovery.py
@@ -3,13 +3,12 @@ import json
import socket
import time
import traceback
-from typing import List, Dict, Callable, Tuple, Coroutine, Type
+from typing import List, Dict, Callable, Tuple, Coroutine
from .discovery import Discovery
from .peer_handle import PeerHandle
from exo.topology.device_capabilities import DeviceCapabilities, device_capabilities, UNKNOWN_DEVICE_CAPABILITIES
from exo import DEBUG_DISCOVERY
-
class ListenProtocol(asyncio.DatagramProtocol):
def __init__(self, on_message: Callable[[bytes, Tuple[str, int]], Coroutine]):
super().__init__()
@@ -29,30 +28,30 @@ class UDPDiscovery(Discovery):
node_id: str,
node_port: int,
listen_port: int,
+ broadcast_port: int,
create_peer_handle: Callable[[str, str, DeviceCapabilities], PeerHandle],
- broadcast_port: int = None,
broadcast_interval: int = 1,
- device_capabilities: DeviceCapabilities = UNKNOWN_DEVICE_CAPABILITIES,
discovery_timeout: int = 30,
+ device_capabilities: DeviceCapabilities = UNKNOWN_DEVICE_CAPABILITIES,
):
self.node_id = node_id
self.node_port = node_port
- self.device_capabilities = device_capabilities
self.listen_port = listen_port
- self.broadcast_port = broadcast_port if broadcast_port is not None else listen_port
+ self.broadcast_port = broadcast_port
+ self.create_peer_handle = create_peer_handle
self.broadcast_interval = broadcast_interval
+ self.discovery_timeout = discovery_timeout
+ self.device_capabilities = device_capabilities
self.known_peers: Dict[str, Tuple[PeerHandle, float, float]] = {}
self.broadcast_task = None
self.listen_task = None
self.cleanup_task = None
- self.discovery_timeout = discovery_timeout
- self.create_peer_handle = create_peer_handle
async def start(self):
self.device_capabilities = device_capabilities()
self.broadcast_task = asyncio.create_task(self.task_broadcast_presence())
- self.listen_task = asyncio.create_task(self.task_listen_for_peers())
- self.cleanup_task = asyncio.create_task(self.task_cleanup_peers())
+ # self.listen_task = asyncio.create_task(self.task_listen_for_peers())
+ # self.cleanup_task = asyncio.create_task(self.task_cleanup_peers())
async def stop(self):
if self.broadcast_task: self.broadcast_task.cancel()
@@ -62,7 +61,6 @@ class UDPDiscovery(Discovery):
await asyncio.gather(self.broadcast_task, self.listen_task, self.cleanup_task, return_exceptions=True)
async def discover_peers(self, wait_for_peers: int = 0) -> List[PeerHandle]:
- if DEBUG_DISCOVERY >= 2: print("Starting peer discovery process...")
if wait_for_peers > 0:
while len(self.known_peers) < wait_for_peers:
if DEBUG_DISCOVERY >= 2: print(f"Current peers: {len(self.known_peers)}/{wait_for_peers}. Waiting for more peers...")
@@ -70,25 +68,26 @@ class UDPDiscovery(Discovery):
return [peer_handle for peer_handle, _, _ in self.known_peers.values()]
async def task_broadcast_presence(self):
- transport, _ = await asyncio.get_event_loop().create_datagram_endpoint(lambda: asyncio.DatagramProtocol(), local_addr=("0.0.0.0", 0), family=socket.AF_INET)
- sock = transport.get_extra_info("socket")
- sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
-
- message = json.dumps({
- "type": "discovery",
- "node_id": self.node_id,
- "grpc_port": self.node_port,
- "device_capabilities": self.device_capabilities.to_dict(),
- }).encode("utf-8")
-
while True:
try:
+ message = json.dumps({
+ "type": "discovery",
+ "node_id": self.node_id,
+ "grpc_port": self.node_port,
+ "device_capabilities": self.device_capabilities.to_dict(),
+ }).encode("utf-8")
if DEBUG_DISCOVERY >= 3: print(f"Broadcast presence: {message}")
+
+ transport, _ = await asyncio.get_event_loop().create_datagram_endpoint(lambda: asyncio.DatagramProtocol(), local_addr=("0.0.0.0", 0), family=socket.AF_INET)
+ sock = transport.get_extra_info("socket")
+ sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
transport.sendto(message, ("<broadcast>", self.broadcast_port))
except Exception as e:
print(f"Error in broadcast presence: {e}")
print(traceback.format_exc())
finally:
+ if transport:
+ transport.close()
await asyncio.sleep(self.broadcast_interval)
async def on_listen_message(self, data, addr):
← 15b5043d test for reconnect
·
back to Exo
·
test kill pids dcb3ac76 →