"""Server-sent events: one LISTEN connection per process, fanned out to every subscriber queue, plus presence. A thread holds a sync psycopg connection (app role) doing LISTEN ps_events and pushes each NOTIFY payload to the event loop with call_soon_threadsafe. It reconnects with backoff when the database drops. Presence is kept in process (one uvicorn worker); phase 8 completes it. """ from __future__ import annotations import asyncio import contextlib import json import logging import threading import uuid from collections.abc import AsyncIterator, Awaitable, Callable from dataclasses import dataclass, field from typing import Any log = logging.getLogger(__name__) CHANNEL = "ps_events" PING_SECONDS = 25.0 QUEUE_MAX = 1000 BACKOFF_MAX = 30.0 Item = tuple[str, dict[str, Any]] def publish(conn: Any, payload: dict[str, Any]) -> None: """NOTIFY ps_events with a JSON payload. Inside a transaction it is delivered on commit. `payload["event"]` names the SSE event (default `changed`); the rest is the event data. """ conn.execute("SELECT ps_notify(%s, %s)", (CHANNEL, json.dumps(payload, default=str, separators=(",", ":")))) @dataclass(eq=False) class Subscriber: queue: asyncio.Queue[Item] loop: asyncio.AbstractEventLoop session_id: str user_id: str display_name: str screen: str | None = None id: str = field(default_factory=lambda: uuid.uuid4().hex) closed: bool = False def put(self, item: Item) -> None: """Called on the event loop.""" try: self.queue.put_nowait(item) except asyncio.QueueFull: log.warning("event queue full for a subscriber; dropping one event") class EventHub: def __init__(self, connect: Callable[[], Any] | None = None): self._connect = connect self._subs: dict[str, Subscriber] = {} self._lock = threading.Lock() self._stop = threading.Event() self._thread: threading.Thread | None = None self.connected = False # ------------------------------------------------ subscribers def subscribe(self, *, session_id: Any, user_id: Any, display_name: str, screen: str | None = None) -> Subscriber: """Register a subscriber on the running loop and broadcast presence.""" sub = Subscriber( queue=asyncio.Queue(maxsize=QUEUE_MAX), loop=asyncio.get_running_loop(), session_id=str(session_id), user_id=str(user_id), display_name=display_name, screen=screen, ) with self._lock: self._subs[sub.id] = sub self._broadcast_presence() return sub def unsubscribe(self, sub: Subscriber) -> None: sub.closed = True with self._lock: removed = self._subs.pop(sub.id, None) if removed is not None: self._broadcast_presence() def presence(self) -> dict[str, Any]: with self._lock: subs = list(self._subs.values()) seen: dict[tuple[str, str | None], dict[str, Any]] = {} for s in subs: seen.setdefault( (s.user_id, s.screen), {"user_id": s.user_id, "display_name": s.display_name, "screen": s.screen} ) users = sorted(seen.values(), key=lambda u: (u["display_name"].lower(), u["screen"] or "")) return {"users": users} def _send(self, sub: Subscriber, item: Item) -> None: try: sub.loop.call_soon_threadsafe(sub.put, item) except RuntimeError: # loop closed pass def _broadcast_presence(self) -> None: item: Item = ("presence", self.presence()) with self._lock: subs = list(self._subs.values()) for s in subs: self._send(s, item) def _deliver(self, payload: str) -> None: """Fan one NOTIFY payload out. `session` events go only to the subscribers of that session.""" try: data = json.loads(payload) except ValueError: log.warning("ignoring a non-JSON ps_events payload") return if not isinstance(data, dict): return name = str(data.pop("event", "changed")) with self._lock: subs = list(self._subs.values()) if name == "session": target = str(data.get("session_id")) subs = [s for s in subs if s.session_id == target] for s in subs: self._send(s, (name, data)) # ------------------------------------------------ listener thread def start(self) -> None: if self._connect is None or (self._thread and self._thread.is_alive()): return self._stop.clear() self._thread = threading.Thread(target=self._run, name="ps-events", daemon=True) self._thread.start() def stop(self, timeout: float = 5.0) -> None: self._stop.set() if self._thread is not None: self._thread.join(timeout) self._thread = None def _run(self) -> None: assert self._connect is not None delay = 1.0 while not self._stop.is_set(): conn = None try: conn = self._connect() conn.execute(f"LISTEN {CHANNEL}") self.connected = True delay = 1.0 log.info("events hub listening on %s", CHANNEL) while not self._stop.is_set(): for n in conn.notifies(timeout=1.0): self._deliver(n.payload) except Exception as e: if not self._stop.is_set(): log.warning("events hub: %s; retrying in %.0f s", type(e).__name__, delay) finally: self.connected = False if conn is not None: with contextlib.suppress(Exception): conn.close() if self._stop.wait(delay): break delay = min(delay * 2, BACKOFF_MAX) def format_sse(event: str, data: dict[str, Any]) -> str: return f"event: {event}\ndata: {json.dumps(data, default=str, separators=(',', ':'))}\n\n" async def stream( hub: EventHub, sub: Subscriber, *, hello: dict[str, Any], is_disconnected: Callable[[], Awaitable[bool]], ping_seconds: float = PING_SECONDS, ) -> AsyncIterator[str]: """The SSE body: `hello`, then queued events, with a `: ping` comment when idle. Ends on own revoke.""" try: yield format_sse("hello", hello) while True: if await is_disconnected(): break try: name, data = await asyncio.wait_for(sub.queue.get(), timeout=ping_seconds) except TimeoutError: yield ": ping\n\n" continue yield format_sse(name, data) if name == "session" and data.get("session_id") == sub.session_id and data.get("action") == "revoked": break finally: hub.unsubscribe(sub)