mirror of
https://github.com/maziggy/bambuddy.git
synced 2026-09-30 03:01:21 +02:00
maziggy
62a64006b8
The single-connection barrier from the last round was correct and was being bypassed. shutdown_broadcaster() popped the broadcaster out of the registry and only then awaited its teardown, so while the socket was still closing the slot sat empty: a /camera/stream request landing in that window minted a broadcaster with no predecessor and dialled port 6000 immediately. A page reload fires /camera/stop and the new stream request concurrently, so a P1S ended up holding two connections, kept feeding the orphan, and starved the live viewer until its TCP keepalive reaped the dead one ~20 min later. The stopped broadcaster now stays in the registry so the successor chains behind its socket close. The camera page also rendered the <img> src before the stream token arrived whenever auth was disabled, then swapped it once the token landed — aborting the in-flight request and issuing a second one. With auth off both reached the backend, so every load attached two viewers to a one-socket printer. The src now waits for the token query to settle. Subscribers only checked for client disconnect after yielding a frame or on a 30s idle timeout, so a viewer that left during a black stream stayed counted — and /camera/stop trusts that count to decide whether to tear the upstream down.
392 lines
18 KiB
Python
392 lines
18 KiB
Python
"""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")
|