"""MJPEG fan-out broadcaster for camera streams. Most Bambu Lab printers only allow one concurrent camera connection: the RTSP socket on X1/H2/P2 models, the chamber-image socket on port 6000 on A1/P1 models. Without fan-out, opening a second viewer either fails or kicks the first viewer off — see issue #1089. This module owns a single upstream connection per printer and pushes each frame to N independent subscriber queues. New viewers tap the existing upstream; no new printer connection is opened. When the last subscriber leaves, the upstream is torn down after a short grace window so that a quick page refresh or second-tab open does not pay a reconnect. """ from __future__ import annotations import asyncio import logging from collections.abc import AsyncGenerator, Awaitable, Callable logger = logging.getLogger(__name__) # How long to keep the upstream pump alive after the last subscriber leaves. # A short grace window absorbs page refreshes and "open camera in new tab" # without paying a fresh ffmpeg/RTSP handshake (which can take several seconds # on some firmwares and is the very reconnect cost we are trying to avoid). _GRACE_SECONDS = 5.0 # Upper bound on how long a new broadcaster waits for a displaced one to finish # tearing down before proceeding anyway (#2521). Teardown is normally sub-second # (cancel pump + close socket); the cap only guards a wedged upstream close. _TEARDOWN_WAIT_SECONDS = 10.0 # Per-subscriber queue depth. Small on purpose: if a viewer can't keep up # with the printer's frame rate we drop frames for that viewer rather than # blocking the broadcaster. Live video — old frames have no value. _SUBSCRIBER_QUEUE_SIZE = 4 # Sentinel pushed to subscriber queues when the upstream pump exits, so each # subscriber's read loop can break out cleanly instead of hanging on get(). _UPSTREAM_GONE = b"" # How often a subscriber that isn't receiving frames re-checks whether its # client is still connected. Only pays a cost when the stream is *not* producing # frames — the normal path returns from queue.get() as soon as a frame lands and # checks after the yield. Kept short because the subscriber count derived from # it is what /camera/stop uses to decide whether to tear the upstream down. _DISCONNECT_POLL_SECONDS = 1.0 UpstreamFactory = Callable[[asyncio.Event], AsyncGenerator[bytes, None]] class MjpegBroadcaster: """Single upstream MJPEG stream, fanned out to N subscribers.""" def __init__(self, key: str, factory: UpstreamFactory, predecessor: MjpegBroadcaster | None = None) -> None: self._key = key self._factory = factory self._subscribers: list[asyncio.Queue[bytes]] = [] self._lock = asyncio.Lock() self._pump_task: asyncio.Task | None = None self._grace_task: asyncio.Task | None = None # Disconnect event passed to the upstream generator so we can ask it to # stop reconnecting when the last subscriber leaves. self._upstream_disconnect = asyncio.Event() self._stopped = False # Most recent chunk pumped to subscribers. New (late) subscribers are # primed with it so the browser renders a frame immediately instead of # waiting for the next upstream frame — critical on slow chamber-image # cams where the wait looked like a permanent black screen (#2521). self._last_chunk: bytes | None = None # Set once teardown is fully complete (pump cancelled AND the upstream # socket closed). A successor broadcaster waits on this before dialing # so a single-connection printer never sees two sockets at once — the # overlap stranded frames on an orphaned socket for the ~20 min it took # the printer's TCP keepalive to reap it (#2521). self._teardown_complete = asyncio.Event() # The stopped broadcaster this one replaces, if any. The pump waits for # its socket to close before opening ours. Guarding at the pump (not at # get_or_create) keeps it correct when concurrent viewers race to # replace the same stopped broadcaster — only the single pump dials. self._predecessor = predecessor @property def key(self) -> str: return self._key @property def subscriber_count(self) -> int: return len(self._subscribers) @property def stopped(self) -> bool: return self._stopped async def subscribe(self) -> asyncio.Queue[bytes]: """Add a subscriber and ensure the upstream pump is running.""" async with self._lock: if self._stopped: raise RuntimeError(f"broadcaster {self._key!r} is stopped") # Cancel any pending grace-window shutdown — a viewer just rejoined. if self._grace_task is not None and not self._grace_task.done(): self._grace_task.cancel() self._grace_task = None queue: asyncio.Queue[bytes] = asyncio.Queue(maxsize=_SUBSCRIBER_QUEUE_SIZE) self._subscribers.append(queue) # Prime a late joiner with the last frame so it renders instantly # (#2521). The very first subscriber has nothing to prime yet — it # starts the pump below. if self._last_chunk is not None: try: queue.put_nowait(self._last_chunk) except asyncio.QueueFull: # pragma: no cover — fresh queue pass if self._pump_task is None or self._pump_task.done(): # Reset the disconnect signal in case a previous pump set it. self._upstream_disconnect = asyncio.Event() self._pump_task = asyncio.create_task(self._pump(), name=f"camera-fanout-pump-{self._key}") return queue async def unsubscribe(self, queue: asyncio.Queue[bytes]) -> int: """Remove a subscriber and return the remaining count (atomic). If this was the last subscriber, schedule grace shutdown. """ async with self._lock: try: self._subscribers.remove(queue) except ValueError: return len(self._subscribers) # Already removed (e.g. force_shutdown) remaining = len(self._subscribers) if remaining == 0 and not self._stopped: # Last subscriber left. Schedule grace-window teardown. self._grace_task = asyncio.create_task(self._grace_then_stop(), name=f"camera-fanout-grace-{self._key}") return remaining async def force_shutdown(self) -> None: """Tear down immediately, kick all subscribers. Idempotent.""" pump_task = await self._mark_stopped_locked(notify_subscribers=True) await self._await_pump_cancellation(pump_task) # Upstream socket is now closed (pump's finally ran) — release anyone # waiting to open a replacement broadcaster (#2521). self._teardown_complete.set() async def wait_until_torn_down(self) -> None: """Block until this broadcaster's upstream socket has fully closed. Only meaningful for a stopped broadcaster; on a live one this never returns. get_or_create_broadcaster gates a replacement on it so the old and new upstream sockets never overlap (#2521). """ await self._teardown_complete.wait() async def _grace_then_stop(self) -> None: try: await asyncio.sleep(_GRACE_SECONDS) except asyncio.CancelledError: return # New subscriber arrived during grace # Re-check under the lock — a subscriber may have rejoined between # the sleep finishing and us acquiring the lock. pump_task: asyncio.Task | None = None async with self._lock: if self._subscribers or self._stopped: return self._upstream_disconnect.set() pump_task = self._pump_task self._pump_task = None self._grace_task = None self._stopped = True await self._await_pump_cancellation(pump_task) # Upstream socket is now closed — release any pending replacement (#2521). self._teardown_complete.set() async def _mark_stopped_locked(self, *, notify_subscribers: bool) -> asyncio.Task | None: """Mark the broadcaster stopped and detach the pump task. Caller MUST NOT hold ``self._lock`` (we acquire it here). Returns the pump task (if any) so the caller can await its cancellation OUTSIDE the lock — the pump's ``finally`` block needs the lock to wake up subscribers, so we'd deadlock if we awaited it under the lock. """ async with self._lock: if self._stopped and self._pump_task is None: return None self._upstream_disconnect.set() if notify_subscribers: for queue in self._subscribers: try: queue.put_nowait(_UPSTREAM_GONE) except asyncio.QueueFull: pass self._subscribers.clear() pump_task = self._pump_task self._pump_task = None self._stopped = True if self._grace_task is not None and not self._grace_task.done(): self._grace_task.cancel() self._grace_task = None return pump_task async def _await_pump_cancellation(self, pump_task: asyncio.Task | None) -> None: if pump_task is None or pump_task.done(): return pump_task.cancel() try: await pump_task except (asyncio.CancelledError, Exception): # Pump exceptions are already logged inside _pump; swallow here so # teardown can never propagate a stray crash. pass async def _pump(self) -> None: """Drive the upstream generator and broadcast each chunk.""" try: # Don't dial the printer until the broadcaster we're replacing has # closed its socket (#2521). Bounded so a wedged teardown degrades # to the old overlap behaviour rather than never producing a frame. predecessor = self._predecessor self._predecessor = None if predecessor is not None: try: await asyncio.wait_for(predecessor.wait_until_torn_down(), timeout=_TEARDOWN_WAIT_SECONDS) except asyncio.TimeoutError: logger.warning("Prior broadcaster %r didn't tear down in time; dialing anyway", self._key) async for chunk in self._factory(self._upstream_disconnect): # Snapshot subscribers under lock so we don't iterate a list # mutated by subscribe()/unsubscribe() while we are putting. # Remember the frame under the same lock so subscribe() can # prime a late joiner with a consistent last-chunk value (#2521). async with self._lock: self._last_chunk = chunk targets = list(self._subscribers) for queue in targets: try: queue.put_nowait(chunk) except asyncio.QueueFull: # Slow viewer — drop this frame for them. They'll catch # up on the next frame. Don't unsubscribe: a brief # browser stall shouldn't end the stream. pass except asyncio.CancelledError: raise except Exception: logger.exception("Camera fan-out pump crashed for %s", self._key) finally: # Pump is exiting — wake up any subscribers still hanging on get(). async with self._lock: for queue in self._subscribers: try: queue.put_nowait(_UPSTREAM_GONE) except asyncio.QueueFull: pass # Global registry. Keyed by printer_id (as str) so a chamber-mode printer # and an RTSP-mode printer can never collide on the same key. _broadcasters: dict[str, MjpegBroadcaster] = {} _registry_lock = asyncio.Lock() async def get_or_create_broadcaster(key: str, factory: UpstreamFactory) -> MjpegBroadcaster: """Return the live broadcaster for `key`, creating one if needed. A broadcaster that has been stopped (force shutdown or grace timeout) is replaced with a fresh instance — the caller will subscribe to the new one. When replacing a stopped broadcaster, the fresh instance is handed it as a predecessor: its pump waits for the old socket to close before dialing, so a single-connection cam (chamber-image port 6000) never sees two sockets at once. Otherwise the printer keeps feeding the orphaned socket and starves the new one until its TCP keepalive reaps it, ~20 min later (#2521). """ async with _registry_lock: existing = _broadcasters.get(key) if existing is not None and not existing.stopped: return existing # `existing` (if any) is stopped/tearing down — chain the new pump # behind its socket close. new_bc = MjpegBroadcaster(key, factory, predecessor=existing) _broadcasters[key] = new_bc return new_bc async def shutdown_broadcaster(key: str) -> bool: """Force-shutdown the broadcaster for `key`. Returns True if one was running. The stopped broadcaster stays in the registry on purpose. It used to be popped *before* ``force_shutdown()`` was awaited, which vacated the slot while the upstream socket was still closing: a ``/camera/stream`` request landing in that window found nothing, minted a broadcaster with ``predecessor=None``, and dialled the printer immediately. That is exactly the two-sockets-at-once overlap the predecessor gate exists to prevent — the gate only engages when the stopped broadcaster is still *findable*, and popping it here bypassed the gate in the one case it was written for. A page reload fires ``/camera/stop`` and the new stream request concurrently, so a single-connection cam (chamber-image port 6000) ended up with an orphaned socket that the printer kept feeding, starving the live viewer until the printer's TCP keepalive reaped it ~20 min later (#2521). Leaving it in place is safe: ``get_or_create_broadcaster`` replaces a stopped entry (chaining the successor behind its teardown), ``get_subscriber_count`` reports 0 for it, and ``active_broadcaster_keys`` filters it out. There is at most one entry per printer, and it is overwritten by the next viewer. """ async with _registry_lock: bc = _broadcasters.get(key) if bc is None or bc.stopped: return False await bc.force_shutdown() return True async def shutdown_all_broadcasters() -> None: """Tear down every broadcaster (for app shutdown).""" async with _registry_lock: bcs = list(_broadcasters.values()) _broadcasters.clear() await asyncio.gather(*(bc.force_shutdown() for bc in bcs), return_exceptions=True) def active_broadcaster_keys() -> list[str]: """Snapshot of keys with a live (non-stopped) broadcaster. For diagnostics.""" return [k for k, bc in _broadcasters.items() if not bc.stopped] def get_subscriber_count(key: str) -> int: """Return the number of live subscribers attached to ``key``, or 0. Used by ``/camera/stop`` to decide whether to force-shutdown the broadcaster or defer to natural cleanup. Other viewers (cam-wall tile, embedded viewer, popup window) all subscribe to the same broadcaster, so a force-shutdown triggered by one leaving viewer would kill the others' streams. """ bc = _broadcasters.get(key) if bc is None or bc.stopped: return 0 return bc.subscriber_count # --------------------------------------------------------------------------- # AsyncGenerator helper — turns a subscriber queue into an async generator # that yields MJPEG chunks until the upstream signals it's gone. # --------------------------------------------------------------------------- async def iter_subscriber( broadcaster: MjpegBroadcaster, queue: asyncio.Queue[bytes], *, is_disconnected: Callable[[], Awaitable[bool]] | None = None, on_unsubscribe: Callable[[int], None] | None = None, ) -> AsyncGenerator[bytes, None]: """Yield chunks from a subscriber queue until upstream ends or client leaves. Always unsubscribes from the broadcaster on exit, even on cancellation. The optional ``on_unsubscribe`` callback receives the post-unsubscribe subscriber count — useful for accurate detach-log lines that don't race with concurrent unsubscribes. """ try: while True: try: chunk = await asyncio.wait_for(queue.get(), timeout=_DISCONNECT_POLL_SECONDS) except asyncio.TimeoutError: # No frame this tick — is the client still there? This used to # wait 30 s before asking, and the disconnect check after a yield # only fires when frames are actually flowing. So a viewer that # went away while the stream was black stayed *counted* as a # subscriber for up to half a minute — and ``/camera/stop`` # trusts that count to decide whether to tear the upstream down, # so a phantom subscriber could make it skip teardown entirely # (#2521). Poll often enough that the count means something. if is_disconnected is not None and await is_disconnected(): break continue if chunk == _UPSTREAM_GONE: break yield chunk if is_disconnected is not None and await is_disconnected(): break finally: remaining = await broadcaster.unsubscribe(queue) if on_unsubscribe is not None: try: on_unsubscribe(remaining) except Exception: logger.exception("on_unsubscribe callback raised")