neuron-tai/packages/p2p/meshnet_p2p/gossip.py

"""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)