diff --git a/.ai/ARCHITECTURE.md b/.ai/ARCHITECTURE.md index d4bd2ace56..36810ec61a 100644 --- a/.ai/ARCHITECTURE.md +++ b/.ai/ARCHITECTURE.md @@ -959,7 +959,13 @@ the replay window; a subscriber past the trimmed floor gets `SharedStorageGap`). Redis is the single source of truth, so no owner election. Pass a `url` (defaults to `$REDIS_URL`) or an existing `client` to reuse a connection pool. This is the only backend correct for **horizontally-scaled, multi-pod** -deployments. +deployments. On an event loop it uses `redis.asyncio` (one client per loop), +and one reader task per loop serves every async subscription with a single +multi-stream `XREAD`; a new subscription interrupts it through a private wake +stream. Never park a blocking read in the executor per subscription: at a dozen +open downlinks that starved every other store call in the worker (25-50% failed +streams at 25 browsers on uvicorn `--workers 4`). Diskcache's async +subscriptions likewise poll without blocking and sleep on the loop. ### Deployment Topology and Shared Storage @@ -1399,7 +1405,10 @@ back to a downlink of its own. - `stream` (ASGI: Quart, FastAPI): one long-lived NDJSON response per browser. It costs no thread -- the subscription parks the task on a future the store resolves (`StoreEngine.apoll`; asyncio streams to the owner from - other workers) -- so a single uvicorn worker holds thousands. + other workers; the per-loop Redis reader) -- so a single uvicorn worker + holds thousands. The downlink's connection record is written with the + store's async calls, and its closed record from a task of its own, since a + client disconnect cancels the response. - `poll` (WSGI: Flask): a WSGI response holds a worker thread for its whole life, so an open downlink per browser exhausts a thread pool at a few dozen browsers (gunicorn `--threads 2`: the second browser hung everything). @@ -1478,6 +1487,7 @@ the right tool: constant latency and a fraction of the CPU. - `dash/_streaming.py` - `StreamedCallbackResponse` marker, context-safe iteration, NDJSON helpers, keepalives, shutdown flag - `dash/_stream_hub.py` - multiplexed transport: `Downlink`, `poll_downlink`, pumps (`pump_to_storage`/`apump_to_storage`), `cancel_stream`, `shutdown_active_streams`/`install_stream_shutdown_handler` - `dash/_shared_storage/_engine.py`, `local.py` - `poll`/`apoll`, loop-native `aget`/`aset`/`apublish`, async client connection +- `dash/_shared_storage/redis.py` - `redis.asyncio` path, `_StreamReader` (one multi-stream `XREAD` per loop) - `dash/backends/_flask.py`, `_quart.py`, `_fastapi.py` - streaming dispatch branches - `dash/backends/ws.py` - `make_stream_frame_emitter`, `consume_stream_frames`/`aconsume_stream_frames` - `dash/dash-renderer/src/actions/callbacks.ts` - `applyStreamFrame`, NDJSON reader, WS frame handling diff --git a/dash/_shared_storage/_polling.py b/dash/_shared_storage/_polling.py index d7eb5f2ddb..e5b254b5bc 100644 --- a/dash/_shared_storage/_polling.py +++ b/dash/_shared_storage/_polling.py @@ -6,6 +6,11 @@ ``SharedStorageGap``. ``LocalSharedStorage`` has its own reconnecting subscription (it also handles owner re-election); this serves the diskcache and Redis backends, whose store is always reachable. + +Async iteration never parks a blocking poll in the executor: that holds a +thread per subscription, and a dozen subscriptions starve every other +``run_in_executor`` call in the process. It polls without waiting and sleeps on +the loop between polls instead. """ import asyncio import threading @@ -21,12 +26,18 @@ class PollingSubscription(Subscription): """Iterate a topic by repeatedly polling ``poll_fn`` from a moving cursor.""" def __init__( - self, topic: str, start_seq: int, poll_fn: PollFn, poll_timeout: float + self, + topic: str, + start_seq: int, + poll_fn: PollFn, + poll_timeout: float, + poll_interval: float, ): self._topic = topic self._cursor = start_seq self._poll_fn = poll_fn self._poll_timeout = poll_timeout + self._poll_interval = poll_interval self._closed = threading.Event() def close(self) -> None: @@ -71,17 +82,16 @@ async def _aiter_with_seq(self): while not self._closed.is_set(): try: res = await loop.run_in_executor( - None, - self._poll_fn, - self._topic, - self._cursor, - self._poll_timeout, + None, self._poll_fn, self._topic, self._cursor, 0.0 ) except RuntimeError: # The loop/executor is shutting down -- end cleanly. break if res.gap: raise self._gap() + if not res.messages: + await asyncio.sleep(self._poll_interval) + continue for pair in self._with_seq(res): yield pair self._cursor = res.last_seq diff --git a/dash/_shared_storage/diskcache.py b/dash/_shared_storage/diskcache.py index 8d909ab80e..b127a22d35 100644 --- a/dash/_shared_storage/diskcache.py +++ b/dash/_shared_storage/diskcache.py @@ -138,4 +138,6 @@ def _poll(self, topic: str, after_seq: int, timeout: float) -> PollResult: def subscribe(self, topic: str, replay_from: Optional[int] = None) -> Subscription: start = replay_from if replay_from is not None else self._head(topic) - return PollingSubscription(topic, start, self._poll, _POLL_TIMEOUT) + return PollingSubscription( + topic, start, self._poll, _POLL_TIMEOUT, _POLL_INTERVAL + ) diff --git a/dash/_shared_storage/redis.py b/dash/_shared_storage/redis.py index 84e5175324..bc5a3024a5 100644 --- a/dash/_shared_storage/redis.py +++ b/dash/_shared_storage/redis.py @@ -11,18 +11,31 @@ ``XADD``) so concurrent publishers stay strictly ordered; the stream is capped to a bounded window (``MAXLEN``), and a subscriber that falls past the trimmed floor gets a ``SharedStorageGap``. + +On an event loop everything goes through ``redis.asyncio``, never the executor: +a blocking XREAD per subscription would hold one executor thread each, and at a +dozen open downlinks every other store call in the worker queues behind them. +Instead each loop gets one reader task that serves all of its subscriptions with +a single multi-stream XREAD. """ +import asyncio import os -from typing import Any, Optional +import secrets +import threading +from typing import Any, Dict, List, Optional, Set, Tuple from ._codec import decode, encode from ._engine import DEFAULT_BUFFER, PollResult from ._polling import PollingSubscription -from .base import BaseSharedStorage, Subscription +from .base import BaseSharedStorage, SharedStorageError, Subscription # Redis Stream XREAD blocks server-side, so a longer cycle than the diskcache # sleep-poll is fine; close() latency is bounded by this. _POLL_TIMEOUT = 5.0 +# The async reader's XREAD block. A new subscription wakes it at once through +# its wake stream; this only bounds how long it lingers once all are gone. +_READER_BLOCK_MS = 5000 +_WAKE_TTL_MS = 60_000 _DEFAULT_URL = "redis://localhost:6379" # Allocate the next sequence and append atomically, so concurrent publishers @@ -59,6 +72,210 @@ def _payload_of(fields: dict) -> bytes: return fields[b"m"] if b"m" in fields else fields["m"] +def _key_of(key: Any) -> str: + return key.decode() if isinstance(key, bytes) else key + + +class RedisSubscription(PollingSubscription): + """A Redis topic subscription. Sync iteration long-polls with XREAD on the + calling thread; async iteration is fed by the loop's :class:`_StreamReader`. + """ + + # pylint: disable=protected-access + + def __init__(self, storage: "RedisSharedStorage", topic: str, start_seq: int): + super().__init__(topic, start_seq, storage._poll, _POLL_TIMEOUT, 0.0) + self._storage = storage + self._stream_key = storage._stream(topic) + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._ready: Optional[asyncio.Event] = None + self._pending: List[Tuple[int, Any]] = [] + self._error: Optional[BaseException] = None + + @property + def cursor(self) -> int: + return self._cursor + + @property + def failed(self) -> bool: + return self._error is not None + + def close(self) -> None: + super().close() + loop, ready = self._loop, self._ready + if loop is not None and ready is not None: + try: + loop.call_soon_threadsafe(ready.set) + except RuntimeError: + pass # the loop already closed + + def deliver(self, entries: List[Tuple[int, Any]]) -> None: + """Called by the reader, on the loop, with entries in stream order.""" + for seq, message in entries: + if seq <= self._cursor: + continue + if ( + seq != self._cursor + 1 + or len(self._pending) >= self._storage._buffer_size + ): + # Trimmed before the reader got to it, or the consumer fell a + # whole buffer behind: either way messages are lost. + self.fail(self._gap()) + return + self._pending.append((seq, message)) + self._cursor = seq + self._wake() + + def fail(self, error: BaseException) -> None: + if self._error is None: + self._error = error + self._wake() + + def _wake(self) -> None: + if self._ready is not None: + self._ready.set() + + async def _aiter_with_seq(self): + self._loop = asyncio.get_running_loop() + self._ready = asyncio.Event() + if self._closed.is_set(): + return + reader = self._storage._loop_state().reader + try: + if self._cursor > await self._storage._ahead(self._topic): + raise self._gap() + await reader.add(self) + while not self._closed.is_set(): + if self._pending: + pairs, self._pending = self._pending, [] + for pair in pairs: + yield pair + continue + if self._error is not None: + raise self._error + self._ready.clear() + await self._ready.wait() + finally: + reader.remove(self) + self.close() + + +class _StreamReader: + """Serves every async subscription of one storage on one event loop. + + One task XREADs all subscribed streams at once, each from the lowest cursor + among its subscribers, and hands the entries to each subscription. A new + subscription interrupts a blocked XREAD through a private wake stream, so it + is never stuck behind the block timeout. + """ + + # pylint: disable=protected-access + + def __init__(self, storage: "RedisSharedStorage", client: Any): + self._storage = storage + self._client = client + self._subs: Dict[str, Set[RedisSubscription]] = {} + self._topics: Dict[str, str] = {} + self._wake_key = f"{storage._prefix}:wake:{secrets.token_hex(8)}" + self._wake_id = "0-0" + self._task: Optional["asyncio.Task[None]"] = None + + async def add(self, sub: RedisSubscription) -> None: + self._subs.setdefault(sub._stream_key, set()).add(sub) + self._topics[sub._stream_key] = sub._topic + if self._task is None: + self._task = asyncio.ensure_future(self._run()) + return + async with self._client.pipeline(transaction=False) as pipe: + pipe.xadd(self._wake_key, {"w": 1}, maxlen=1, approximate=False) + pipe.pexpire(self._wake_key, _WAKE_TTL_MS) + await pipe.execute() + + def remove(self, sub: RedisSubscription) -> None: + subs = self._subs.get(sub._stream_key) + if subs is not None: + subs.discard(sub) + if not subs: + del self._subs[sub._stream_key] + del self._topics[sub._stream_key] + + def _starts(self) -> Dict[str, int]: + """Where to read each stream from: the lowest cursor among its live + subscribers. A failed one waits for its consumer to drop it, and must + not hold the read back meanwhile.""" + starts = {} + for key, subs in self._subs.items(): + cursors = [sub.cursor for sub in subs if not sub.failed] + if cursors: + starts[key] = min(cursors) + return starts + + async def _check_heads(self, starts: Dict[str, int]) -> None: + """On a quiet cycle, catch a stream whose sequence was reset under its + subscribers (the keys flushed or evicted): XREAD from their cursor would + wait until the new sequence climbs back past it.""" + keys = [key for key in starts if key in self._topics] + async with self._client.pipeline(transaction=False) as pipe: + for key in keys: + pipe.get(self._storage._seq(self._topics[key])) + heads = await pipe.execute() + for key, raw in zip(keys, heads): + if starts[key] > (int(raw) if raw is not None else 0): + for sub in self._subs.get(key, ()): + sub.fail(sub._gap()) + + def _dispatch(self, result: Any, starts: Dict[str, int]) -> None: + for key, items in result: + key = _key_of(key) + if key == self._wake_key: + self._wake_id = items[-1][0] + continue + entries = [ + (_seq_of(entry_id), decode(_payload_of(fields))) + for entry_id, fields in items + ] + for sub in self._subs.get(key, ()): + # Joined behind this read's start while it was in flight: the + # next read starts from its cursor. + if sub.cursor >= starts[key]: + sub.deliver(entries) + + def _fail_all(self, error: BaseException) -> None: + for subs in self._subs.values(): + for sub in subs: + sub.fail(error) + self._subs.clear() + self._topics.clear() + + async def _run(self) -> None: + error: Optional[BaseException] = None + try: + while True: + starts = self._starts() + if not starts: + return + streams = {key: f"{start}-0" for key, start in starts.items()} + streams[self._wake_key] = self._wake_id + result = await self._client.xread(streams, block=_READER_BLOCK_MS) + if result: + self._dispatch(result, starts) + else: + await self._check_heads(starts) + except asyncio.CancelledError: + error = SharedStorageError("the subscription reader stopped") + raise + except Exception as exc: # pylint: disable=broad-except + # The store is unreachable: end every subscription with the error, + # as a sync XREAD would, so each consumer can reconnect. + error = exc + finally: + # No await from here on: an add() that lands after this sees no + # reader and starts a fresh one. + self._task = None + if error is not None: + self._fail_all(error) + + class RedisSharedStorage(BaseSharedStorage): """Shared storage backed by Redis. Works across processes and containers. @@ -84,9 +301,14 @@ def __init__( url = url or os.environ.get("REDIS_URL") or _DEFAULT_URL self._redis = redis.Redis.from_url(url) self._owns_client = True + self._url = url if client is None else None self._prefix = key_prefix self._buffer_size = buffer_size self._publish_script: Any = None + # One redis.asyncio client (and reader) per event loop: their + # connections belong to the loop that opened them. + self._loops: Dict[asyncio.AbstractEventLoop, _LoopState] = {} + self._loops_lock = threading.Lock() def start(self) -> None: if self._publish_script is None: @@ -98,6 +320,43 @@ def close(self) -> None: self._redis.close() except Exception: # pylint: disable=broad-except pass + with self._loops_lock: + states = list(self._loops.items()) + self._loops.clear() + for loop, state in states: + # A closed loop's connections can no longer be closed politely. + if not loop.is_closed(): + try: + asyncio.run_coroutine_threadsafe(state.aclose(), loop) + except RuntimeError: + pass # it closed just now + + def _new_async_client(self) -> Any: + _require_redis() + from redis import asyncio as aredis # type: ignore[import-not-found,import-untyped] # pylint: disable=import-outside-toplevel,import-error + + if self._url is not None: + return aredis.Redis.from_url(self._url) + # Same server and options as the client we were handed. + pool = self._redis.connection_pool + connection_class = getattr( + aredis.connection, pool.connection_class.__name__, aredis.Connection + ) + return aredis.Redis( + connection_pool=aredis.ConnectionPool( + connection_class=connection_class, **pool.connection_kwargs + ) + ) + + def _loop_state(self) -> "_LoopState": + loop = asyncio.get_running_loop() + with self._loops_lock: + state = self._loops.get(loop) + if state is None: + for old in [old for old in self._loops if old.is_closed()]: + del self._loops[old] + state = self._loops[loop] = _LoopState(self, self._new_async_client()) + return state # --- key layout -------------------------------------------------------- def _kv(self, key: str) -> str: @@ -123,6 +382,17 @@ def set(self, key: str, value: Any, ttl: Optional[float] = None) -> None: def delete(self, key: str) -> None: self._redis.delete(self._kv(key)) + async def aget(self, key: str, default: Any = None) -> Any: + raw = await self._loop_state().client.get(self._kv(key)) + return default if raw is None else decode(raw) + + async def aset(self, key: str, value: Any, ttl: Optional[float] = None) -> None: + px = max(1, round(ttl * 1000)) if ttl is not None else None + await self._loop_state().client.set(self._kv(key), encode(value), px=px) + + async def adelete(self, key: str) -> None: + await self._loop_state().client.delete(self._kv(key)) + # --- pub/sub ----------------------------------------------------------- def publish(self, topic: str, message: Any) -> None: self.start() @@ -131,10 +401,20 @@ def publish(self, topic: str, message: Any) -> None: args=[encode(message), self._buffer_size], ) + async def apublish(self, topic: str, message: Any) -> None: + await self._loop_state().publish_script( + keys=[self._seq(topic), self._stream(topic)], + args=[encode(message), self._buffer_size], + ) + def _head(self, topic: str) -> int: raw = self._redis.get(self._seq(topic)) return int(raw) if raw is not None else 0 + async def _ahead(self, topic: str) -> int: + raw = await self._loop_state().client.get(self._seq(topic)) + return int(raw) if raw is not None else 0 + def _poll(self, topic: str, after_seq: int, timeout: float) -> PollResult: stream = self._stream(topic) # Cursor past the head: it was minted before the stream was reset (the @@ -159,4 +439,16 @@ def _poll(self, topic: str, after_seq: int, timeout: float) -> PollResult: def subscribe(self, topic: str, replay_from: Optional[int] = None) -> Subscription: start = replay_from if replay_from is not None else self._head(topic) - return PollingSubscription(topic, start, self._poll, _POLL_TIMEOUT) + return RedisSubscription(self, topic, start) + + +class _LoopState: # pylint: disable=too-few-public-methods + def __init__(self, storage: RedisSharedStorage, client: Any): + self.client = client + self.publish_script = client.register_script(_PUBLISH_LUA) + self.reader = _StreamReader(storage, client) + + async def aclose(self) -> None: + # Also the pool a client built around a passed one's options, which + # aclose() leaves open by default. + await self.client.aclose(close_connection_pool=True) diff --git a/dash/_stream_hub.py b/dash/_stream_hub.py index 3f6a1878fb..ee4fd01cb5 100644 --- a/dash/_stream_hub.py +++ b/dash/_stream_hub.py @@ -177,10 +177,11 @@ async def apublish_frame( class Downlink: """One long-lived downlink: its topic subscription plus its lifecycle record. - ``close`` is thread-safe and idempotent. Closing only rewrites the connection - record if it still carries this downlink's token, so a downlink that was - replaced by a reconnect cannot mark the new one closed when it finally winds - down. + The relay (``envelopes`` / ``aenvelopes``) marks the record open when it + starts and closed when it ends. ``close`` is thread-safe and idempotent. + Closing only rewrites the connection record if it still carries this + downlink's token, so a downlink that was replaced by a reconnect cannot mark + the new one closed when it finally winds down. """ def __init__( @@ -192,45 +193,56 @@ def __init__( self.storage = storage self.connection_id = connection_id self.subscription = storage.subscribe(stream_topic(connection_id), replay_from) + self._key = connection_key(connection_id) self._token = secrets.token_hex(8) self._closed = False - storage.set( - connection_key(connection_id), - { - "mode": "stream", - "open": True, - "at": time.time(), - "token": self._token, - }, - ) with _registry_lock: _active_subscriptions.add(self.subscription) - def close(self) -> None: + def _record(self, is_open: bool) -> Any: + return { + "mode": "stream", + "open": is_open, + "at": time.time(), + "token": self._token, + } + + def _is_mine(self, record: Any) -> bool: + return isinstance(record, dict) and record.get("token") == self._token + + def _release(self) -> bool: + """End the subscription; False if this downlink was already closed.""" if self._closed: - return + return False self._closed = True with _registry_lock: _active_subscriptions.discard(self.subscription) self.subscription.close() - key = connection_key(self.connection_id) + return True + + def open(self) -> None: + self.storage.set(self._key, self._record(True)) + + async def aopen(self) -> None: + await self.storage.aset(self._key, self._record(True)) + + def close(self) -> None: + if not self._release(): + return with contextlib.suppress(Exception): - current = self.storage.get(key) - if isinstance(current, dict) and current.get("token") == self._token: - self.storage.set( - key, - { - "mode": "stream", - "open": False, - "at": time.time(), - "token": self._token, - }, - ) + if self._is_mine(self.storage.get(self._key)): + self.storage.set(self._key, self._record(False)) + + async def _arecord_closed(self) -> None: + with contextlib.suppress(Exception): + if self._is_mine(await self.storage.aget(self._key)): + await self.storage.aset(self._key, self._record(False)) def envelopes(self) -> Iterator[Any]: """Sync relay: yield envelopes until the subscription ends or is closed. A lost buffer surfaces as a single reset envelope.""" try: + self.open() for seq, message in self.subscription.iter_with_seq(): yield {**message, "seq": seq} except SharedStorageGap: @@ -241,12 +253,25 @@ def envelopes(self) -> Iterator[Any]: async def aenvelopes(self) -> AsyncIterator[Any]: """Async counterpart of :meth:`envelopes` for ASGI backends.""" try: + await self.aopen() async for seq, message in self.subscription.aiter_with_seq(): yield {**message, "seq": seq} except SharedStorageGap: yield dict(RESET_ENVELOPE) finally: - self.close() + if self._release(): + # Its own task: when the client disconnects, the server cancels + # this response, and an await here would be cancelled with it, + # leaving the record open so the pumps never learn the browser + # left. + with contextlib.suppress(RuntimeError): # the loop is gone + task = asyncio.ensure_future(self._arecord_closed()) + _closing_downlinks.add(task) + task.add_done_callback(_closing_downlinks.discard) + + +# Downlinks writing their closed record (keeps the tasks referenced). +_closing_downlinks: "set[asyncio.Task]" = set() def downlink_gone( diff --git a/tests/shared_storage/test_diskcache_backend.py b/tests/shared_storage/test_diskcache_backend.py index 9b833c5a6b..04fb65f2d0 100644 --- a/tests/shared_storage/test_diskcache_backend.py +++ b/tests/shared_storage/test_diskcache_backend.py @@ -4,9 +4,11 @@ and cross-process visibility over a shared cache directory. Requires the diskcache extra, which the shared-storage CI job installs. """ +import asyncio import multiprocessing as mp import threading import time +from concurrent.futures import ThreadPoolExecutor import pytest @@ -144,3 +146,37 @@ def test_pubsub_across_processes(tmp_path): sub.close() store.close() assert received == ["m0", "m1", "m2"] + + +def test_async_subscriptions_hold_no_executor_threads(store): + """Waiting async subscriptions must leave the loop's executor free for the + other store calls on an ASGI worker.""" + + async def first(sub): + async for pair in sub.aiter_with_seq(): + return pair + return None + + async def scenario(): + executor = ThreadPoolExecutor(max_workers=1) + asyncio.get_running_loop().set_default_executor(executor) + subs = [store.subscribe(f"t{i}", replay_from=0) for i in range(20)] + try: + await asyncio.wait_for(check(subs), timeout=5) + finally: + for sub in subs: + sub.close() + executor.shutdown(wait=False, cancel_futures=True) + + async def check(subs): + readers = [asyncio.ensure_future(first(sub)) for sub in subs] + await asyncio.sleep(0.3) + + started = time.monotonic() + for i in range(20): + await store.apublish(f"t{i}", f"m{i}") + got = await asyncio.wait_for(asyncio.gather(*readers), timeout=5) + assert got == [(1, f"m{i}") for i in range(20)] + assert time.monotonic() - started < 2.0 + + asyncio.run(scenario()) diff --git a/tests/shared_storage/test_redis_backend.py b/tests/shared_storage/test_redis_backend.py index f86a650c02..c5933a8a41 100644 --- a/tests/shared_storage/test_redis_backend.py +++ b/tests/shared_storage/test_redis_backend.py @@ -4,10 +4,12 @@ tests skip when none is available. The CI shared-storage job provides one. Each test uses a unique key prefix so a shared Redis stays isolated. """ +import asyncio import os import threading import time import uuid +from concurrent.futures import ThreadPoolExecutor import pytest @@ -142,3 +144,162 @@ def test_two_instances_share_state(store): sub.close() other.close() assert received == ["m0", "m1", "m2"] + + +# --- asyncio path ------------------------------------------------------------- + + +async def _first(sub): + async for pair in sub.aiter_with_seq(): + return pair + return None + + +def test_async_subscriptions_hold_no_executor_threads(store): + """Open async subscriptions must not occupy the loop's executor: on an ASGI + worker every other store call (the pumps' publishes, the downlink records) + would queue behind them.""" + + async def scenario(): + executor = ThreadPoolExecutor(max_workers=1) + asyncio.get_running_loop().set_default_executor(executor) + subs = [store.subscribe(f"t{i}", replay_from=0) for i in range(20)] + try: + await asyncio.wait_for(check(subs), timeout=5) + finally: + for sub in subs: + sub.close() + executor.shutdown(wait=False, cancel_futures=True) + + async def check(subs): + readers = [asyncio.ensure_future(_first(sub)) for sub in subs] + await asyncio.sleep(0.3) # every reader is now waiting + + started = time.monotonic() + await store.aset("k", 1) + assert await store.aget("k") == 1 + for i in range(20): + await store.apublish(f"t{i}", f"m{i}") + got = await asyncio.wait_for(asyncio.gather(*readers), timeout=2) + assert got == [(1, f"m{i}") for i in range(20)] + assert time.monotonic() - started < 1.0 + await store.adelete("k") + assert await store.aget("k", "gone") == "gone" + + asyncio.run(scenario()) + + +def test_async_replay_and_live_messages_in_order(store): + store.publish("t", "m1") + store.publish("t", "m2") + + async def scenario(): + out = [] + sub = store.subscribe("t", replay_from=1) + async for pair in sub.aiter_with_seq(): + out.append(pair) + if len(out) == 1: + await store.apublish("t", "m3") + if len(out) == 2: + await store.apublish("t", "m4") + if len(out) == 3: + break + return out + + assert asyncio.run(scenario()) == [(2, "m2"), (3, "m3"), (4, "m4")] + + +def test_async_subscribers_with_different_cursors_share_a_topic(store): + for i in range(3): + store.publish("t", f"m{i}") + + async def scenario(): + late = store.subscribe("t") # at the head: seq 3 + early = store.subscribe("t", replay_from=0) + late_first = asyncio.ensure_future(_first(late)) + assert await _first(early) == (1, "m0") + await asyncio.sleep(0.1) + await store.apublish("t", "m3") + return await asyncio.wait_for(late_first, timeout=2) + + assert asyncio.run(scenario()) == (4, "m3") + + +def test_async_gap_when_buffer_overruns(): + prefix = f"dash:sstest:{uuid.uuid4().hex[:12]}" + store = RedisSharedStorage(url=REDIS_URL, key_prefix=prefix, buffer_size=2) + for i in range(5): + store.publish("t", f"m{i}") # stream trimmed to seqs 4,5 + + async def scenario(replay_from): + sub = store.subscribe("t", replay_from=replay_from) + with pytest.raises(SharedStorageGap): + await _first(sub) + + asyncio.run(scenario(1)) # wants seq 2, trimmed away + asyncio.run(scenario(9)) # cursor past the head: the store was reset + store.close() + + +def test_async_subscription_close_from_another_thread_ends_iteration(store): + async def scenario(): + sub = store.subscribe("t") + reader = asyncio.ensure_future(_first(sub)) + await asyncio.sleep(0.2) + threading.Thread(target=sub.close).start() + return await asyncio.wait_for(reader, timeout=1) + + assert asyncio.run(scenario()) is None + + +def test_async_path_with_a_passed_client(): + client = redis.Redis.from_url(REDIS_URL) + store = RedisSharedStorage( + client=client, key_prefix=f"dash:sstest:{uuid.uuid4().hex[:12]}" + ) + + async def scenario(): + await store.aset("a", {"x": 1}) + reader = asyncio.ensure_future(_first(store.subscribe("t"))) + await asyncio.sleep(0.1) + await store.apublish("t", "m") + return await store.aget("a"), await asyncio.wait_for(reader, timeout=2) + + assert asyncio.run(scenario()) == ({"x": 1}, (1, "m")) + assert store.get("a") == {"x": 1} + store.close() + client.close() + + +def test_async_subscriber_joining_behind_a_read_in_flight_gets_no_false_gap(store): + """A reconnecting downlink replays from behind while another subscription to + the same topic is parked in the reader's XREAD: the read in flight started + past the newcomer's cursor, which must not count as lost messages.""" + for i in range(5): + store.publish("t", f"m{i}") + + async def head(_topic): + return 10**9 + + async def scenario(): + parked = store.subscribe("t") # at the head: seq 5 + parked_first = asyncio.ensure_future(_first(parked)) + await asyncio.sleep(0.3) # the reader is blocked reading from seq 5 + store._ahead = head # joins without an await before the reader sees it + store.publish("t", "m5") + out = [] + async for pair in store.subscribe("t", replay_from=1).aiter_with_seq(): + out.append(pair) + if len(out) == 5: + break + return out, await asyncio.wait_for(parked_first, timeout=2) + + out, parked = asyncio.run(scenario()) + assert out == [(2, "m1"), (3, "m2"), (4, "m3"), (5, "m4"), (6, "m5")] + assert parked == (6, "m5") + + +def test_async_clients_of_closed_loops_are_dropped(store): + for _ in range(3): + asyncio.run(store.aget("k")) + assert len(store._loops) == 1 diff --git a/tests/streaming/test_stream_asgi_redis.py b/tests/streaming/test_stream_asgi_redis.py new file mode 100644 index 0000000000..cc37e765dc --- /dev/null +++ b/tests/streaming/test_stream_asgi_redis.py @@ -0,0 +1,168 @@ +"""Streaming on multi-worker uvicorn with Redis shared storage. + +Every browser holds an open downlink. Those must cost a worker nothing but a +waiting task: if each one parks a blocking read in the loop's executor, a dozen +of them starve the worker, and every publish and store call behind them waits +out the read's block timeout. Frames then arrive seconds late, or not at all. +""" +import os +import re +import signal +import socket +import subprocess +import sys +import textwrap +import time +import uuid +from concurrent.futures import ThreadPoolExecutor + +import pytest +import requests + +from dash.testing.wait import until + +pytest.importorskip("uvicorn") +redis = pytest.importorskip("redis") + +REDIS_URL = os.environ.get("REDIS_URL", "redis://localhost:6379") + + +def _redis_available(): + try: + client = redis.Redis.from_url(REDIS_URL) + client.ping() + client.close() + return True + except Exception: # pylint: disable=broad-except + return False + + +pytestmark = [ + pytest.mark.skipif( + sys.platform == "win32", reason="multi-worker uvicorn is POSIX only here" + ), + pytest.mark.skipif( + not _redis_available(), reason="no Redis reachable at REDIS_URL" + ), +] +# Per worker, comfortably more than the default executor's threads (at most 32). +IDLE_DOWNLINKS = 80 + +APP = textwrap.dedent( + """ + import asyncio + import os + from dash import Dash, Input, Output, RedisSharedStorage, html + + app = Dash( + __name__, + backend=os.environ["TEST_BACKEND"], + shared_storage=RedisSharedStorage( + url=os.environ["REDIS_URL"], key_prefix=os.environ["TEST_PREFIX"] + ), + ) + # Every worker must verify the others' signed end_id. + app.server.secret_key = "stream-test" + app.layout = html.Div([html.Button("go", id="btn", n_clicks=0), html.Div(id="out")]) + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def stream(n): + if not n: + return + for i in range(20): + yield f"token {i} " + await asyncio.sleep(0.05) + + server = app.server + """ +) + + +def _free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _end_id(base): + page = requests.get(base, timeout=5).text + return re.search(r'"end_id":\s*"([^"]+)"', page).group(1) + + +def _open_downlink(url): + resp = requests.post( + url, json={"streamDownlink": {"from": 0}}, stream=True, timeout=10 + ) + resp.raise_for_status() + return resp + + +@pytest.mark.parametrize("backend", ["fastapi", "quart"]) +def test_star001_stream_promptly_beside_many_open_downlinks(dash_br, tmp_path, backend): + pytest.importorskip(backend) + (tmp_path / "asgi_app.py").write_text(APP) + port = _free_port() + base = f"http://127.0.0.1:{port}" + log = (tmp_path / "uvicorn.log").open("wb") + proc = subprocess.Popen( # pylint: disable=consider-using-with + [ + sys.executable, + "-m", + "uvicorn", + "asgi_app:server", + "--host", + "127.0.0.1", + "--port", + str(port), + "--workers", + "2", + ], + cwd=tmp_path, + env={ + **os.environ, + "TEST_BACKEND": backend, + "REDIS_URL": REDIS_URL, + "TEST_PREFIX": f"dash:asgitest:{uuid.uuid4().hex[:12]}", + }, + stdout=subprocess.DEVNULL, + stderr=log, + start_new_session=True, + ) + downlinks = [] + try: + + def up(): + try: + return requests.get(base, timeout=1).ok + except requests.RequestException: + return False + + until(up, timeout=30, poll=0.2, msg="uvicorn never came up") + + # Other browsers, quiet but connected, spread over both workers. + url = f"{base}/_dash-update-component?endId={_end_id(base)}" + with ThreadPoolExecutor(max_workers=16) as pool: + downlinks = list(pool.map(_open_downlink, [url] * IDLE_DOWNLINKS)) + time.sleep(0.5) + + dash_br.server_url = base + dash_br.find_element("#btn").click() + started = time.monotonic() + dash_br.wait_for_contains_text("#out", "token 19", timeout=15) + # 20 frames 50 ms apart: about 1 s when each frame is delivered as it + # is produced. A starved worker takes 5 s or more for each frame. + assert time.monotonic() - started < 4 + assert dash_br.get_logs() == [] + except Exception: + print((tmp_path / "uvicorn.log").read_text(errors="replace")[-4000:]) + raise + finally: + for resp in downlinks: + resp.close() + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + log.close() diff --git a/tests/streaming/test_stream_transport.py b/tests/streaming/test_stream_transport.py index 458cb93fda..5cc7c2a788 100644 --- a/tests/streaming/test_stream_transport.py +++ b/tests/streaming/test_stream_transport.py @@ -6,6 +6,7 @@ dispatch on all three backends (Flask WSGI, Quart + FastAPI ASGI). """ import asyncio +import contextlib import json import threading import time @@ -217,6 +218,7 @@ def test_downlink_gone_semantics(monkeypatch): assert hub.downlink_gone(storage, cid, started - 2.0, grace=1.0) downlink = hub.Downlink(storage, cid) + downlink.open() assert not hub.downlink_gone(storage, cid, started - 100, grace=0.0) downlink.close() # Closed long before this pump started: the client may be about to open a @@ -233,13 +235,48 @@ def test_replaced_downlink_does_not_mark_the_new_one_closed(): storage = _storage("replace") cid = "c2" old = hub.Downlink(storage, cid) + old.open() new = hub.Downlink(storage, cid, replay_from=0) # the client reconnected + new.open() old.close() # the stale relay winds down late assert storage.get(hub.connection_key(cid))["open"] is True new.close() assert storage.get(hub.connection_key(cid))["open"] is False +def test_cancelled_async_downlink_still_records_closed(): + """A client disconnect cancels the ASGI response mid-relay; the closed record + must still be written, or pumps never learn the browser left.""" + from dash import _stream_hub as hub + + storage = _storage("acancel") + cid = "c3" + + async def scenario(): + downlink = hub.Downlink(storage, cid) + envelopes = downlink.aenvelopes() + + async def relay(): + async for _env in envelopes: + pass + + task = asyncio.ensure_future(relay()) + await asyncio.sleep(0.2) + assert (await storage.aget(hub.connection_key(cid)))["open"] is True + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + assert downlink.subscription not in hub._active_subscriptions + for _ in range(50): + if not (await storage.aget(hub.connection_key(cid)))["open"]: + return True + await asyncio.sleep(0.02) + return False + + assert asyncio.run(scenario()) + storage.close() + + def _cancellation_probe(): """An endless async frame generator that records its cancellation.""" state = {"cancelled": False, "frames": 0} @@ -307,6 +344,7 @@ def test_async_pump_cancels_callback_when_downlink_gone(monkeypatch): storage = _storage("apump") cid = "c5" downlink = hub.Downlink(storage, cid) + downlink.open() async def scenario(): state, frames = _cancellation_probe() @@ -341,6 +379,7 @@ def test_stream_cancel_stops_sync_pump_while_downlink_stays_open(monkeypatch): storage = _storage("cancel-pump") cid = "c6" downlink = hub.Downlink(storage, cid) # stays open: other tabs still stream + downlink.open() state, frames = _cancellation_probe() marker = StreamedCallbackResponse(frames, is_async=True) pump = hub.pump_to_storage(storage, cid, "r1", marker) @@ -527,7 +566,7 @@ def test_wsgi_pumps_share_one_loop_thread(): from dash._streaming import StreamedCallbackResponse storage = _storage("pumploop") - hub.Downlink(storage, "c10") # a browser is present + hub.Downlink(storage, "c10").open() # a browser is present pumps = [] for i in range(25): _state, frames = _cancellation_probe() @@ -556,7 +595,7 @@ def test_wsgi_pump_loop_survives_an_exception_raised_into_it(): from dash._streaming import StreamedCallbackResponse storage = _storage("pumprestart") - hub.Downlink(storage, "c11") + hub.Downlink(storage, "c11").open() hub._shared_pump_loop() # pylint: disable=protected-access thread = hub._pump_thread # pylint: disable=protected-access _state, frames = _cancellation_probe()