"""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): one entry per (user, screen, packet) with the display name and the role label. A locked session drops out of presence until it is unlocked; a revoked one drops out at once. A browser that goes away drops out when its stream ends (the stream's `finally`, at the latest on the next ping). """ 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 datetime import UTC, datetime 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 packet_id: int | None = None # the packet the subscriber has open, for "Also here" on the packet page role: str = "" # the label shown beside the name: the screen open, else the person's primary role since: datetime = field(default_factory=lambda: datetime.now(UTC)) locked: bool = False # the session is locked: left out of presence until it is unlocked revoked: bool = False # the session was revoked: left out of presence while the stream winds down id: str = field(default_factory=lambda: uuid.uuid4().hex) closed: bool = False @property def present(self) -> bool: return not (self.closed or self.locked or self.revoked) 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, packet_id: int | None = None, role: str = "", locked: bool = False, ) -> 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, packet_id=packet_id, role=role, locked=locked, ) 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 _users(self) -> list[dict[str, Any]]: """One entry per (user, screen, packet) among the present subscribers, the earliest `since` kept.""" with self._lock: subs = [s for s in self._subs.values() if s.present] seen: dict[tuple[str, str | None, int | None], dict[str, Any]] = {} for s in sorted(subs, key=lambda x: x.since): seen.setdefault( (s.user_id, s.screen, s.packet_id), { "user_id": s.user_id, "display_name": s.display_name, "role": s.role, "screen": s.screen, "packet_id": s.packet_id, "since": s.since.isoformat(), }, ) return sorted(seen.values(), key=lambda u: (u["display_name"].lower(), u["screen"] or "", u["packet_id"] or 0)) def presence(self) -> dict[str, Any]: """{users: [{user_id, display_name, role, screen, packet_id, since}]}, locked and revoked sessions left out.""" return {"users": self._users()} def presence_by_screen(self) -> dict[str, list[dict[str, Any]]]: """The same users grouped by the screen they have open ("none" for a stream without a screen), for the Admin "who's in" view.""" out: dict[str, list[dict[str, Any]]] = {} for u in self._users(): out.setdefault(u["screen"] or "none", []).append(u) return out 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()) changed = False if name == "session": target = str(data.get("session_id")) subs = [s for s in subs if s.session_id == target] changed = self._mark(subs, str(data.get("action") or "")) for s in subs: self._send(s, (name, data)) if changed: self._broadcast_presence() def _mark(self, subs: list[Subscriber], action: str) -> bool: """Apply a session event to that session's subscribers. True when presence changed.""" changed = False with self._lock: for s in subs: before = s.present if action == "locked": s.locked = True elif action == "unlocked": s.locked = False elif action == "revoked": s.revoked = True changed = changed or before != s.present return changed # ------------------------------------------------ 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)