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