from __future__ import annotations import logging import queue import threading from collections.abc import Iterator from typing import Self from extensions.redis_names import serialize_redis_name from libs.broadcast_channel.channel import Producer, Subscriber, Subscription from libs.broadcast_channel.exc import SubscriptionClosedError from redis import Redis, RedisCluster logger = logging.getLogger(__name__) class StreamsBroadcastChannel: """ Redis Streams based broadcast channel implementation. Characteristics: - At-least-once delivery for late subscribers within the stream retention window. - Each topic is stored as a dedicated Redis Stream key. - The stream key expires `retention_seconds` after the last event is published (to bound storage). """ def __init__( self, redis_client: Redis | RedisCluster, *, retention_seconds: int = 600, join_timeout_ms: int = 2000, ): self._client = redis_client self._retention_seconds = max(int(retention_seconds or 0), 0) # Max time close() will wait for the listener thread to finish. # See `_StreamsSubscription._join_timeout_ms` for the rationale. self._join_timeout_ms = max(int(join_timeout_ms or 0), 0) def topic(self, topic: str) -> StreamsTopic: return StreamsTopic( self._client, topic, retention_seconds=self._retention_seconds, join_timeout_ms=self._join_timeout_ms, ) class StreamsTopic: def __init__( self, redis_client: Redis | RedisCluster, topic: str, *, retention_seconds: int = 600, join_timeout_ms: int = 2000, ): self._client = redis_client self._topic = topic self._key = serialize_redis_name(f"stream:{topic}") self._retention_seconds = retention_seconds self._join_timeout_ms = max(int(join_timeout_ms or 0), 0) self.max_length = 5000 def as_producer(self) -> Producer: return self def publish(self, payload: bytes) -> None: self._client.xadd(self._key, {b"data": payload}, maxlen=self.max_length) if self._retention_seconds > 0: try: self._client.expire(self._key, self._retention_seconds) except Exception as e: logger.warning("Failed to set expire for stream key %s: %s", self._key, e, exc_info=True) def as_subscriber(self) -> Subscriber: return self def subscribe(self) -> Subscription: return _StreamsSubscription(self._client, self._key, join_timeout_ms=self._join_timeout_ms) class _StreamsSubscription(Subscription): _SENTINEL = object() def __init__(self, client: Redis | RedisCluster, key: str, *, join_timeout_ms: int = 2000): self._client = client self._key = key # Max time close() will wait for the listener thread to finish before # returning. Bounds SSE close tail latency: the listener blocks on # XREAD with BLOCK=1000ms, so close() naturally waits up to ~1s for # the thread to notice _closed. Setting this lower lets close() # return promptly while the daemon listener exits on its own within # one BLOCK window - safe because the listener holds no critical # state. ``0`` means close() does not wait at all. self._join_timeout_ms = max(int(join_timeout_ms or 0), 0) self._queue: queue.Queue[object] = queue.Queue() # The `_lock` lock is used to # # 1. protect the _listener attribute # 2. prevent repeated releases of underlying resoueces. (The _closed flag.) # # INVARIANT: the implementation must hold the lock while # reading and writing the _listener / `_closed` attribute. self._lock = threading.Lock() self._closed: bool = False # self._closed = threading.Event() self._listener: threading.Thread | None = None def _listen(self) -> None: """The `_listen` method handles the message retrieval loop. It requires a dedicated thread and is not intended for direct invocation. The thread is started by `_start_if_needed`. """ # since this method runs in a dedicated thread, acquiring `_lock` inside this method won't cause # deadlock. # Setting initial last id to `$` to signal redis that we only want new messages. # # ref: https://redis.io/docs/latest/commands/xread/#the-special--id last_id = "$" try: while True: with self._lock: if self._closed: break streams = self._client.xread({self._key: last_id}, block=1000, count=100) if not streams: continue for _, entries in streams: for entry_id, fields in entries: data = None if isinstance(fields, dict): data = fields.get(b"data") data_bytes: bytes | None = None if isinstance(data, str): data_bytes = data.encode() elif isinstance(data, (bytes, bytearray)): data_bytes = bytes(data) if data_bytes is not None: self._queue.put_nowait(data_bytes) last_id = entry_id finally: self._queue.put_nowait(self._SENTINEL) with self._lock: self._listener = None self._closed = True def _start_if_needed(self) -> None: """This method must be called with `_lock` held.""" if self._listener is not None: return # Ensure only one listener thread is created under concurrent calls if self._listener is not None or self._closed: return self._listener = threading.Thread( target=self._listen, name=f"redis-streams-sub-{self._key}", daemon=True, ) self._listener.start() def __iter__(self) -> Iterator[bytes]: # Iterator delegates to receive with timeout; stops on closure. with self._lock: self._start_if_needed() while True: with self._lock: if self._closed: return try: item = self.receive(timeout=1) except SubscriptionClosedError: return if item is not None: yield item def receive(self, timeout: float | None = 0.1) -> bytes | None: with self._lock: if self._closed: raise SubscriptionClosedError("The Redis streams subscription is closed") self._start_if_needed() try: if timeout is None: item = self._queue.get() else: item = self._queue.get(timeout=timeout) except queue.Empty: return None if item is self._SENTINEL: raise SubscriptionClosedError("The Redis streams subscription is closed") assert isinstance(item, (bytes, bytearray)), "Unexpected item type in stream queue" return bytes(item) def close(self) -> None: with self._lock: if self._closed: return self._closed = True listener = self._listener if listener is not None: self._listener = None # We close the listener outside of the with block to avoid holding the # lock for a long time. if listener is not None and listener.is_alive(): listener.join(timeout=self._join_timeout_ms / 1000.0) if listener.is_alive(): logger.debug( "Streams subscription listener for key %s did not stop within %dms; " "daemon thread will exit on its own within one poll window.", self._key, self._join_timeout_ms, ) # Context manager helpers def __enter__(self) -> Self: with self._lock: self._start_if_needed() return self def __exit__(self, exc_type, exc_value, traceback) -> bool | None: self.close() return None