PacketStation/server/packetstation/events.py

"""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
    packet_id: int | None = None  # the packet the subscriber has open, for "Also here" on the packet page
    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,
        packet_id: int | 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,
            packet_id=packet_id,
        )
        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, int | None], dict[str, Any]] = {}
        for s in subs:
            seen.setdefault(
                (s.user_id, s.screen, s.packet_id),
                {"user_id": s.user_id, "display_name": s.display_name, "screen": s.screen, "packet_id": s.packet_id},
            )
        users = sorted(seen.values(), key=lambda u: (u["display_name"].lower(), u["screen"] or "", u["packet_id"] or 0))
        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)