"""WebSocket gossip client — connects to relay, publish/subscribe to topics.""" from __future__ import annotations import asyncio import json import logging import threading import time import uuid from collections import defaultdict from typing import Callable log = logging.getLogger(__name__) # Message envelope topics TOPIC_NODE_JOIN = "node-join" TOPIC_NODE_LEAVE = "node-leave" TOPIC_COVERAGE_UPDATE = "coverage-update" TOPIC_HEARTBEAT = "heartbeat" TOPIC_PEER_LIST = "peer-list" TOPIC_RELAY_ANNOUNCE = "relay-announce" _MSG_TTL = 3 # max re-broadcast hops def _make_envelope(topic: str, payload: dict, from_peer: str, ttl: int = _MSG_TTL) -> dict: return { "topic": topic, "version": 1, "from_peer": from_peer, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "msg_id": str(uuid.uuid4()), "ttl": ttl, "payload": payload, } class GossipClient: """Thread-safe WebSocket gossip client. Usage:: client = GossipClient(relay_url="ws://relay:8765", peer_id="abc123") client.subscribe("node-join", lambda env: print(env["payload"])) client.start() client.publish("node-join", {"addr": "http://192.168.1.42:8001", ...}) ... client.stop() """ def __init__(self, relay_url: str, peer_id: str, reconnect_interval: float = 3.0): self.relay_url = relay_url self.peer_id = peer_id self.reconnect_interval = reconnect_interval self._handlers: dict[str, list[Callable]] = defaultdict(list) self._seen: set[str] = set() self._loop: asyncio.AbstractEventLoop | None = None self._thread: threading.Thread | None = None self._ws = None self._running = False self._connected = threading.Event() self._stop_event: asyncio.Event | None = None def subscribe(self, topic: str, handler: Callable) -> None: """Register a sync callback for messages on topic.""" self._handlers[topic].append(handler) def publish(self, topic: str, payload: dict) -> None: """Send a gossip message to all peers via the relay. Thread-safe.""" envelope = _make_envelope(topic, payload, self.peer_id) if self._loop and self._running: asyncio.run_coroutine_threadsafe(self._send(envelope), self._loop) def start(self) -> None: """Start the gossip client in a background thread.""" self._running = True self._loop = asyncio.new_event_loop() self._thread = threading.Thread(target=self._run_loop, daemon=True, name="gossip") self._thread.start() def wait_connected(self, timeout: float = 5.0) -> bool: """Block until connected to relay or timeout. Returns True if connected.""" return self._connected.wait(timeout) def stop(self) -> None: """Stop the gossip client and clean up.""" self._running = False if self._loop and self._stop_event is not None: self._loop.call_soon_threadsafe(self._stop_event.set) if self._thread: self._thread.join(timeout=3.0) # ------------------------------------------------------------------ # Internal asyncio methods (run inside the background event loop) # ------------------------------------------------------------------ def _run_loop(self) -> None: asyncio.set_event_loop(self._loop) self._stop_event = asyncio.Event() try: self._loop.run_until_complete(self._connect_loop()) except Exception: log.debug("Gossip loop exited", exc_info=True) async def _connect_loop(self) -> None: import websockets # type: ignore[import] while self._running and not (self._stop_event and self._stop_event.is_set()): try: async with websockets.connect( self.relay_url, ping_interval=20, ping_timeout=10, open_timeout=5, ) as ws: self._ws = ws self._connected.set() log.debug("Gossip connected to %s", self.relay_url) # Send peer registration await ws.send(json.dumps( _make_envelope( "peer-register", {"peer_id": self.peer_id}, self.peer_id, ) )) await self._receive_loop(ws) except Exception as exc: self._connected.clear() if self._running: log.debug("Gossip disconnected (%s); reconnecting in %ss", exc, self.reconnect_interval) await asyncio.sleep(self.reconnect_interval) async def _receive_loop(self, ws) -> None: async for raw in ws: try: envelope = json.loads(raw) except (json.JSONDecodeError, TypeError): continue msg_id = envelope.get("msg_id", "") if msg_id in self._seen: continue self._seen.add(msg_id) if len(self._seen) > 10_000: # Trim seen set to avoid unbounded growth self._seen = set(list(self._seen)[-5_000:]) topic = envelope.get("topic", "") for handler in self._handlers.get(topic, []): try: handler(envelope) except Exception: log.debug("Gossip handler error for topic %s", topic, exc_info=True) async def _send(self, envelope: dict) -> None: if self._ws is not None: try: await self._ws.send(json.dumps(envelope)) except Exception as exc: log.debug("Gossip send failed: %s", exc)