iot-dboard/backend/events.py

"""Server-Sent Event bus for pushing updates to displays.

Events are scoped by user_id so that each user's display only receives
updates about their own data (alarms, calendar, pomodoro, ...).

Cross-process broadcast:
The app runs under multiple Gunicorn workers; each worker has its own
in-process EventBus AND its own SSE subscribers. A publish() on worker A
would never reach an SSE connection held by worker B. To fix this,
publish() additionally sends the event over a PostgreSQL LISTEN/NOTIFY
channel; every worker runs a listener thread that re-dispatches remote
events to its local subscribers. Oversized payloads (rare, > ~7KB) are
delivered as a ``db_changed`` signal so the display re-fetches full state.
"""
import json
import os
import queue
import threading

import psycopg2

_NOTIFY_CHANNEL = "iot_events"
_NOTIFY_MAX_LEN = 7000  # stay well under PostgreSQL's 8KB NOTIFY payload limit


class EventBus:
    """Simple pub/sub event system with SSE + cross-process broadcast support.

    publish(event_type, data, user_id=None): when user_id is given, only
    subscribers registered for that user receive the event.
    """

    def __init__(self):
        self._subscribers: list[tuple[int | None, queue.Queue]] = []
        self._lock = threading.Lock()
        self._notify_conn = None
        self._notify_lock = threading.Lock()
        self._listener = None

    def attach(self, db_url: str):
        """Enable cross-process broadcast for THIS process: opens a NOTIFY
        connection and starts a LISTEN thread. Call once per worker (start.py
        / run.py). Without attach, publish() still works locally (tests)."""
        if self._notify_conn is not None:
            return
        try:
            self._notify_conn = psycopg2.connect(db_url, connect_timeout=5)
            self._notify_conn.autocommit = True
        except Exception:
            self._notify_conn = None
            return
        self._listener = threading.Thread(
            target=self._listen_loop, args=(db_url,), daemon=True,
            name="eventbus-listener",
        )
        self._listener.start()

    def subscribe(self, user_id: int | None = None) -> queue.Queue:
        """Subscribe to events. user_id=None receives ALL events (broadcast)."""
        q: queue.Queue = queue.Queue()
        with self._lock:
            self._subscribers.append((user_id, q))
        return q

    def unsubscribe(self, q: queue.Queue):
        with self._lock:
            self._subscribers = [(uid, sq) for uid, sq in self._subscribers if sq is not q]

    def publish(self, event_type: str, data: dict, user_id: int | None = None):
        """Deliver locally AND broadcast to other workers (best-effort).
        The NOTIFY carries this process's PID so our own listener thread can
        skip the echo (each subscriber must receive an event exactly once)."""
        self._dispatch(event_type, data, user_id)
        self._broadcast(event_type, data, user_id)

    # ─── local delivery ───────────────────────
    def _dispatch(self, event_type: str, data: dict, user_id: int | None = None):
        msg = json.dumps({"type": event_type, "data": data})
        with self._lock:
            dead = []
            for sub_uid, q in self._subscribers:
                if sub_uid is not None and user_id is not None and sub_uid != user_id:
                    continue
                try:
                    q.put_nowait(msg)
                except queue.Full:
                    dead.append((sub_uid, q))
            for item in dead:
                self._subscribers.remove(item)

    # ─── cross-process broadcast ──────────────
    def _broadcast(self, event_type: str, data: dict, user_id: int | None):
        if self._notify_conn is None:
            return
        try:
            payload = json.dumps(
                {"pid": os.getpid(), "type": event_type, "data": data,
                 "user_id": user_id},
                default=str,
            )
            if len(payload) > _NOTIFY_MAX_LEN:
                # Too big for NOTIFY: receivers re-fetch full state from the DB.
                payload = json.dumps(
                    {"pid": os.getpid(), "type": event_type,
                     "oversized": True, "user_id": user_id}
                )
            with self._notify_lock:
                cur = self._notify_conn.cursor()
                cur.execute("SELECT pg_notify(%s, %s)", (_NOTIFY_CHANNEL, payload))
                self._notify_conn.commit()
                cur.close()
        except Exception:
            # Broadcast is best-effort (local delivery already happened).
            # Drop the dead connection so the next publish reconnects once.
            try:
                self._notify_conn.close()
            except Exception:
                pass
            self._notify_conn = None

    def handle_remote(self, payload: str):
        """Listener-thread entry: re-dispatch a remote event locally (no NOTIFY).
        Events broadcast by THIS process are skipped (already delivered locally)."""
        try:
            msg = json.loads(payload)
            if msg.get("pid") == os.getpid():
                return  # echo of our own broadcast
            if msg.get("oversized"):
                self._dispatch("db_changed", {}, msg.get("user_id"))
            else:
                self._dispatch(
                    msg.get("type", "db_changed"),
                    msg.get("data") or {},
                    msg.get("user_id"),
                )
        except Exception:
            pass

    def _listen_loop(self, db_url: str):
        """Blocking LISTEN loop with reconnect on DB failure."""
        import select
        import time
        while True:
            conn = None
            try:
                conn = psycopg2.connect(db_url, connect_timeout=5)
                conn.autocommit = True
                cur = conn.cursor()
                cur.execute(f"LISTEN {_NOTIFY_CHANNEL}")
                cur.close()
                while True:
                    if select.select([conn], [], [], 1.0)[0]:
                        conn.poll()
                        while conn.notifies:
                            n = conn.notifies.pop(0)
                            self.handle_remote(n.payload)
            except Exception:
                time.sleep(2)  # DB temporarily unavailable: retry
            finally:
                if conn is not None:
                    try:
                        conn.close()
                    except Exception:
                        pass