Fix SSE fanout packet loss on reconnect

utils/sse.py

@@ -27,6 +27,14 @@ _fanout_channels_lock = threading.Lock()
 def _run_fanout(channel: _QueueFanoutChannel) -> None:
     """Drain source queue and fan out each message to all subscribers."""
     while True:
+        with channel.lock:
+            has_subscribers = bool(channel.subscribers)
+
+        if not has_subscribers:
+            # Do not drain source_queue when no clients are connected.
+            time.sleep(channel.source_timeout)
+            continue
+
         try:
             msg = channel.source_queue.get(timeout=channel.source_timeout)
         except queue.Empty:
@@ -35,6 +43,14 @@ def _run_fanout(channel: _QueueFanoutChannel) -> None:
         with channel.lock:
             subscribers = tuple(channel.subscribers)
 
+        if not subscribers:
+            # Subscriber set changed after we dequeued; requeue best-effort.
+            try:
+                channel.source_queue.put_nowait(msg)
+            except queue.Full:
+                pass
+            continue
+
         for subscriber in subscribers:
             try:
                 subscriber.put_nowait(msg)
@@ -52,13 +68,24 @@ def _ensure_fanout_channel(
     source_queue: queue.Queue,
     source_timeout: float,
 ) -> _QueueFanoutChannel:
-    """Get/create a fanout channel and ensure distributor thread is running."""
+    """Get/create a fanout channel."""
     with _fanout_channels_lock:
         channel = _fanout_channels.get(channel_key)
         if channel is None:
             channel = _QueueFanoutChannel(source_queue=source_queue, source_timeout=source_timeout)
             _fanout_channels[channel_key] = channel
 
+        if channel.source_queue is not source_queue:
+            # Keep channel in sync if source queue object is replaced.
+            channel.source_queue = source_queue
+        channel.source_timeout = source_timeout
+
+    return channel
+
+
+def _ensure_distributor_running(channel: _QueueFanoutChannel, channel_key: str) -> None:
+    """Ensure fanout distributor thread is running for a channel."""
+    with _fanout_channels_lock:
         if channel.distributor is None or not channel.distributor.is_alive():
             channel.distributor = threading.Thread(
                 target=_run_fanout,
@@ -68,8 +95,6 @@ def _ensure_fanout_channel(
             )
             channel.distributor.start()
 
-    return channel
-
 
 def subscribe_fanout_queue(
     source_queue: queue.Queue,
@@ -89,6 +114,9 @@ def subscribe_fanout_queue(
     with channel.lock:
         channel.subscribers.add(subscriber)
 
+    # Start distributor only after subscriber is registered to avoid initial-loss race.
+    _ensure_distributor_running(channel, channel_key)
+
     def _unsubscribe() -> None:
         with channel.lock:
             channel.subscribers.discard(subscriber)