DocsGPT/docsgpt/streaming/broadcast_channel.py

"""Redis pub/sub Topic abstraction for SSE fan-out.

A Topic is a named channel for one-shot live event delivery. Canonical uses:

- ``user:{user_id}`` for per-user notifications
- ``channel:{message_id}`` for per-chat-message streams

Subscription is race-free via ``on_subscribe``: the callback fires only
after Redis acknowledges ``SUBSCRIBE``, so a publisher dispatched inside
the callback cannot lose its first event to a not-yet-registered
subscriber.

The subscribe iterator yields ``None`` on poll timeout so the caller can
emit SSE keepalive comments without spawning a separate timer thread.
"""

from __future__ import annotations

import logging
from typing import Callable, Iterator, Optional

import redis as redis_lib

from docsgpt.cache import get_pubsub_redis_instance, get_redis_instance

logger = logging.getLogger(__name__)


class Topic:
    """A pub/sub channel identified by a string name."""

    def __init__(self, name: str) -> None:
        self.name = name

    def publish(self, payload: str | bytes) -> int:
        """Fan out a payload to currently subscribed clients.

        Returns the number Redis reports as receiving the message (limited
        to subscribers connected to *this* Redis instance), or 0 if Redis
        is unavailable. Never raises.
        """
        redis = get_redis_instance()
        if redis is None:
            logger.debug("Redis unavailable; dropping publish to %s", self.name)
            return 0
        try:
            return int(redis.publish(self.name, payload))
        except Exception:
            logger.exception("Topic.publish failed for %s", self.name)
            return 0

    def subscribe(
        self,
        on_subscribe: Optional[Callable[[], None]] = None,
        poll_timeout: float = 1.0,
    ) -> Iterator[Optional[bytes]]:
        """Subscribe to the topic; yield raw payloads or ``None`` on tick.

        Yields ``None`` every ``poll_timeout`` seconds while idle so the
        caller can emit keepalive frames or check cancellation. Yields
        ``bytes`` for each delivered message.

        ``on_subscribe`` runs synchronously after Redis acknowledges the
        SUBSCRIBE — use it to seed any state (e.g. read backlog) that
        must be ordered after the subscriber is live but before the
        first pub/sub message is processed.

        If Redis is unavailable, returns immediately without yielding.
        Cleanly unsubscribes on ``GeneratorExit`` (client disconnect).

        Uses the pub/sub-dedicated client (bounded ``socket_timeout``): a
        subscriber whose connection went half-open must fail within seconds
        and release its WSGI thread, not block in ``get_message`` until the
        worker restarts.
        """
        redis = get_pubsub_redis_instance()
        if redis is None:
            logger.debug("Redis unavailable; subscribe to %s yielded nothing", self.name)
            return
        pubsub = None
        on_subscribe_fired = False
        try:
            pubsub = redis.pubsub()
            try:
                pubsub.subscribe(self.name)
            except Exception:
                # Subscribe failure (transient Redis hiccup, conn reset, etc.)
                # is treated like "Redis unavailable": yield nothing, let the
                # caller fall back to its own resilience strategy. The finally
                # block will still tear down the pubsub object cleanly.
                logger.exception("pubsub.subscribe failed for %s", self.name)
                return
            while True:
                try:
                    msg = pubsub.get_message(timeout=poll_timeout)
                except redis_lib.exceptions.TimeoutError:
                    # A bounded read timed out mid-message/health-check: the
                    # connection is half-open (NAT/IPVS dropped the flow).
                    # End the subscription so the SSE client reconnects on a
                    # fresh connection; expected during network churn, so no
                    # stack trace.
                    logger.info(
                        "pubsub read timed out for %s; closing subscriber",
                        self.name,
                    )
                    return
                except Exception:
                    logger.exception("pubsub.get_message failed for %s", self.name)
                    return
                if msg is None:
                    yield None
                    continue
                msg_type = msg.get("type")
                if msg_type == "subscribe":
                    if not on_subscribe_fired and on_subscribe is not None:
                        try:
                            on_subscribe()
                        except Exception:
                            logger.exception(
                                "on_subscribe callback failed for %s", self.name
                            )
                    on_subscribe_fired = True
                    continue
                if msg_type != "message":
                    continue
                data = msg.get("data")
                if data is None:
                    continue
                yield data if isinstance(data, bytes) else str(data).encode("utf-8")
        finally:
            if pubsub is not None:
                if on_subscribe_fired:
                    try:
                        pubsub.unsubscribe(self.name)
                    except Exception:
                        logger.debug(
                            "pubsub unsubscribe error for %s",
                            self.name,
                            exc_info=True,
                        )
                try:
                    pubsub.close()
                except Exception:
                    logger.debug("pubsub close error for %s", self.name, exc_info=True)