diff --git a/.ai/ARCHITECTURE.md b/.ai/ARCHITECTURE.md index 4a8a379d18..d4bd2ace56 100644 --- a/.ai/ARCHITECTURE.md +++ b/.ai/ARCHITECTURE.md @@ -1284,6 +1284,207 @@ async def validate_message(websocket, message): - `@plotly/dash-websocket-worker/src/worker.ts` - SharedWorker entry point - `dash/backends/_fastapi.py` - Server-side WebSocket handler +## Streaming Callbacks + +A callback defined as a generator function (or async generator function) +streams: its yields are pushed to the browser as they are produced — for LLM +token streaming, progress feeds, and long computations. There is no opt-in +keyword; `dash._callback.register_callback` infers it from the decorated +function (`inspect.isgeneratorfunction` / `isasyncgenfunction`) and registers +the streaming wrapper instead of the regular one. + +```python +import asyncio +from dash import callback, Output, Input, Patch + +@callback( + Output('log', 'children'), + Input('btn', 'n_clicks'), + prevent_initial_call=True, +) +async def run(n): + yield 'Starting...' # replaces children immediately + async for token in llm(): + p = Patch() + p += token + yield p # appends to children (incremental) + yield 'Done' # last yield = final value +``` + +### Semantics + +- Each yield has the same shape as a regular return value (one value per + `Output`) and **replaces** the outputs. Yield `dash.Patch` objects for + incremental updates. +- `no_update` works per-output within a yield; a yield where nothing updates + produces no frame. Raising `PreventUpdate` mid-stream ends the stream + cleanly. +- `set_props()` between yields is folded into the next frame's `sideUpdate` + (HTTP) or streams immediately (WebSocket transport). +- Intermediate frames are applied through the same renderer path as + `set_props`, so dependent callbacks fire per frame and loading states stay + on until the stream completes (`Updating...` title for the whole stream). +- `on_error` applies per-stream: its return value becomes a final frame. + Without it, an exception mid-stream sends an error frame shown in devtools; + frames already applied stay applied. +- The callback must be an `async def` generator on every backend; a + synchronous generator is rejected at registration, since it would occupy + a server worker (or WS executor thread) for the whole stream. +- HTTP streams emit a blank keepalive line every `stream_keepalive_interval` + ms (`Dash(stream_keepalive_interval=15000)`) that the callback spends + between yields, so proxy idle timeouts (nginx `proxy_read_timeout`, 60s by + default) don't close a stream mid-thought; `None` disables it. The + renderer skips blank lines. +- Incompatible with `background=True`, `mcp_enabled` and `api_endpoint` + (validated at registration, when the function is inspected). Clientside + callbacks cannot stream at all. +- `callback_map[callback_id]['stream']` records the inferred flag server-side; + it is not part of the callback spec sent to the client, which detects a + stream from the response instead (NDJSON content type / `stream` frames). + +### Transport & frame protocol + +Transport follows the callback's normal transport selection: if the callback +runs over the WebSocket callback transport (`websocket=True` or +`websocket_callbacks=True`), frames ride the open connection as +`callback_response` messages with `stream: true`; the terminal message is +`{status: 'ok', stream: true, done: true}`. Otherwise the HTTP POST response +streams NDJSON (`application/x-ndjson`), one frame per line: + +``` +{"multi": true, "response": {"": {"": }}, "sideUpdate": {...}?} +{"done": true} <- terminal frame +{"done": true, "error": {"message": "..."}} <- error terminal frame +``` + +The renderer applies each frame on arrival (via the `sideUpdate` path, so +`Patch` applies exactly once) and resolves the callback's execution promise +with an empty result on the terminal frame. + +### Multiplexed downlink and the stream SharedWorker + +When the app has a shared-storage backend (the default `LocalSharedStorage`), +HTTP streams do not each hold their own response. The callback's POST carries +`streamConnection: {requestId}` and returns a fast ack; a *pump* +(`dash/_stream_hub.py`) drives the generator as an asyncio task and publishes +each frame, tagged with the request id, to the connection's shared-storage +topic. The browser's *downlink* (`streamDownlink: {from}`) reads that topic +and the client routes `{rid, frame, seq}` envelopes back by request id. +Callback and downlink can be on different workers -- the store is the broker +-- and a downlink always resumes from its last `seq`, replaying from the +store's buffer. If the buffer no longer covers the cursor (restart, owner +re-election) the server sends `{reset: true}` and the client fails its +in-flight streams instead of silently skipping frames. + +The connection id is never chosen by the client: every stream request rides on +`?endId=`, the server-signed per-page-load token, and the backend derives the +id from it (`get_stream_connection_id`), answering 403 when it is missing or +forged -- otherwise a client could read or inject into another page's topic. +Across worker processes every worker must resolve the same signing secret +(`secret_key`). + +The downlink is hosted in a SharedWorker (`dash-stream-worker.js`, served like +the WebSocket worker; `config.stream.worker_url`) so **one connection per +browser** serves every tab: browsers cap HTTP/1.1 connections per host at +about six, and a downlink per tab stalls the sixth tab. The worker pins the +`endId` of the tab that opened the downlink while streams are in flight (all +tabs' frames flow through that one topic). The page talks to the worker +through `SharedStreamClient` (`utils/streamClient.ts`); the worker runs the +real `StreamClient` behind `attachStreamWorkerHost` +(`utils/streamWorkerHost.ts`). Without SharedWorker support the page falls +back to a downlink of its own. + +**Two downlink modes** (`config.stream.mode`, from `backend.downlink_mode`): + +- `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. +- `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). + Instead each downlink request returns the frames queued since the cursor + and ends at once (`poll_downlink`, `Subscription.poll(0)`), taking a thread + for milliseconds. The worker re-polls every `stream_poll_interval` ms + (default 100) while frames flow, backs off to five times that after two + empty polls (bounding a slow stream's frame latency so frames don't bunch + into one poll), and polls immediately when a new stream starts. Pumps are + tasks on one event-loop thread per WSGI process (`pump_to_storage`), using + the store's loop-native `aget`/`apublish`, not a thread per stream. + +**Lifecycle.** Each downlink records its state under the connection's key in +shared storage: open/closed for a long-lived downlink, a heartbeat (at most +once a second per worker) for a polling one. Every pump checks it about every +2s and cancels its callback at its current `await` once the browser is gone: +a downlink closed for `DOWNLINK_GRACE` (10s), or no poll for `POLL_GRACE` +(30s -- wide, because an overloaded pool delays polls and overload must cost +latency, never the stream). A tab closing while other tabs keep the shared +downlink sends `streamCancel: {requestId}` per stream instead; +the same pump check picks up the per-request key. A pump that stops publishes +a terminal `{"done": true}` so a late-reconnecting client resolves. + +**Shutdown.** ASGI servers drain in-flight responses before stopping and a +long-lived downlink never ends by itself, so `_stream_hub` installs a +SIGINT/SIGTERM handler (`install_stream_shutdown_handler`, at import and again +from backend startup since uvicorn replaces handlers) that runs +`shutdown_active_streams` -- sets the shutdown flag, cancels every pump on its +own loop, closes every open downlink subscription -- then chains to the +server's own handler. WSGI pumps also stop from an `atexit` hook. The pump +loop thread shrugs off exceptions raised into it (dash.testing's runner stops +every thread an app started) and is recreated if it ever dies. + +**Scale** (this dev box, 8 cores shared with the load clients; a streaming +callback per browser yielding every 0.5s; delivery = server yield to client +receipt): + +| server | browsers | frame delay p50 / p95 | +|---|---|---| +| uvicorn, 1 worker (FastAPI) | 1000 | 3 ms / 22 ms | +| uvicorn, 4 workers | 1000 | 1 ms / 3 ms | +| gunicorn `-w 4 --threads 8` (Flask, poll) | 300 | 105 ms / 200 ms | +| gunicorn `-w 8 --threads 8` | 1000 | 180 ms / 3.6 s (CPU-bound) | +| gunicorn `-w 1` (sync worker) | 50 | 50 ms / 100 ms | + +Flask works and degrades gracefully -- the cost is a poll per browser per +interval, so plan roughly one gunicorn worker per 150 concurrently streaming +browsers -- but for thousands of concurrent streams the ASGI backends are +the right tool: constant latency and a fraction of the CPU. + +### Caveats + +- Streaming is inferred from the decorated function, so another decorator + between `@callback` and the generator hides it: if that decorator returns a + plain function, Dash registers a regular callback and the returned generator + object fails to serialize (`InvalidCallbackReturnValue: type generator`). +- Long streams should check `ctx.websocket.is_shutdown` (WS transport) in + loops; on HTTP, client disconnect raises `GeneratorExit` into the user + generator at its current `yield`. +- Proxies and compression middleware (nginx buffering, flask-compress/gzip, + Jupyter proxies) can buffer NDJSON and defeat streaming. Dash sets + `X-Accel-Buffering: no`, but middleware configuration may still be needed. +- Streamed frames bypass persistence (`prunePersistence`/`applyPersistence`). +- Wrapping a streamed output in `dcc.Loading` hides it for the entire stream + (loading stays on by design). +- Flask + async generator requires `dash[async]`; frames are bridged from a + private event-loop thread. +- `flask.request` inside a streamed callback body only works on the pure-WSGI + Flask path (no `dash[async]`/`use_async`); under async dispatch the request + context cannot be carried into the stream. Use Dash's `ctx` (cookies, + headers, args are captured at dispatch) instead. + +### Key Files + +- `dash/_callback.py` - `add_context_stream`/`async_add_context_stream` wrappers, frame builders +- `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/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 +- `dash/dash-renderer/src/utils/workerClient.ts` - stream-aware `callback_response` handling +- `dash/dash-renderer/src/utils/streamClient.ts` - `StreamClient` (downlink + uplinks), `SharedStreamClient` (page side of the worker), `getStreamClient` +- `dash/dash-renderer/src/utils/streamWorkerHost.ts`, `src/workers/streamWorker.ts` - the stream SharedWorker + ## Security ### XSS Protection diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 1258bbc266..12b3bd6a56 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -20,6 +20,7 @@ jobs: dcc_paths_changed: ${{ steps.filter.outputs.dcc_related_paths }} html_paths_changed: ${{ steps.filter.outputs.html_related_paths }} websocket_changed: ${{ steps.filter.outputs.websocket_paths }} + streaming_changed: ${{ steps.filter.outputs.streaming_paths }} steps: - name: Checkout repository uses: actions/checkout@v4 @@ -79,6 +80,14 @@ jobs: - '@dash-websocket-worker/**' - 'dash/dash-renderer/src/**' - 'tests/websocket/**' + streaming_paths: + - *shared_paths + - 'dash/_callback.py' + - 'dash/_streaming.py' + - 'dash/_callback_context.py' + - 'dash/backends/**' + - 'dash/dash-renderer/src/**' + - 'tests/streaming/**' lint-unit: name: Lint & Unit Tests (Python ${{ matrix.python-version }}) @@ -696,6 +705,69 @@ jobs: touch __init__.py pytest --headless --nopercyfinalize tests/websocket -v -s + streaming-tests: + name: Streaming Callback Tests (Python ${{ matrix.python-version }}) + needs: [build, changes_filter] + if: | + (github.event_name == 'push' && (github.ref == 'refs/heads/master' || github.ref == 'refs/heads/dev')) || + needs.changes_filter.outputs.streaming_changed == 'true' + timeout-minutes: 30 + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.9", "3.12"] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Node.js + uses: actions/setup-node@v4 + with: + node-version: '24' + cache: 'npm' + + - name: Install Node.js dependencies + run: npm ci + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + cache-dependency-path: requirements/*.txt + + - name: Download built Dash packages + uses: actions/download-artifact@v4 + with: + name: dash-packages + path: packages/ + + - name: Install Dash packages + # Streaming callbacks are async generators; the async extra pulls in + # flask[async], which they require on the Flask backend. + run: | + python -m pip install --upgrade pip wheel + python -m pip install "setuptools<80.0.0" + find packages -name dash-*.whl -print -exec sh -c 'pip install "{}[async,ci,testing,dev,fastapi,quart]"' \; + + - name: Setup Chrome and ChromeDriver + uses: browser-actions/setup-chrome@v1 + with: + chrome-version: stable + + - name: Build/Setup test components + run: npm run setup-tests.py + + - name: Run streaming tests + run: | + mkdir streamtests + cp -r tests streamtests/tests + cd streamtests + touch __init__.py + pytest --headless --nopercyfinalize tests/streaming -v -s + test-main: name: Main Dash Tests (Python ${{ matrix.python-version }}, React ${{ matrix.react-version }}, Group ${{ matrix.test-group }}) needs: build diff --git a/CHANGELOG.md b/CHANGELOG.md index d5652f17c8..83824d0bd9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ This project adheres to [Semantic Versioning](https://semver.org/). - `DiskcacheSharedStorage`: on a `diskcache.Cache`, shared by every process on one host, for single-machine deployments (not for multi-pod ones with ephemeral disks). - `RedisSharedStorage`: on Redis, using Redis Streams for the ordered pub/sub: the backend for horizontally-scaled deployments behind a load balancer, e.g. apps scaled across pods. - `LocalSharedStorage` additionally accepts `mode=` to make its key/value store durable: `"memory"` (default, in-memory only), `"persist"` (write-through to disk on every change), or `"persist-reset"` (in-memory speed with a periodic flush every `flush_interval` seconds and on clean exit). Persistent modes recover on start and on owner re-election, so state survives a process restart or a crashed owner. Data is stored in a chunked, atomically-written msgpack store (a per-namespace folder under the user cache directory by default, overridable via `path=`). TTLs are preserved across restarts; pub/sub remains transient. +- [#3931](https://github.com/plotly/dash/pull/3931) Streaming callbacks: a callback defined as an `async def` generator streams its yields to the browser as they are produced (`dash.Patch` yields give incremental updates). All of a browser's streams share one connection hosted in a SharedWorker, so they don't count against the per-host connection limit; closing a tab cancels its streams. - [#3977](https://github.com/plotly/dash/pull/3977) Add partial WebSocket prop reads with `get_prop(..., path=...)`. Closes [#3975](https://github.com/plotly/dash/issues/3975). - [#3765](https://github.com/plotly/dash/pull/3765) Add opt-in partial pattern matching for callback `Input`, `Output`, and `State` dependencies via `partial_pattern=True`. Dictionary ID patterns can now match component IDs containing additional keys, and partial patterns can be combined with `ALL` and `MATCH` wildcards. Fixes [#3764](https://github.com/plotly/dash/issues/3764). - [#3646](https://github.com/plotly/dash/pull/3646) Experimental support for React 19. The default is still React 18.3.1; to use React 19 set the environment variable `REACT_VERSION=19.2.4` before running your app, or call `dash._dash_renderer._set_react_version("19.2.4")` inside the app. React 19 has no official UMD builds, so Dash serves the [`umd-react`](https://www.npmjs.com/package/umd-react) package, together with a compatibility shim loaded after react-dom and before any component package. The shim keeps component libraries built against React <=18 (e.g. dash-bootstrap-components, dash-mantine-components) working under React 19: it stubs the removed `ReactCurrentOwner` internals, redirects the legacy element `$$typeof` symbol so pre-bundled React 18 jsx-runtimes produce elements React 19 accepts (error #525), and exposes a global `react/jsx-runtime` (`window.ReactJSXRuntime`) that Dash's own component bundles externalize to. Component library authors adopting this convention should copy the defensive `jsxRuntimeExternal` webpack external from `components/dash-core-components/webpack.config.js` rather than a bare `'ReactJSXRuntime'` string: it falls back to a `React.createElement`-based runtime when the global is missing, so the same build also works on Dash versions older than this release. diff --git a/dash/_callback.py b/dash/_callback.py index 1a24bdf621..1b92907df1 100644 --- a/dash/_callback.py +++ b/dash/_callback.py @@ -1,6 +1,7 @@ import collections import hashlib import inspect +import logging import warnings from functools import wraps from typing import Callable, Optional, Any, List, Tuple, Union, Dict, TypeVar, cast @@ -22,6 +23,7 @@ MissingLongCallbackManagerError, BackgroundCallbackError, ImportedInsideCallbackError, + StreamCallbackError, ) from ._get_app import get_app from . import _callback_signing @@ -41,6 +43,7 @@ from .background_callback.managers import BaseBackgroundCallbackManager from ._callback_context import context_value +from ._streaming import StreamedCallbackResponse from .types import CallbackExecutionResponse from ._no_update import NoUpdate from . import _validate @@ -60,6 +63,8 @@ def _invoke_callback(func, *args, **kwargs): # used to mark the frame for the d return func(*args, **kwargs) # %% callback invoked %% +logger = logging.getLogger(__name__) + GLOBAL_CALLBACK_LIST: List[Any] = [] GLOBAL_CALLBACK_MAP: Dict[str, Any] = {} GLOBAL_INLINE_SCRIPTS: List[Any] = [] @@ -111,6 +116,15 @@ def callback( not to fire when its outputs are first added to the page. Defaults to `False` and unlike `app.callback` is not configurable at the app level. + Decorating an async generator function (`async def` with `yield`) registers + a streaming callback: each yielded value has the same shape as a regular + return value (one value per `Output`) and is pushed to the browser + immediately; yield `dash.Patch` objects for incremental updates. Streams + over the WebSocket callback transport when active, otherwise over the HTTP + response (NDJSON). Synchronous generators are not supported (they would + occupy a server worker for the whole stream). Streaming callbacks cannot be + combined with `background=True`, `mcp_enabled=True` or `api_endpoint`. + :Keyword Arguments: :param background: Mark the callback as a background callback to execute in a manager for @@ -269,6 +283,32 @@ def callback( ) +def _validate_stream_callback(callback_id, background, kwargs, is_sync_gen): + """Reject options a streaming (generator) callback cannot be combined with.""" + if is_sync_gen: + raise StreamCallbackError( + f"Streaming callback '{callback_id}' is a synchronous generator, " + "which is not supported: a sync generator occupies a server worker " + "for the whole stream. Define it with 'async def' so it streams on " + "the event loop instead." + ) + if background is not None: + raise BackgroundCallbackError( + f"Streaming callback '{callback_id}' cannot be combined with " + "background=True: background callbacks return a single result." + ) + if kwargs.get("mcp_enabled"): + raise StreamCallbackError( + f"Streaming callback '{callback_id}' cannot be combined with " + "mcp_enabled=True: MCP tools expect a single JSON result." + ) + if kwargs.get("api_endpoint"): + raise StreamCallbackError( + f"Streaming callback '{callback_id}' cannot be combined with " + "api_endpoint: API endpoints expect a single JSON result." + ) + + def validate_background_inputs(deps): for dep in deps: if dep.has_wildcard(): @@ -371,6 +411,9 @@ def insert_callback( "allow_dynamic_callbacks": dynamic_creator, "no_output": no_output, "websocket": websocket, + # Flipped to True by register_callback when the decorated function + # turns out to be a generator (streaming callback). + "stream": False, "mcp_enabled": mcp_enabled, "mcp_expose_docstring": mcp_expose_docstring, "compress_payload": compress_payload, @@ -428,6 +471,24 @@ def get_request_end_id(secret: bytes): return _callback_signing.unsign(secret, _callback_signing.END_SCOPE, token) +def get_stream_connection_id() -> "str | None": + """Return the verified streaming connection id for the request, or ``None``. + + The multiplexed streaming transport keys each page's downlink topic on this + id, so it must be unforgeable: a client that could name an arbitrary topic + could read another page's stream or inject frames into it. It is derived from + the server-signed ``end_id`` (the same per-page-load token background-callback + handles are bound to), never from anything the client picks. A missing or + forged token yields ``None``, and the backend refuses the request (403). + + ``end_id`` is signed with the server secret, so across worker processes every + worker must resolve the same secret: set a ``secret_key`` on the server, or + cross-worker stream requests will not verify. Single-process apps are fine + with no configuration. + """ + return get_request_end_id(_get_signing_secret()) + + def _get_signing_secret() -> bytes: return get_app()._get_signing_secret() # pylint: disable=protected-access @@ -802,9 +863,23 @@ def register_callback( compress_payload=_kwargs.get("compress_payload", False), compress_threshold=_kwargs.get("compress_threshold", 5_000), ) + # The client-facing spec insert_callback just appended. Streaming is flagged + # on it below (once the function is known to be a generator) so the renderer + # can keep streaming callbacks out of its concurrent-request budget. + client_spec = callback_list[-1] # pylint: disable=too-many-locals def wrap_func(func): + # A generator (or async generator) callback streams its yields; that is + # inferred from the function itself, there is no opt-in keyword. + is_gen_func = inspect.isgeneratorfunction(func) + is_async_gen_func = inspect.isasyncgenfunction(func) + is_stream = is_gen_func or is_async_gen_func + if is_stream: + _validate_stream_callback( + callback_id, background, _kwargs, is_sync_gen=is_gen_func + ) + if _kwargs.get("api_endpoint"): api_endpoint = _kwargs.get("api_endpoint") GLOBAL_API_PATHS[api_endpoint] = func @@ -963,7 +1038,143 @@ async def async_add_context(*args, **kwargs): return jsonResponse - if inspect.iscoroutinefunction(func): + def _build_stream_frame( + output_value, output_spec, callback_ctx, app, original_packages + ): + """Build one stream frame from a yielded value. + + Returns None when the yield produced no update (PreventUpdate or + all no_update with no set_props). + """ + response: CallbackExecutionResponse = {"multi": True} + try: + _prepare_response( + output_value, + output_spec, + multi, + response, + callback_ctx, + app, + original_packages, + None, + False, + has_output, + output, + callback_id, + allow_dynamic_callbacks, + ) + except PreventUpdate: + return None + finally: + # set_props between yields were folded into this frame's + # sideUpdate; don't resend them with later frames. + callback_ctx.updated_props.clear() + if not response.get("response") and not response.get("sideUpdate"): + return None + return response + + def _stream_error_frame(err): + logger.exception("Exception raised in streamed callback") + return {"done": True, "error": {"message": str(err) or repr(err)}} + + async def _astream_frames( + user_gen, error_handler, output_spec, callback_ctx, app, original_packages + ): + # Set the callback context var around each resumption of the user + # generator so dash.ctx/set_props resolve while its body runs. It is + # set per step rather than once for the whole stream because the + # keepalive drivers resume each __anext__ in a freshly copied + # context, where a token taken in an earlier step could not be reset. + async def _next(): + token = context_value.set(callback_ctx) + try: + return await user_gen.__anext__() + finally: + context_value.reset(token) + + def _handle_error(err): + # Run the on_error handler under the callback context too so + # dash.ctx/set_props resolve inside it, matching a step. + token = context_value.set(callback_ctx) + try: + return error_handler(err) + finally: + context_value.reset(token) + + try: + while True: + frame = None + try: + output_value = await _next() + frame = _build_stream_frame( + output_value, + output_spec, + callback_ctx, + app, + original_packages, + ) + except (StopAsyncIteration, PreventUpdate): + break + except Exception as err: # pylint: disable=broad-exception-caught + if error_handler: + output_value = _handle_error(err) + if output_value is not None: + frame = _build_stream_frame( + output_value, + output_spec, + callback_ctx, + app, + original_packages, + ) + if frame is not None: + yield frame + break + yield _stream_error_frame(err) + return + if frame is not None: + yield frame + yield {"done": True} + finally: + await user_gen.aclose() + + @wraps(func) + async def async_add_context_stream(*args, **kwargs): + """Handles streaming callbacks defined as async generators.""" + error_handler = on_error or kwargs.pop("app_on_error", None) + + ( + output_spec, + callback_ctx, + func_args, + func_kwargs, + app, + original_packages, + _, + ) = _initialize_context( + args, kwargs, inputs_state_indices, has_output, insert_output + ) + + user_gen = _invoke_callback(func, *func_args, **func_kwargs) + frames = _astream_frames( + user_gen, + error_handler, + output_spec, + callback_ctx, + app, + original_packages, + ) + return StreamedCallbackResponse(frames, is_async=True) + + if is_stream: + # Only async generators reach here; sync generators are rejected in + # _validate_stream_callback above. + callback_map[callback_id]["stream"] = True + # Server-inferred flag (not the removed stream=True keyword): the + # renderer reads it to exclude long-lived streams from its + # concurrent-request limit. + client_spec["stream"] = True + callback_map[callback_id]["callback"] = async_add_context_stream + elif inspect.iscoroutinefunction(func): callback_map[callback_id]["callback"] = async_add_context else: # A persistent, no-output callback streams via set_props and typically diff --git a/dash/_dash_renderer.py b/dash/_dash_renderer.py index b53b7867d4..0cac1b4b80 100644 --- a/dash/_dash_renderer.py +++ b/dash/_dash_renderer.py @@ -100,4 +100,9 @@ def _set_react_version(v_react, v_reactdom=None): "namespace": "dash", "dynamic": True, }, + { + "relative_package_path": "dash-renderer/build/dash-stream-worker.js", + "namespace": "dash", + "dynamic": True, + }, ] diff --git a/dash/_shared_storage/_engine.py b/dash/_shared_storage/_engine.py index 2df3313523..fd1fc8b4f0 100644 --- a/dash/_shared_storage/_engine.py +++ b/dash/_shared_storage/_engine.py @@ -9,8 +9,12 @@ silent hole. The engine is thread-safe and transport-agnostic; sockets live one layer up. +``poll`` blocks a thread; ``apoll`` parks an asyncio task on a future that +``publish`` resolves from whichever thread it runs on, so an ASGI server can +hold thousands of subscriptions without an executor thread each. """ +import asyncio import threading import time from collections import deque @@ -31,13 +35,23 @@ class PollResult(NamedTuple): gap: bool +_Waiter = Tuple[asyncio.AbstractEventLoop, "asyncio.Future[None]"] + + class _Topic: # pylint: disable=too-few-public-methods - __slots__ = ("seq", "buf", "cond") + __slots__ = ("seq", "buf", "cond", "waiters") def __init__(self, maxlen: int): self.seq = 0 self.buf: Deque[Tuple[int, Any]] = deque(maxlen=maxlen) self.cond = threading.Condition() + # asyncio tasks parked in apoll(), woken by the next publish/close. + self.waiters: List[_Waiter] = [] + + +def _wake(fut: "asyncio.Future[None]") -> None: + if not fut.done(): + fut.set_result(None) class StoreEngine: @@ -59,6 +73,10 @@ def attach_persistence(self, persistence: Any) -> None: populates the store first does not re-mark every restored key dirty).""" self._persistence = persistence + @property + def closed(self) -> bool: + return self._closed + # --- key/value --------------------------------------------------------- def get(self, key: str, default: Any = None) -> Any: with self._data_lock: @@ -139,7 +157,11 @@ def publish(self, topic: str, message: Any) -> int: t.seq += 1 t.buf.append((t.seq, message)) t.cond.notify_all() - return t.seq + waiters, t.waiters = t.waiters, [] + seq = t.seq + for loop, fut in waiters: + loop.call_soon_threadsafe(_wake, fut) + return seq def head_seq(self, topic: str) -> int: """Current highest sequence -- where a fresh subscription starts.""" @@ -147,6 +169,27 @@ def head_seq(self, topic: str) -> int: with t.cond: return t.seq + def _ready(self, t: _Topic, after_seq: int) -> Optional[PollResult]: + """Under ``t.cond``: the result available right now, or None to wait.""" + if self._closed: + return PollResult([], after_seq, False) + # The cursor points past every sequence this topic has ever produced. + # That can only happen when the cursor was minted by a previous + # incarnation of the topic (the owner was re-elected, or the server + # restarted with an empty store) -- treat it as a gap so the consumer + # resets instead of stalling until the fresh sequence climbs back past + # the stale cursor. + if after_seq > t.seq: + return PollResult([], after_seq, True) + # The next message we want is after_seq + 1; if the buffer's oldest is + # newer than that, it was evicted -> gap. + if t.buf and after_seq + 1 < t.buf[0][0]: + return PollResult([], after_seq, True) + fresh = [m for (s, m) in t.buf if s > after_seq] + if fresh: + return PollResult(fresh, t.buf[-1][0], False) + return None + def poll(self, topic: str, after_seq: int, timeout: float) -> PollResult: """Return messages with sequence > ``after_seq``, waiting up to ``timeout`` seconds for at least one. An empty result means the wait @@ -157,21 +200,41 @@ def poll(self, topic: str, after_seq: int, timeout: float) -> PollResult: deadline = time.monotonic() + timeout with t.cond: while True: - if self._closed: - return PollResult([], after_seq, False) - # The next message we want is after_seq + 1; if the buffer's - # oldest is newer than that, it was evicted -> gap. - if t.buf and after_seq + 1 < t.buf[0][0]: - return PollResult([], after_seq, True) - fresh = [m for (s, m) in t.buf if s > after_seq] - if fresh: - last = t.buf[-1][0] - return PollResult(fresh, last, False) + res = self._ready(t, after_seq) + if res is not None: + return res remaining = deadline - time.monotonic() if remaining <= 0: return PollResult([], after_seq, False) t.cond.wait(remaining) + async def apoll(self, topic: str, after_seq: int, timeout: float) -> PollResult: + """:meth:`poll` for asyncio: parks the task on a future instead of + blocking a thread; ``publish`` (from any thread) or ``close`` wakes it.""" + t = self._topic(topic) + loop = asyncio.get_running_loop() + deadline = time.monotonic() + timeout + while True: + with t.cond: + res = self._ready(t, after_seq) + if res is not None: + return res + remaining = deadline - time.monotonic() + if remaining <= 0: + return PollResult([], after_seq, False) + fut: "asyncio.Future[None]" = loop.create_future() + waiter = (loop, fut) + t.waiters.append(waiter) + try: + await asyncio.wait_for(asyncio.shield(fut), remaining) + except asyncio.TimeoutError: + return PollResult([], after_seq, False) + finally: + if not fut.done(): + with t.cond: + if waiter in t.waiters: + t.waiters.remove(waiter) + def close(self) -> None: self._closed = True if self._persistence is not None: @@ -182,3 +245,6 @@ def close(self) -> None: for t in topics: with t.cond: t.cond.notify_all() + waiters, t.waiters = t.waiters, [] + for loop, fut in waiters: + loop.call_soon_threadsafe(_wake, fut) diff --git a/dash/_shared_storage/_polling.py b/dash/_shared_storage/_polling.py index 146b9315b2..d7eb5f2ddb 100644 --- a/dash/_shared_storage/_polling.py +++ b/dash/_shared_storage/_polling.py @@ -35,21 +35,37 @@ def close(self) -> None: def _gap(self) -> SharedStorageGap: return SharedStorageGap(f"replay buffer overran on topic {self._topic!r}") - def __iter__(self): + @staticmethod + def _with_seq(res: PollResult): + # Poll batches are contiguous, ending at last_seq, so each message's + # sequence follows from its position. + first = res.last_seq - len(res.messages) + 1 + for offset, message in enumerate(res.messages): + yield first + offset, message + + def poll(self, timeout: float = 0.0): + res = self._poll_fn(self._topic, self._cursor, timeout) + if res.gap: + raise self._gap() + pairs = list(self._with_seq(res)) + self._cursor = res.last_seq + return pairs + + def iter_with_seq(self): try: while not self._closed.is_set(): res = self._poll_fn(self._topic, self._cursor, self._poll_timeout) if res.gap: raise self._gap() - yield from res.messages + yield from self._with_seq(res) self._cursor = res.last_seq finally: self.close() - def __aiter__(self): - return self._aiter() + def aiter_with_seq(self): + return self._aiter_with_seq() - async def _aiter(self): + async def _aiter_with_seq(self): loop = asyncio.get_running_loop() try: while not self._closed.is_set(): @@ -66,8 +82,8 @@ async def _aiter(self): break if res.gap: raise self._gap() - for message in res.messages: - yield message + for pair in self._with_seq(res): + yield pair self._cursor = res.last_seq finally: self.close() diff --git a/dash/_shared_storage/_transport.py b/dash/_shared_storage/_transport.py index 78ddb4fb7d..fcfe5f3877 100644 --- a/dash/_shared_storage/_transport.py +++ b/dash/_shared_storage/_transport.py @@ -18,10 +18,11 @@ connection), which is what makes long-poll subscriptions cheap. """ +import asyncio import socket import struct import threading -from typing import Any, Optional +from typing import Any, Optional, Tuple from ._codec import decode, encode @@ -57,6 +58,22 @@ def recv_frame(sock: socket.socket) -> Any: return decode(body) +async def asend_frame(writer: asyncio.StreamWriter, obj: Any) -> None: + data = encode(obj) + writer.write(struct.pack("!I", len(data)) + data) + await writer.drain() + + +async def arecv_frame(reader: asyncio.StreamReader) -> Any: + try: + header = await reader.readexactly(4) + (length,) = struct.unpack("!I", header) + body = await reader.readexactly(length) + except asyncio.IncompleteReadError: + return EOF + return decode(body) + + class OwnerServer: """Serves one StoreEngine to client workers over ``listen_sock``.""" @@ -129,6 +146,23 @@ def close(self) -> None: self._engine.close() +async def aconnect_to_owner( + family: int, address, token: str, timeout: float = 5.0 +) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]: + """Async counterpart of :func:`connect_to_owner` (asyncio streams).""" + if family == socket.AF_INET: + host, port = address + opener = asyncio.open_connection(host, port) + else: + opener = asyncio.open_unix_connection(address) + reader, writer = await asyncio.wait_for(opener, timeout) + await asend_frame(writer, token) + if await arecv_frame(reader) != OK: + writer.close() + raise ConnectionError("shared-storage owner rejected the handshake") + return reader, writer + + def connect_to_owner(family: int, address, token: str, timeout: float = 5.0): """Open a client connection and complete the token handshake.""" sock = socket.socket(family, socket.SOCK_STREAM) diff --git a/dash/_shared_storage/base.py b/dash/_shared_storage/base.py index a2d2d8aad3..2f0ac37d19 100644 --- a/dash/_shared_storage/base.py +++ b/dash/_shared_storage/base.py @@ -18,7 +18,8 @@ """ import abc -from typing import Any, AsyncIterator, Iterator, Optional +import asyncio +from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple class SharedStorageError(Exception): @@ -42,20 +43,47 @@ class Subscription(abc.ABC): messages buffered since the subscription's cursor replay first. Iteration ends when the subscription is closed. Raises ``SharedStorageGap`` if the buffer overran while the consumer was behind. + + ``iter_with_seq`` / ``aiter_with_seq`` yield ``(sequence, message)`` pairs so + a consumer can record its position and resume a later subscription from it + (via ``replay_from``) -- how the streaming downlink survives a reconnect + without losing frames. The plain message iterators are built on these. + + ``poll`` is the one-shot primitive underneath: everything published since + the cursor, waiting at most ``timeout`` seconds for something to arrive. + With ``timeout=0`` it never blocks, which is what lets a WSGI request serve + a downlink without holding its worker thread. """ @abc.abstractmethod - def __iter__(self) -> Iterator[Any]: + def poll(self, timeout: float = 0.0) -> List[Tuple[int, Any]]: + """Return the ``(sequence, message)`` pairs published since the cursor + and advance the cursor past them, waiting up to ``timeout`` seconds for + at least one. Raises ``SharedStorageGap`` if the buffer overran.""" + + @abc.abstractmethod + def iter_with_seq(self) -> Iterator[Any]: ... @abc.abstractmethod - def __aiter__(self) -> AsyncIterator[Any]: + def aiter_with_seq(self) -> AsyncIterator[Any]: ... @abc.abstractmethod def close(self) -> None: ... + def __iter__(self) -> Iterator[Any]: + for _seq, message in self.iter_with_seq(): + yield message + + def __aiter__(self) -> AsyncIterator[Any]: + return self._messages() + + async def _messages(self) -> AsyncIterator[Any]: + async for _seq, message in self.aiter_with_seq(): + yield message + def __enter__(self) -> "Subscription": return self @@ -108,6 +136,30 @@ def delete(self, key: str) -> None: def publish(self, topic: str, message: Any) -> None: """Append ``message`` to ``topic``; delivered to every current subscriber.""" + # --- asyncio variants -------------------------------------------------- + # Code running on an event loop (ASGI request handlers, the streaming + # pumps) must not block the loop on a round trip to the store. Backends + # should override these with loop-native I/O; the defaults run the sync + # operation on the loop's default executor. + + async def aget(self, key: str, default: Any = None) -> Any: + return await asyncio.get_running_loop().run_in_executor( + None, self.get, key, default + ) + + async def aset(self, key: str, value: Any, ttl: Optional[float] = None) -> None: + await asyncio.get_running_loop().run_in_executor( + None, self.set, key, value, ttl + ) + + async def adelete(self, key: str) -> None: + await asyncio.get_running_loop().run_in_executor(None, self.delete, key) + + async def apublish(self, topic: str, message: Any) -> None: + await asyncio.get_running_loop().run_in_executor( + None, self.publish, topic, message + ) + @abc.abstractmethod def subscribe(self, topic: str, replay_from: Optional[int] = None) -> Subscription: """Subscribe to ``topic``. diff --git a/dash/_shared_storage/diskcache.py b/dash/_shared_storage/diskcache.py index f19ec4fa3e..8d909ab80e 100644 --- a/dash/_shared_storage/diskcache.py +++ b/dash/_shared_storage/diskcache.py @@ -113,6 +113,12 @@ def _poll(self, topic: str, after_seq: int, timeout: float) -> PollResult: deadline = time.monotonic() + timeout while True: head = self._head(topic) + # Cursor past the head: it was minted before the store was reset + # (the cache was cleared, or the counter evicted under the cache's + # size limit). Gap so the consumer resets rather than blocking until + # the sequence climbs back past the cursor. + if after_seq > head: + return PollResult([], after_seq, True) if head > after_seq: floor = max(1, head - self._buffer_size + 1) if after_seq + 1 < floor: diff --git a/dash/_shared_storage/local.py b/dash/_shared_storage/local.py index 10496b10f4..4d838799b0 100644 --- a/dash/_shared_storage/local.py +++ b/dash/_shared_storage/local.py @@ -13,6 +13,7 @@ import asyncio import atexit +import contextlib import hashlib import json import os @@ -22,11 +23,20 @@ import tempfile import threading import time -from typing import Any, Optional +from typing import Any, Optional, Tuple from ._engine import DEFAULT_BUFFER, PollResult, StoreEngine from ._persistence import _Persistence, default_store_dir -from ._transport import EOF, OwnerServer, connect_to_owner, recv_frame, send_frame +from ._transport import ( + EOF, + OwnerServer, + aconnect_to_owner, + arecv_frame, + asend_frame, + connect_to_owner, + recv_frame, + send_frame, +) from .base import BaseSharedStorage, SharedStorageError, SharedStorageGap, Subscription _HAS_AF_UNIX = hasattr(socket, "AF_UNIX") @@ -197,6 +207,10 @@ def connect(self) -> socket.socket: assert self._family is not None and self._token is not None return connect_to_owner(self._family, self._address, self._token) + async def aconnect(self): + self.ensure() + return await aconnect_to_owner(self._family, self._address, self._token) + def on_owner_lost(self) -> None: """The owner became unreachable; re-elect (may promote us to owner).""" with self._lock: @@ -226,11 +240,16 @@ class _LocalSubscription(Subscription): cursor on reconnect so no buffered message is missed. """ - def __init__(self, coord: _Coordinator, topic: str, start_seq: int): + def __init__(self, coord: _Coordinator, topic: str, start_seq: int, call): self._coord = coord self._topic = topic self._cursor = start_seq + self._call = call # the storage's short request/response channel self._conn: Optional[socket.socket] = None + # Async (ASGI) client-role connection: reader/writer + its loop, so + # close() from another thread can shut it via the loop. + self._aconn: Optional[Tuple[asyncio.StreamReader, asyncio.StreamWriter]] = None + self._aloop: Optional[asyncio.AbstractEventLoop] = None self._closed = threading.Event() def close(self) -> None: @@ -241,15 +260,26 @@ def close(self) -> None: conn.close() except OSError: pass - - def _poll_once(self) -> PollResult: + aconn, self._aconn = self._aconn, None + if aconn is not None and self._aloop is not None: + _reader, writer = aconn + with contextlib.suppress(RuntimeError): + self._aloop.call_soon_threadsafe(writer.close) + + def _poll_once(self, timeout: Optional[float] = None) -> PollResult: + """One poll cycle. ``timeout`` overrides the role's long-poll default + (``0`` never blocks).""" if self._coord.is_owner(): engine = self._coord.engine - assert engine is not None - return engine.poll(self._topic, self._cursor, _OWNER_POLL_TIMEOUT) - return self._client_poll() + if engine is None or engine.closed: + self._closed.set() # owner gone -> end iteration, don't busy-loop + return PollResult([], self._cursor, False) + if timeout is None: + timeout = _OWNER_POLL_TIMEOUT + return engine.poll(self._topic, self._cursor, timeout) + return self._client_poll(_CLIENT_POLL_TIMEOUT if timeout is None else timeout) - def _client_poll(self) -> PollResult: + def _client_poll(self, timeout: float) -> PollResult: for attempt in range(4): if self._closed.is_set(): return PollResult([], self._cursor, False) @@ -272,7 +302,7 @@ def _client_poll(self) -> PollResult: try: send_frame( self._conn, - ["poll", self._topic, self._cursor, _CLIENT_POLL_TIMEOUT], + ["poll", self._topic, self._cursor, timeout], ) resp = recv_frame(self._conn) if resp is EOF: @@ -290,7 +320,30 @@ def _client_poll(self) -> PollResult: self._conn = None # reconnect on the next attempt return PollResult([], self._cursor, False) - def __iter__(self): + @staticmethod + def _with_seq(res: PollResult): + # Poll batches are contiguous, ending at last_seq, so each message's + # sequence follows from its position. + first = res.last_seq - len(res.messages) + 1 + for offset, message in enumerate(res.messages): + yield first + offset, message + + def poll(self, timeout: float = 0.0): + if timeout <= 0 and not self._coord.is_owner(): + # Non-blocking: a plain request on the worker's shared connection + # to the owner, like get/set -- no dedicated socket (and no thread + # on the owner) per poll. Blocking polls keep their own connection + # so they don't hold the shared one. + res = PollResult(*self._call(["poll", self._topic, self._cursor, 0.0])) + else: + res = self._poll_once(timeout) + if res.gap: + raise SharedStorageGap(f"replay buffer overran on topic {self._topic!r}") + pairs = list(self._with_seq(res)) + self._cursor = res.last_seq + return pairs + + def iter_with_seq(self): try: while not self._closed.is_set(): res = self._poll_once() @@ -298,35 +351,110 @@ def __iter__(self): raise SharedStorageGap( f"replay buffer overran on topic {self._topic!r}" ) - yield from res.messages + yield from self._with_seq(res) self._cursor = res.last_seq finally: self.close() - def __aiter__(self): - return self._aiter() + def aiter_with_seq(self): + return self._aiter_with_seq() - async def _aiter(self): - loop = asyncio.get_running_loop() + async def _apoll_once(self) -> PollResult: + """One long-poll cycle without leaving the event loop: the owner engine's + ``apoll`` (owner role) or an asyncio-streams connection to the owner + (client role) -- no executor thread per subscription either way.""" + if self._coord.is_owner(): + engine = self._coord.engine + if engine is None or engine.closed: + self._closed.set() + return PollResult([], self._cursor, False) + return await engine.apoll(self._topic, self._cursor, _OWNER_POLL_TIMEOUT) + return await self._aclient_poll(_CLIENT_POLL_TIMEOUT) + + async def _aclient_poll(self, timeout: float) -> PollResult: + self._aloop = asyncio.get_running_loop() + for attempt in range(4): + if self._closed.is_set(): + return PollResult([], self._cursor, False) + aconn = self._aconn + if aconn is None: + prev_token = self._coord.token + try: + aconn = await self._coord.aconnect() + except (OSError, asyncio.TimeoutError) as exc: + if attempt >= 2: + self._coord.on_owner_lost() + if self._coord.token != prev_token: + raise SharedStorageGap( + "shared-storage owner changed; buffered " + "messages were lost" + ) from exc + await asyncio.sleep(0.1 * (attempt + 1)) + continue + self._aconn = aconn + reader, writer = aconn + try: + await asend_frame(writer, ["poll", self._topic, self._cursor, timeout]) + resp = await arecv_frame(reader) + if resp is EOF: + raise ConnectionError("owner closed the connection") + status, val = resp + if status == "err": + raise SharedStorageError(val) + return PollResult(*val) + except (OSError, EOFError, ConnectionError): + with contextlib.suppress(Exception): + writer.close() + self._aconn = None # reconnect on the next attempt + return PollResult([], self._cursor, False) + + async def _aiter_with_seq(self): try: while not self._closed.is_set(): - try: - res = await loop.run_in_executor(None, self._poll_once) - except RuntimeError: - # The loop/executor is shutting down (client disconnected or - # the app is stopping) -- end the subscription cleanly. - break + res = await self._apoll_once() if res.gap: raise SharedStorageGap( f"replay buffer overran on topic {self._topic!r}" ) - for message in res.messages: - yield message + for pair in self._with_seq(res): + yield pair self._cursor = res.last_seq finally: self.close() +class _AsyncConn: + """A client worker's asyncio-streams request/response channel to the owner, + serialized by an asyncio.Lock (one request in flight per connection).""" + + def __init__(self, coord: _Coordinator): + self._coord = coord + self._streams = None + self._lock: Optional[asyncio.Lock] = None + + async def call(self, req): + if self._lock is None: + self._lock = asyncio.Lock() + async with self._lock: + if self._streams is None: + self._streams = await self._coord.aconnect() + reader, writer = self._streams + try: + await asend_frame(writer, req) + resp = await arecv_frame(reader) + if resp is EOF: + raise ConnectionError("owner closed the connection") + except (OSError, EOFError, ConnectionError, asyncio.TimeoutError): + with contextlib.suppress(Exception): + writer.close() + self._streams = None + raise + status, val = resp + if status == "err": + raise SharedStorageError(val) + return val + + class LocalSharedStorage(BaseSharedStorage): """In-memory shared storage, elected to a single owner process per machine. @@ -373,6 +501,10 @@ def __init__( self._coord = _Coordinator(ns, buffer_size, mode, path, flush_interval) self._conn: Optional[socket.socket] = None self._conn_lock = threading.Lock() + # Client role, asyncio callers: one asyncio-streams connection per + # event loop (with its own asyncio.Lock), so a loop never blocks on the + # sync connection or contends with request threads for it. + self._aconns: "dict[asyncio.AbstractEventLoop, _AsyncConn]" = {} def start(self) -> None: self._coord.ensure() @@ -437,6 +569,36 @@ def _remote(self, req): raise SharedStorageError(val) return val + async def _acall(self, req): + if self._coord.is_owner(): + return self._local(req) + loop = asyncio.get_running_loop() + conn = self._aconns.get(loop) + if conn is None: + conn = self._aconns[loop] = _AsyncConn(self._coord) + last_err: Optional[Exception] = None + for _ in range(3): + try: + return await conn.call(req) + except (OSError, EOFError, ConnectionError, asyncio.TimeoutError) as err: + last_err = err + self._coord.on_owner_lost() + if self._coord.is_owner(): + return self._local(req) + raise SharedStorageError(f"shared-storage owner unreachable: {last_err}") + + async def aget(self, key: str, default: Any = None) -> Any: + return await self._acall(["get", key, default]) + + async def aset(self, key: str, value: Any, ttl: Optional[float] = None) -> None: + await self._acall(["set", key, value, ttl]) + + async def adelete(self, key: str) -> None: + await self._acall(["delete", key]) + + async def apublish(self, topic: str, message: Any) -> None: + await self._acall(["publish", topic, message]) + def get(self, key: str, default: Any = None) -> Any: return self._call(["get", key, default]) @@ -455,4 +617,4 @@ def _head(self, topic: str) -> int: def subscribe(self, topic: str, replay_from: Optional[int] = None) -> Subscription: self._coord.ensure() start = replay_from if replay_from is not None else self._head(topic) - return _LocalSubscription(self._coord, topic, start) + return _LocalSubscription(self._coord, topic, start, self._call) diff --git a/dash/_shared_storage/redis.py b/dash/_shared_storage/redis.py index 2cea201da0..84e5175324 100644 --- a/dash/_shared_storage/redis.py +++ b/dash/_shared_storage/redis.py @@ -137,6 +137,12 @@ def _head(self, topic: str) -> int: 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 + # key was flushed, or evicted under a maxmemory policy). Gap so the + # consumer resets rather than blocking on XREAD until the sequence climbs + # back past the cursor. + if after_seq > self._head(topic): + return PollResult([], after_seq, True) # Gap: the next wanted sequence sits below the trimmed floor. Checked # before XREAD, which would otherwise silently resume at the floor. first = self._redis.xrange(stream, count=1) diff --git a/dash/_stream_hub.py b/dash/_stream_hub.py new file mode 100644 index 0000000000..3f6a1878fb --- /dev/null +++ b/dash/_stream_hub.py @@ -0,0 +1,672 @@ +"""Multiplexed streaming over shared storage. + +A browser holds a single downlink identified by a ``connection_id``. Every +streaming callback publishes its frames -- each tagged with the callback's +``request_id`` -- to that connection's shared-storage topic; the downlink reads +the topic and relays the frames to the client, which routes them back to the +right callback by ``request_id`` and closes the downlink once no streams remain +running. + +Because the frames travel through the shared store (not the HTTP response of the +callback that produced them), the worker that runs a callback and the worker +that holds the downlink do not have to be the same process -- the store is the +broker. Reconnecting a dropped downlink resumes from its cursor, so the store's +replay buffer covers the gap without losing frames. + +The connection id is the page's server-signed ``end_id`` (verified by the +backend, never taken from the client), so a page can only ever read or write its +own topic. The renderer hosts the downlink in a SharedWorker so every tab of the +browser shares one connection: the worker pins the ``end_id`` of the first tab +that streams and sends it with every request for that connection. + +Downlink line shape (one JSON object per NDJSON line):: + + {"rid": "", "frame": {}, "seq": } + +where ``frame`` is a ``CallbackExecutionResponse`` frame or a ``{"done": true}`` +terminal, exactly as the single-callback NDJSON transport emits today. A +``{"reset": true}`` envelope tells the client its cursor is stale (the store lost +this connection's frames) and it must reset to the head. + +Two downlink shapes. On ASGI the downlink is one long-lived response per +browser: it costs no thread, so it simply stays open (``async_downlink_marker``). +On WSGI a response holds a worker thread for its whole life, so a long-lived +downlink per browser would exhaust any thread pool at a few dozen browsers; +there the browser *polls* instead (``poll_downlink``): each request returns +whatever frames are queued since its cursor and ends at once, taking a thread +for milliseconds. The client re-polls at its poll interval while frames flow and +backs off while quiet, so latency stays near the interval and a small pool +serves many browsers. The pumps themselves are tasks on one event-loop thread +per WSGI process (``pump_to_storage``), not a thread per stream. + +Lifecycle. The pump that drives a callback and the downlink that relays its +frames may live on different workers, so "the browser went away" has to travel +through the store too: each downlink records its state under the connection's +key -- open/closed for a long-lived downlink, a heartbeat per poll otherwise -- +and every pump checks that record periodically. A downlink closed, silent, or +never opened for longer than the grace period means the browser is gone, and +the pump cancels its callback rather than running it to completion for nobody. +A closing tab, while other tabs keep the shared downlink open, instead sends an +explicit ``streamCancel`` for each of its requests, recorded under a +per-request key the same pump check picks up. + +Shutdown. A server drains in-flight responses before it stops, and a long-lived +downlink never ends on its own, so Ctrl+C would wait forever. A SIGINT/SIGTERM +handler (``install_stream_shutdown_handler``) closes every open subscription +and cancels every pump, then hands the signal on to the server's own handler. +""" + +import asyncio +import atexit +import contextlib +import json +import logging +import secrets +import signal +import threading +import time +from typing import Any, AsyncIterator, Callable, Iterator, List, Optional + +from ._shared_storage.base import BaseSharedStorage, SharedStorageGap, Subscription +from ._streaming import ( + StreamedCallbackResponse, + _shutdown as _streaming_shutdown, + to_json, +) + +logger = logging.getLogger(__name__) + +_TOPIC_PREFIX = "_dash_stream:" +_CONN_PREFIX = "_dash_stream_conn:" +_CANCEL_PREFIX = "_dash_stream_cancel:" + +# How long a downlink may stay closed (a reconnect in progress) before the +# pumps on its connection give up on the client. The renderer reconnects one +# second after a drop, so this is generous. +DOWNLINK_GRACE = 10.0 +# How long a polling browser may go without a poll reaching the server before +# it counts as gone. It polls at least once a second while it has streams, but +# on an overloaded WSGI pool its polls can queue for many seconds -- and an +# overload must cost latency, never the stream itself. +POLL_GRACE = 30.0 +# How often a pump consults the connection record while a callback runs. +DOWNLINK_CHECK_INTERVAL = 2.0 + +# The uplink's fast acknowledgement -- the streaming callback's POST returns this +# immediately; its outputs arrive on the downlink, not this response. +STREAM_ACK = {"multi": True, "stream": True} +# Acknowledgement of a ``streamCancel`` request. +STREAM_CANCEL_ACK = {"multi": True, "stream": True, "cancelled": True} + +# Control envelope telling the client its cursor is stale (its frames were lost +# to an owner re-election or a server restart) and it must reset to the head and +# resubscribe, rather than stall waiting for the fresh sequence to pass it. +RESET_ENVELOPE = {"reset": True} + +# Terminal frame a pump publishes when it cancels a callback because the +# downlink went away: a client that reconnects late resolves the request +# instead of waiting forever (and holding its downlink open for it). +_CANCELLED_FRAME = {"done": True} +# Terminal frame a pump publishes when a shutdown cancels it mid-stream. With an +# external store (Redis) it outlives the process, so a downlink reconnecting +# after the restart replays it and the client settles that callback instead of +# waiting on a pump that no longer exists. +_INTERRUPTED_FRAME = { + "done": True, + "error": { + "message": "Streaming callback interrupted: " + "the server shut down while it was running" + }, +} + + +def stream_topic(connection_id: str) -> str: + return f"{_TOPIC_PREFIX}{connection_id}" + + +def connection_key(connection_id: str) -> str: + return f"{_CONN_PREFIX}{connection_id}" + + +def cancel_key(connection_id: str, request_id: str) -> str: + return f"{_CANCEL_PREFIX}{connection_id}:{request_id}" + + +def _envelope(request_id: str, frame: Any) -> Any: + """Reduce a frame to plain JSON inside its downlink envelope. + + A frame may carry ``dash.Patch`` objects (and components) that only Dash's + JSON encoder understands; reduce it to a plain JSON structure here, before it + reaches shared storage, whose wire codec is data-only. This also matches what + the single-connection NDJSON path emits, so the client applies frames + identically either way. + """ + return {"rid": request_id, "frame": json.loads(to_json(frame))} + + +def publish_frame( + storage: BaseSharedStorage, + connection_id: str, + request_id: str, + frame: Any, +) -> None: + """Publish one streaming frame onto a connection's downlink topic.""" + storage.publish(stream_topic(connection_id), _envelope(request_id, frame)) + + +async def apublish_frame( + storage: BaseSharedStorage, + connection_id: str, + request_id: str, + frame: Any, +) -> None: + """:func:`publish_frame` for the pumps: never blocks their event loop.""" + await storage.apublish(stream_topic(connection_id), _envelope(request_id, frame)) + + +# --- downlink lifecycle record --------------------------------------------- + +# Open downlink subscriptions, so a server shutdown can close them (each one +# otherwise blocks its worker in a long poll, stalling a graceful shutdown). +# Guarded by a lock: subscriptions open/close on worker threads while a shutdown +# hook iterates the set, and a plain set is not safe against that. +_active_subscriptions: "set[Subscription]" = set() +_registry_lock = threading.Lock() + + +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. + """ + + def __init__( + self, + storage: BaseSharedStorage, + connection_id: str, + replay_from: Optional[int] = None, + ): + self.storage = storage + self.connection_id = connection_id + self.subscription = storage.subscribe(stream_topic(connection_id), replay_from) + 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: + if self._closed: + return + self._closed = True + with _registry_lock: + _active_subscriptions.discard(self.subscription) + self.subscription.close() + key = connection_key(self.connection_id) + 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, + }, + ) + + 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: + for seq, message in self.subscription.iter_with_seq(): + yield {**message, "seq": seq} + except SharedStorageGap: + yield dict(RESET_ENVELOPE) + finally: + self.close() + + async def aenvelopes(self) -> AsyncIterator[Any]: + """Async counterpart of :meth:`envelopes` for ASGI backends.""" + try: + async for seq, message in self.subscription.aiter_with_seq(): + yield {**message, "seq": seq} + except SharedStorageGap: + yield dict(RESET_ENVELOPE) + finally: + self.close() + + +def downlink_gone( + storage: BaseSharedStorage, + connection_id: str, + since: float, + grace: Optional[float] = None, +) -> bool: + """Whether a connection's downlink has been away for longer than ``grace``. + + A long-lived downlink is away once it recorded itself closed; a polling + browser is away once its last poll is older than ``grace``. ``since`` is + when the asking pump started: a record closed or last polled *before* that + (the client's previous streams finished and it went quiet) does not count + until the new downlink has had ``grace`` to show up, and a missing record -- + the downlink racing the uplink -- likewise gets ``grace`` to appear. + """ + return _judge_gone(storage.get(connection_key(connection_id)), since, grace) + + +def _judge_gone(record: Any, since: float, grace: Optional[float]) -> bool: + if not isinstance(record, dict): + last_alive = 0.0 + elif record.get("mode") == "poll": + last_alive = record.get("at", 0.0) # every poll is a heartbeat + if grace is None: + grace = POLL_GRACE + elif record.get("open"): + return False + else: + last_alive = record.get("at", 0.0) # when it closed + if grace is None: + grace = DOWNLINK_GRACE # read at call time so it stays tunable + return time.time() - max(last_alive, since) > grace + + +def cancel_stream( + storage: BaseSharedStorage, connection_id: str, request_id: str +) -> None: + """Ask the pump driving one callback to stop. + + The client sends this when a request's consumer is gone while the shared + downlink stays open for other tabs -- a browser tab closed -- so the + callback does not run to completion for nobody. The pump on whichever + worker runs it notices within ``DOWNLINK_CHECK_INTERVAL``. + """ + storage.set(cancel_key(connection_id, request_id), time.time()) + + +def stream_cancelled( + storage: BaseSharedStorage, connection_id: str, request_id: str +) -> bool: + return storage.get(cancel_key(connection_id, request_id)) is not None + + +def _stop_check( + storage: BaseSharedStorage, connection_id: str, request_id: str +) -> Callable[[], Any]: + """Whether a pump should give up: its request was cancelled, or the + connection's downlink has been gone for longer than the grace period. + Async, and through the store's loop-native operations: pumps share one + event loop, which a blocking round trip would stall for every stream.""" + started = time.time() + + async def stop() -> bool: + try: + if await storage.aget(cancel_key(connection_id, request_id)) is not None: + return True + record = await storage.aget(connection_key(connection_id)) + return _judge_gone(record, started, None) + except Exception: # pylint: disable=broad-exception-caught + # The store is unreachable; the pump's own publish will surface + # that. Don't cancel a callback over a transient lookup failure. + return False + + return stop + + +# --- downlinks --------------------------------------------------------------- + + +def subscribe_envelopes( + storage: BaseSharedStorage, + connection_id: str, + replay_from: Optional[int] = None, +) -> Iterator[Any]: + """Yield a connection's downlink envelopes until the subscription ends. + + Convenience over :class:`Downlink` for consumers that drive the relay + themselves (tests, tooling). Each envelope carries its ``seq`` so the client + can resume from it after a reconnect without losing frames. + """ + return Downlink(storage, connection_id, replay_from).envelopes() + + +# Last heartbeat written per connection by this process: a browser polls many +# times a second, the pumps only need to hear from it every POLL_GRACE. +_HEARTBEAT_INTERVAL = 1.0 +_last_heartbeat: "dict[str, float]" = {} +_HEARTBEAT_CACHE_LIMIT = 50_000 + + +def _heartbeat(storage: BaseSharedStorage, connection_id: str) -> None: + now = time.time() + if now - _last_heartbeat.get(connection_id, 0.0) < _HEARTBEAT_INTERVAL: + return + if len(_last_heartbeat) > _HEARTBEAT_CACHE_LIMIT: + _last_heartbeat.clear() # bounded; a miss only costs one extra write + _last_heartbeat[connection_id] = now + storage.set( + connection_key(connection_id), {"mode": "poll", "open": True, "at": now} + ) + + +def poll_downlink( + storage: BaseSharedStorage, + connection_id: str, + replay_from: Optional[int] = None, +) -> List[Any]: + """One poll of a connection's downlink for the WSGI path: the envelopes + published since ``replay_from``, without waiting, and (at most once a + second) a heartbeat on the connection record so the pumps know the + browser is still there. A lost buffer yields a single reset envelope. + + Costs a worker thread for milliseconds rather than for the browser's + whole visit, which is what lets a WSGI pool serve many browsers. + """ + sub = storage.subscribe(stream_topic(connection_id), replay_from or 0) + try: + pairs = sub.poll(0.0) + except SharedStorageGap: + return [dict(RESET_ENVELOPE)] + finally: + sub.close() + _heartbeat(storage, connection_id) + return [{**message, "seq": seq} for seq, message in pairs] + + +def async_downlink_marker( + storage: BaseSharedStorage, + connection_id: str, + replay_from: Optional[int] = None, +) -> StreamedCallbackResponse: + """A downlink as a ``StreamedCallbackResponse`` for the ASGI NDJSON path.""" + downlink = Downlink(storage, connection_id, replay_from) + return StreamedCallbackResponse(downlink.aenvelopes(), is_async=True) + + +# --- pumps ------------------------------------------------------------------- + +# In-flight pump tasks (ASGI server loop or the WSGI pump loop): keeps them +# referenced so their loop doesn't GC them mid-stream, and lets shutdown cancel +# them. +_pending_pumps: "set[asyncio.Task]" = set() + + +async def _publish_terminal(storage, connection_id, request_id, frame): + with contextlib.suppress(Exception): + await apublish_frame(storage, connection_id, request_id, frame) + + +async def _forget_cancel(storage, connection_id, request_id): + with contextlib.suppress(Exception): + await storage.adelete(cancel_key(connection_id, request_id)) + + +async def apump_to_storage( + storage: BaseSharedStorage, + connection_id: str, + request_id: str, + marker: StreamedCallbackResponse, +) -> None: + """Drive an async streaming callback and publish each frame to the topic. + + Runs as a task -- on the server's loop under ASGI, on the process's shared + pump loop under WSGI (:func:`pump_to_storage`) -- so the callback's POST + can return immediately. The frame generator already emits the terminal + ``{"done": True}``; publishing it lets the client resolve that request. + + Gives up once the request is cancelled or the downlink has been gone for + the grace period: the pending step is cancelled, which raises into the + user generator at its current ``await``, and a plain terminal frame is + published. A shutdown cancelling the task mid-stream publishes an error + terminal instead, best effort, then re-raises. + """ + stop = _stop_check(storage, connection_id, request_id) + iterator = marker.frames.__aiter__() + next_check = time.monotonic() + DOWNLINK_CHECK_INTERVAL + completed = False + step = None + try: + while True: + step = asyncio.ensure_future(iterator.__anext__()) + while True: + done, _ = await asyncio.wait( + {step}, timeout=max(0.0, next_check - time.monotonic()) + ) + if done: + break + next_check = time.monotonic() + DOWNLINK_CHECK_INTERVAL + if await stop(): + return + try: + frame = step.result() + except StopAsyncIteration: + return + finally: + step = None + await apublish_frame(storage, connection_id, request_id, frame) + completed = bool(frame.get("done")) + if time.monotonic() >= next_check: + next_check = time.monotonic() + DOWNLINK_CHECK_INTERVAL + if await stop(): + return + except asyncio.CancelledError: + completed = True + await _publish_terminal(storage, connection_id, request_id, _INTERRUPTED_FRAME) + raise + except Exception: # pylint: disable=broad-exception-caught + logger.exception("Streaming callback pump failed") + completed = True + await _publish_terminal(storage, connection_id, request_id, _CANCELLED_FRAME) + finally: + if step is not None and not step.done(): + step.cancel() + with contextlib.suppress(BaseException): + await step + with contextlib.suppress(Exception): + await marker.frames.aclose() + if not completed: + await _publish_terminal( + storage, connection_id, request_id, _CANCELLED_FRAME + ) + await _forget_cancel(storage, connection_id, request_id) + + +def spawn_async_pump( + storage: BaseSharedStorage, + connection_id: str, + request_id: str, + marker: StreamedCallbackResponse, +) -> None: + """Run the pump as a fire-and-forget task on the ASGI event loop, so the + callback's request returns immediately while frames keep flowing. + """ + task = asyncio.ensure_future( + apump_to_storage(storage, connection_id, request_id, marker) + ) + _pending_pumps.add(task) + task.add_done_callback(_pending_pumps.discard) + + +# --- WSGI pumps: one event-loop thread per process --------------------------- +# +# A pump drives an async generator, so it is naturally a task. Under WSGI there +# is no server loop to put it on, so the process runs one of its own: every +# pump in the worker is a task on that loop, and a worker can drive thousands +# of streams without a thread (let alone a thread pair and a private loop) per +# stream. + +_pump_loop: Optional[asyncio.AbstractEventLoop] = None +_pump_thread: Optional[threading.Thread] = None +_pump_loop_lock = threading.Lock() +_pump_exit_hook_installed = False +_pump_loop_stopping = False + + +def _run_pump_loop(loop: asyncio.AbstractEventLoop, ready: threading.Event) -> None: + asyncio.set_event_loop(loop) + ready.set() + while not loop.is_closed(): + try: + loop.run_forever() + return # stopped on purpose + except BaseException: # pylint: disable=broad-exception-caught + # Something was raised *into* this thread (a test harness that + # stops "every thread the app started", a stray async exception). + # This thread carries every stream in the process; it does not + # exit on anyone's behalf but its own shutdown. + if _pump_loop_stopping: + return + logger.warning("Stream pump loop interrupted; resuming", exc_info=True) + + +def _shared_pump_loop() -> asyncio.AbstractEventLoop: + global _pump_loop, _pump_thread # pylint: disable=global-statement + global _pump_exit_hook_installed # pylint: disable=global-statement + with _pump_loop_lock: + thread_alive = _pump_thread is not None and _pump_thread.is_alive() + if _pump_loop is None or _pump_loop.is_closed() or not thread_alive: + if _pump_loop is not None and not _pump_loop.is_closed(): + # The thread died under the loop: its pumps are lost, but + # streaming must not stay dead for the process. + logger.warning("Stream pump loop thread died; starting a new one") + with contextlib.suppress(Exception): + _pump_loop.close() + loop = asyncio.new_event_loop() + ready = threading.Event() + thread = threading.Thread( + target=_run_pump_loop, + args=(loop, ready), + daemon=True, + name="dash-stream-pumps", + ) + thread.start() + ready.wait() + _pump_loop, _pump_thread = loop, thread + if not _pump_exit_hook_installed: + # End the pumps before interpreter teardown finalizes their + # loop and generators half-way (which logs spurious errors). + _pump_exit_hook_installed = True + atexit.register(_stop_pump_loop) + return _pump_loop + + +def _stop_pump_loop() -> None: + global _pump_loop_stopping # pylint: disable=global-statement + loop, thread = _pump_loop, _pump_thread + if loop is None or loop.is_closed() or thread is None or not thread.is_alive(): + return + _pump_loop_stopping = True + shutdown_active_streams() + + async def drain(): + pending = [t for t in _pending_pumps if not t.done()] + if pending: + await asyncio.wait(pending, timeout=2.0) + + with contextlib.suppress(Exception): + asyncio.run_coroutine_threadsafe(drain(), loop).result(timeout=3.0) + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=2.0) + + +async def _tracked_pump(storage, connection_id, request_id, marker): + task = asyncio.current_task() + _pending_pumps.add(task) + try: + await apump_to_storage(storage, connection_id, request_id, marker) + finally: + _pending_pumps.discard(task) + + +def pump_to_storage( + storage: BaseSharedStorage, + connection_id: str, + request_id: str, + marker: StreamedCallbackResponse, +): + """WSGI entry point: schedule the pump for one streaming callback on the + process's shared pump loop and return at once, so the callback's POST acks + immediately. Returns a ``concurrent.futures.Future`` that resolves when + the pump ends (useful to wait on in tests). + """ + return asyncio.run_coroutine_threadsafe( + _tracked_pump(storage, connection_id, request_id, marker), + _shared_pump_loop(), + ) + + +# --- process shutdown -------------------------------------------------------- + + +def shutdown_active_streams() -> None: + """Stop every in-flight stream so the server can shut down. + + Sets the module-level shutdown flag so keepalive generators exit on their + next timeout, cancels the pump tasks (on whichever loop they run), and + closes open downlink subscriptions. Each downlink otherwise sits in a long + poll that a graceful shutdown would wait on forever. Backends call this + from their shutdown hook. Idempotent and safe to call when nothing is + streaming, and from a signal handler. + """ + _streaming_shutdown.set() + for task in list(_pending_pumps): + with contextlib.suppress(Exception): + task.get_loop().call_soon_threadsafe(task.cancel) + with _registry_lock: + subscriptions = list(_active_subscriptions) + for sub in subscriptions: + with contextlib.suppress(Exception): + sub.close() + + +def install_stream_shutdown_handler(): + """Install SIGINT/SIGTERM handlers that tear down active streams. + + Without this, a streaming response generator blocks the server's + worker thread/task, and the process ignores Ctrl+C: the server's + graceful shutdown waits for connections to drain, connections wait + for the response to finish, and the response waits for the next + frame that will never come. The handler breaks the cycle by setting + the shutdown flag and closing subscriptions before the server even + begins its shutdown sequence, then chains to the handler that was + installed before it so the server's own shutdown runs normally. + + Runs at import, and again from backend startup hooks: uvicorn imports + the app before it installs its own handlers, which replace whatever + was there, so the import-time install is lost and must be redone once + the server is listening. Safe to call repeatedly; a signal whose + handler is already ours is left alone. + """ + if threading.current_thread() is not threading.main_thread(): + return + + for signum in (signal.SIGINT, signal.SIGTERM): + original = signal.getsignal(signum) + if getattr(original, "_dash_stream_shutdown", False): + continue + + def _handler(sig, frame, original=original): + shutdown_active_streams() + if callable(original): + original(sig, frame) + elif original == signal.SIG_DFL: + signal.signal(sig, signal.SIG_DFL) + signal.raise_signal(sig) + + _handler._dash_stream_shutdown = True # pylint: disable=protected-access + signal.signal(signum, _handler) + + +install_stream_shutdown_handler() diff --git a/dash/_streaming.py b/dash/_streaming.py new file mode 100644 index 0000000000..1e3f421ea5 --- /dev/null +++ b/dash/_streaming.py @@ -0,0 +1,341 @@ +"""Transport helpers for streaming callbacks (generator callbacks). + +A streaming callback is a generator (or async generator) whose yields are +converted to "frames" — dicts with the same shape as a regular callback +response (see ``CallbackExecutionResponse``) — followed by a terminal +``{"done": True}`` frame. Frames are delivered to the renderer either as +NDJSON lines on the HTTP response or as individual messages over the +WebSocket callback transport. + +The frame generators are built in ``dash._callback``; this module owns the +marker object the backends dispatch on and the transport-side iteration +helpers. Iteration helpers exist because Dash's callback context lives in a +``contextvars.ContextVar`` and a sync generator runs in whatever context its +consumer drives it from — which, for a streaming HTTP response, is not the +request context the callback started in. + +The NDJSON transports also emit keepalives: a blank line every +``keepalive`` seconds the callback spends between yields. Every proxy in a +typical deployment enforces an idle timeout on the response (nginx's +``proxy_read_timeout`` defaults to 60s), and a callback that thinks for +longer than that gets its connection closed mid-stream. A blank line resets +those timers; the renderer skips empty lines, so it costs nothing on the +client. The WebSocket transport has its own heartbeat +(``websocket_heartbeat_interval``) and does not use these helpers. +""" + +import asyncio +import contextlib +import functools +import logging +import queue +import threading +import time +from typing import cast + +from ._utils import to_json as _to_json + +logger = logging.getLogger(__name__) + +_shutdown = threading.Event() + +STREAM_MIMETYPE = "application/x-ndjson" +# Disable proxy/server buffering so frames reach the browser as they are +# produced (X-Accel-Buffering covers nginx). +STREAM_HEADERS = {"X-Accel-Buffering": "no", "Cache-Control": "no-cache"} + +# Emitted when the callback is quiet for longer than the keepalive interval. +# The renderer's NDJSON reader skips blank lines. +KEEPALIVE_LINE = "\n" + +_SENTINEL = object() +_KEEPALIVE = object() + + +def keepalive_seconds(interval_ms): + """Normalize a configured keepalive interval (ms) to seconds, or None.""" + if not interval_ms or interval_ms <= 0: + return None + return interval_ms / 1000 + + +def to_json(value) -> str: + return cast(str, _to_json(value)) + + +class StreamedCallbackResponse: # pylint: disable=too-few-public-methods + """Marker returned by streaming callback wrappers. + + Backends detect this instead of a JSON string and return a streaming + response. ``frames`` is a generator (async generator when ``is_async``) + of frame dicts. ``ctx`` is the ``contextvars`` snapshot captured when the + callback was invoked; sync frame generators must be driven through it + (``iter_stream_frames``) so ``dash.ctx``/``set_props`` keep working after + the dispatch function has returned. + """ + + def __init__(self, frames, is_async, ctx=None): + self.frames = frames + self.is_async = is_async + self.ctx = ctx + + +def iter_stream_frames(marker): + """Drive a sync frame generator inside its captured context snapshot. + + Closing this generator closes the frame generator too, inside the + callback context so ``dash.ctx`` resolves in its cleanup handlers. + """ + try: + while True: + try: + yield marker.ctx.run(next, marker.frames) + except StopIteration: + return + finally: + with contextlib.suppress(Exception): + marker.ctx.run(marker.frames.close) + + +async def aiter_stream_frames(marker): + """Async wrapper for a sync frame generator (ASGI backends). + + Each step runs on an executor thread through ``marker.ctx`` — Starlette's + own threadpool iteration would use a fresh context copy per chunk and + lose the callback context. + """ + loop = asyncio.get_running_loop() + while True: + frame = await loop.run_in_executor( + None, marker.ctx.run, functools.partial(next, marker.frames, _SENTINEL) + ) + if frame is _SENTINEL: + return + yield frame + + +def _serialize_frame(frame): + """Serialize one frame to an NDJSON line. + + Returns ``(line, fatal)``; a serialization failure produces a terminal + error frame so the client is not left waiting on a silently dead stream. + """ + try: + return to_json(frame) + "\n", False + except TypeError as err: + logger.exception("Failed to serialize streamed callback frame") + return ( + to_json( + { + "done": True, + "error": { + "message": "Non-serializable value in streamed " + f"callback output: {err}" + }, + } + ) + + "\n", + True, + ) + + +def _keepalive_frames(marker, keepalive): + """Yield frames from a sync generator, plus keepalives while it is quiet. + + A blocking ``next()`` cannot be interrupted on a timer, so the frame + generator is driven on a pump thread and this generator waits on a queue + instead. One consequence: a client disconnect no longer raises + ``GeneratorExit`` into the user generator at its current yield — the pump + notices the stop flag once the next frame arrives, and closes it then. The + async path (``async def`` callbacks) cancels at the yield as before, which + is one more reason ``dash._callback`` recommends async for streams. + """ + frames: queue.Queue = queue.Queue(maxsize=1) + stop = threading.Event() + + def put(item): + """Hand one item to the consumer; False if it went away.""" + while not stop.is_set(): + try: + frames.put(item, timeout=0.2) + return True + except queue.Full: + continue + return False + + def pump(): + try: + for frame in iter_stream_frames(marker): + if not put(("item", frame)): + return + put(("end", None)) + except BaseException as err: # pylint: disable=broad-exception-caught + put(("error", err)) + finally: + # Run the user generator's cleanup inside the callback context so + # dash.ctx still resolves in its GeneratorExit/finally handlers. + with contextlib.suppress(Exception): + marker.ctx.run(marker.frames.close) + + thread = threading.Thread(target=pump, daemon=True, name="dash-stream-pump") + thread.start() + poll = min(keepalive, 0.5) if keepalive else 0.5 + try: + last_activity = time.monotonic() + while not _shutdown.is_set(): + try: + kind, value = frames.get(timeout=poll) + except queue.Empty: + if keepalive and time.monotonic() - last_activity >= keepalive: + yield _KEEPALIVE + last_activity = time.monotonic() + continue + last_activity = time.monotonic() + if kind == "item": + yield value + elif kind == "error": + raise value + else: + return + finally: + stop.set() + + +async def _akeepalive_frames(frames, keepalive): + """Yield frames from an async iterator, plus keepalives while it is quiet. + + The pending ``__anext__`` is held across timeouts rather than awaited with + ``asyncio.wait_for``, which would cancel the user generator mid-step every + time a keepalive was due. + """ + poll = min(keepalive, 0.5) if keepalive else 0.5 + iterator = frames.__aiter__() + pending = None + try: + while not _shutdown.is_set(): + pending = asyncio.ensure_future(iterator.__anext__()) + last_activity = time.monotonic() + while not _shutdown.is_set(): + done, _ = await asyncio.wait({pending}, timeout=poll) + if done: + break + if keepalive and time.monotonic() - last_activity >= keepalive: + yield _KEEPALIVE + last_activity = time.monotonic() + if _shutdown.is_set(): + return + try: + frame = pending.result() + except StopAsyncIteration: + return + finally: + pending = None + yield frame + finally: + if pending is not None and not pending.done(): + pending.cancel() + + +def _line(frame): + """Serialize one frame or keepalive; ``(line, fatal)`` as _serialize_frame.""" + if frame is _KEEPALIVE: + return KEEPALIVE_LINE, False + return _serialize_frame(frame) + + +def ndjson_lines(marker, keepalive=None): + """Sync NDJSON body for a sync frame generator (Flask/WSGI).""" + if keepalive: + frames = _keepalive_frames(marker, keepalive) + else: + frames = iter_stream_frames(marker) + try: + for frame in frames: + line, fatal = _line(frame) + yield line + if fatal: + return + finally: + # Reached on client disconnect too (the WSGI server closes the body + # iterator): end the frame source rather than leave it to the GC. + frames.close() + + +async def andjson_lines(frames, keepalive=None): + """Async NDJSON body over an async iterator of frames.""" + if keepalive: + frames = _akeepalive_frames(frames, keepalive) + async for frame in frames: + line, fatal = _line(frame) + yield line + if fatal: + return + + +def marker_ndjson_aiter(marker, keepalive=None): + """Async NDJSON body for either flavor of frame generator.""" + if marker.is_async: + return andjson_lines(marker.frames, keepalive) + return andjson_lines(aiter_stream_frames(marker), keepalive) + + +def sync_iter_asyncgen(agen): + """Iterate an async generator from sync code (Flask + async gen). + + Runs the whole consumption on one task on a private event-loop thread so + contextvars set inside the generator persist across steps. Closing this + generator (client disconnect) cancels the task, which raises into the + user generator at its current yield, and waits briefly for the loop + thread to wind down so the generator's cleanup has run -- and so nothing + is left for interpreter shutdown to finalize noisily. + """ + frame_queue: queue.Queue = queue.Queue() + loop = asyncio.new_event_loop() + task = None + task_ready = threading.Event() + + async def consume(): + try: + async for item in agen: + frame_queue.put(("item", item)) + frame_queue.put(("end", None)) + except BaseException as err: # pylint: disable=broad-exception-caught + frame_queue.put(("error", err)) + finally: + with contextlib.suppress(Exception): + await agen.aclose() + + def run(): + nonlocal task + asyncio.set_event_loop(loop) + task = loop.create_task(consume()) + task_ready.set() + try: + loop.run_until_complete(task) + except BaseException: # pylint: disable=broad-exception-caught + pass + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + thread = threading.Thread(target=run, daemon=True, name="dash-stream-bridge") + thread.start() + try: + while True: + kind, value = frame_queue.get() + if kind == "item": + yield value + elif kind == "error": + if isinstance(value, asyncio.CancelledError): + return + raise value + else: + return + finally: + task_ready.wait(timeout=5) + if task is not None and not task.done(): + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(task.cancel) + # Bounded: a generator that ignores cancellation must not hang the + # request thread that is closing us. + thread.join(timeout=2.0) diff --git a/dash/backends/_fastapi.py b/dash/backends/_fastapi.py index e617d5f20b..86572d26c5 100644 --- a/dash/backends/_fastapi.py +++ b/dash/backends/_fastapi.py @@ -20,7 +20,7 @@ try: from fastapi import FastAPI, Request, Response, Body - from fastapi.responses import JSONResponse, RedirectResponse + from fastapi.responses import JSONResponse, RedirectResponse, StreamingResponse from fastapi.staticfiles import StaticFiles from starlette.responses import Response as StarletteResponse from starlette.datastructures import MutableHeaders @@ -36,6 +36,25 @@ from dash.fingerprint import check_fingerprint from dash import _validate, get_app +from dash._streaming import ( + STREAM_HEADERS, + STREAM_MIMETYPE, + StreamedCallbackResponse, + _shutdown as _streaming_shutdown, + keepalive_seconds, + marker_ndjson_aiter, + to_json, +) +from dash._callback import get_stream_connection_id +from dash._stream_hub import ( + STREAM_ACK, + STREAM_CANCEL_ACK, + async_downlink_marker, + cancel_stream, + install_stream_shutdown_handler, + shutdown_active_streams, + spawn_async_pump, +) from dash.exceptions import PreventUpdate from dash._compression import decompress_payload from .base_server import ( @@ -49,6 +68,7 @@ run_callback_in_executor, run_callback_on_loop, make_callback_done_handler, + make_stream_frame_emitter, shutdown_ws_connection, ) from ._utils import format_traceback_html @@ -57,6 +77,19 @@ from dash import Dash +def _run_subprocess(args, env): + proc = subprocess.Popen(args, env=env) # pylint: disable=R1732 + try: + proc.wait() + except KeyboardInterrupt: + proc.terminate() + try: + proc.wait(timeout=3) + except (KeyboardInterrupt, subprocess.TimeoutExpired): + proc.kill() + proc.wait() + + class FastAPIResponseAdapter(ResponseAdapter): """ A custom Response class that wraps FastAPI's JSONResponse @@ -77,7 +110,8 @@ def set_response(self, **kwargs): """ data = kwargs.get("data") if isinstance(data, (str, bytes, bytearray)): - resp = Response(content=data) + # Already-serialized JSON, like the Flask adapter's default content type. + resp = Response(content=data, media_type="application/json") else: resp = JSONResponse(content=data) if self._headers: @@ -216,7 +250,17 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: except Exception: # pylint: disable=broad-exception-caught traceback.print_exc() await self._initialize_dev_tools() - await self.app(scope, receive, send) + + async def _receive_with_shutdown(): + msg = await receive() + if msg.get("type") == "lifespan.startup": + _streaming_shutdown.clear() + install_stream_shutdown_handler() + elif msg.get("type") == "lifespan.shutdown": + shutdown_active_streams() + return msg + + await self.app(scope, _receive_with_shutdown, send) return # Non-HTTP/WebSocket scopes pass through @@ -257,6 +301,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: class FastAPIDashServer(BaseDashServer[FastAPI]): websocket_capability: bool = True + downlink_mode: str = "stream" def __init__(self, server: FastAPI): super().__init__(server) @@ -472,9 +517,7 @@ def run(self, dash_app: Dash, host, port, debug, **kwargs): # pylint: disable=R # Add any other kwargs as CLI args if needed - # pylint: disable=R1732 - proc = subprocess.Popen(uvicorn_args, env=env) - proc.wait() + _run_subprocess(uvicorn_args, env) def make_response( self, @@ -549,12 +592,53 @@ def add_redirect_rule(self, app, fullname, path): ) def serve_callback(self, dash_app: Dash): - async def _dispatch(request: Request): # pylint: disable=unused-argument + def _ndjson_response(marker): + # pylint: disable=protected-access + return StreamingResponse( + marker_ndjson_aiter( + marker, + keepalive_seconds(dash_app._stream_keepalive_interval), + ), + media_type=STREAM_MIMETYPE, + headers=dict(STREAM_HEADERS), + ) + + async def _dispatch( + request: Request, + ): # pylint: disable=unused-argument,too-many-return-statements # pylint: disable=protected-access if "gzip" in request.headers.get("content-encoding", ""): body = decompress_payload(await self.request_adapter()._request.body()) else: body = self.request_adapter().get_json() + downlink = body.get("streamDownlink") + if downlink is not None: + connection_id = get_stream_connection_id() + if connection_id is None: + return Response(status_code=403) + return _ndjson_response( + async_downlink_marker( + dash_app.shared_storage, + connection_id, + downlink.get("from"), + ) + ) + cancel = body.get("streamCancel") + if cancel is not None: + # A tab closed while the shared downlink stays open for others: + # stop its pump. Same connection-id rule as the downlink: a page + # can only cancel its own streams. + connection_id = ( + get_stream_connection_id() + if dash_app.shared_storage_enabled + else None + ) + if connection_id is None: + return Response(status_code=403) + cancel_stream( + dash_app.shared_storage, connection_id, cancel["requestId"] + ) + return JSONResponse(content=STREAM_CANCEL_ACK) cb_ctx = dash_app._initialize_context( body ) # pylint: disable=protected-access @@ -571,6 +655,29 @@ async def _dispatch(request: Request): # pylint: disable=unused-argument response_data = ctx.run(partial_func) if inspect.iscoroutine(response_data): response_data = await response_data + if isinstance(response_data, StreamedCallbackResponse): + stream_conn = body.get("streamConnection") + if stream_conn is not None: + # A streamConnection asserts "multiplex me": the connection + # id is derived from the signed end_id (never from the + # client), so a page can only publish to its own topic. An + # unverified connection is refused, never run some other way, + # so no frame reaches a topic without a valid token. + connection_id = ( + get_stream_connection_id() + if dash_app.shared_storage_enabled + else None + ) + if connection_id is None: + return Response(status_code=403) + spawn_async_pump( + dash_app.shared_storage, + connection_id, + stream_conn["requestId"], + response_data, + ) + return cb_ctx.dash_response.set_response(data=to_json(STREAM_ACK)) + return _ndjson_response(response_data) return cb_ctx.dash_response.set_response(data=response_data) return _dispatch @@ -817,6 +924,12 @@ async def websocket_handler(websocket: WebSocket): renderer_id, shutdown_event, ) + stream_emitter = make_stream_frame_emitter( + outbound_queue, + request_id, + renderer_id, + shutdown_event, + ) if is_async: task = asyncio.create_task( @@ -825,6 +938,7 @@ async def websocket_handler(websocket: WebSocket): payload, ws_cb, FastAPIResponseAdapter(), + stream_emitter, ) ) task.add_done_callback(done_handler) @@ -837,6 +951,7 @@ async def websocket_handler(websocket: WebSocket): payload, ws_cb, FastAPIResponseAdapter(), + stream_emitter, ) # Set up done callback to send response future.add_done_callback(done_handler) diff --git a/dash/backends/_flask.py b/dash/backends/_flask.py index 024875bdb8..8054a25ac7 100644 --- a/dash/backends/_flask.py +++ b/dash/backends/_flask.py @@ -1,11 +1,11 @@ from __future__ import annotations import asyncio +import inspect +import mimetypes import pkgutil import sys -import mimetypes import time -import inspect import traceback from contextvars import copy_context @@ -22,14 +22,36 @@ g as flask_g, has_request_context, redirect, + stream_with_context, ) from werkzeug.debug import tbtools from dash.fingerprint import check_fingerprint from dash import _validate from dash.exceptions import PreventUpdate, InvalidResourceError -from dash._callback import _invoke_callback, _async_invoke_callback +from dash._callback import ( + _invoke_callback, + _async_invoke_callback, + get_stream_connection_id, +) from dash._compression import decompress_payload +from dash._streaming import ( + STREAM_HEADERS, + STREAM_MIMETYPE, + StreamedCallbackResponse, + _shutdown as _streaming_shutdown, + andjson_lines, + keepalive_seconds, + ndjson_lines, + sync_iter_asyncgen, + to_json, +) +from dash._stream_hub import ( + STREAM_CANCEL_ACK, + cancel_stream, + poll_downlink, + pump_to_storage, +) from dash._utils import parse_version from .base_server import BaseDashServer, RequestAdapter, ResponseAdapter @@ -167,6 +189,7 @@ def has_request_context(self) -> bool: return has_request_context() def run(self, dash_app: Dash, host: str, port: int, debug: bool, **kwargs: Any): + _streaming_shutdown.clear() self.server.run(host=host, port=port, debug=debug, **kwargs) def make_response( @@ -249,13 +272,110 @@ def add_redirect_rule(self, app, fullname, path): self._create_redirect_function(app.get_relative_path(path)), ) - # pylint: disable=unused-argument + # pylint: disable=unused-argument,too-many-statements def serve_callback(self, dash_app: Dash): - def _dispatch(): - if "gzip" in request.headers.get("Content-Encoding", ""): - body = decompress_payload(request.data) + def _stream_response( + marker: StreamedCallbackResponse, with_request_ctx: bool + ) -> Response: + keepalive = keepalive_seconds( + dash_app._stream_keepalive_interval # pylint: disable=protected-access + ) + if marker.is_async: + # Drive the async frame generator on a private event-loop + # thread; the response iterator drains it synchronously. + body = sync_iter_asyncgen(andjson_lines(marker.frames, keepalive)) else: - body = request.get_json() + body = ndjson_lines(marker, keepalive) + if with_request_ctx: + # Keep flask.request usable while the body is iterated (the + # generator runs after the view returns). Only valid on the + # sync dispatch path: under an async view (asgiref) the + # request context lives in a different contextvars context + # and stream_with_context would corrupt its teardown. Dash's + # own callback context always works — it travels in the + # marker's context snapshot. + body = stream_with_context(body) + return Response( + body, + content_type=STREAM_MIMETYPE, + headers=dict(STREAM_HEADERS), + ) + + def _serve_downlink(downlink): + # One poll of the browser's multiplexed downlink: the frames queued + # since its cursor (published by streaming callbacks, possibly on + # other workers, via shared storage), as an NDJSON body that ends at + # once. A WSGI response holds a worker thread for its whole life, + # so the browser polls rather than keeping a connection open. + # The connection id is derived from the signed end_id, never taken + # from the client, so a page can only ever read its own topic. + connection_id = get_stream_connection_id() + if connection_id is None: + return Response(status=403) + envelopes = poll_downlink( + dash_app.shared_storage, connection_id, downlink.get("from") + ) + body = "".join(to_json(envelope) + "\n" for envelope in envelopes) + return Response( + body, content_type=STREAM_MIMETYPE, headers=dict(STREAM_HEADERS) + ) + + def _serve_cancel(cancel): + # A tab closed while the shared downlink stays open for others: + # stop the pump driving this request, wherever it runs. Same + # connection-id rule as the downlink: a page can only cancel its own. + connection_id = ( + get_stream_connection_id() if dash_app.shared_storage_enabled else None + ) + if connection_id is None: + return Response(status=403) + cancel_stream(dash_app.shared_storage, connection_id, cancel["requestId"]) + return Response(to_json(STREAM_CANCEL_ACK), content_type="application/json") + + def _serve_uplink(marker, body, cb_ctx): + # Multiplexed uplink: a streamConnection asserts "multiplex me". Pump + # the frames onto that connection's topic (a task on the process's + # pump loop) and return immediately, so this request does not hold a + # connection for the stream's life. The connection id is derived + # from the signed end_id, never from the client, so a page can only + # publish to its own topic; an unverified connection is refused + # (403), never run some other way, so no frame can reach a topic + # without a valid token. Returns None only when there is no + # multiplex intent (no streamConnection), to fall back to inline + # NDJSON. + stream_conn = body.get("streamConnection") + if stream_conn is None: + return None + connection_id = ( + get_stream_connection_id() if dash_app.shared_storage_enabled else None + ) + if connection_id is None: + return Response(status=403) + pump_to_storage( + dash_app.shared_storage, + connection_id, + stream_conn["requestId"], + marker, + ) + return cb_ctx.dash_response.set_response( + data=to_json({"multi": True, "stream": True}) + ) + + def _read_body(): + return ( + decompress_payload(request.data) + if "gzip" in request.headers.get("Content-Encoding", "") + else request.get_json() + ) + + def _dispatch(): + body = _read_body() + downlink = body.get("streamDownlink") + if downlink is not None: + return _serve_downlink(downlink) + cancel = body.get("streamCancel") + if cancel is not None: + return _serve_cancel(cancel) # pylint: disable=protected-access cb_ctx = dash_app._initialize_context(body) func = dash_app._prepare_callback(cb_ctx, body) @@ -265,6 +385,11 @@ def _dispatch(): func, args, cb_ctx.outputs_list, cb_ctx ) response_data = ctx.run(partial_func) + if isinstance(response_data, StreamedCallbackResponse): + uplink = _serve_uplink(response_data, body, cb_ctx) + if uplink is not None: + return uplink + return _stream_response(response_data, with_request_ctx=True) if asyncio.iscoroutine(response_data): raise Exception( "You are trying to use a coroutine without dash[async]. " @@ -274,10 +399,13 @@ def _dispatch(): return cb_ctx.dash_response.set_response(data=response_data) async def _dispatch_async(): - if "gzip" in request.headers.get("Content-Encoding", ""): - body = decompress_payload(request.data) - else: - body = request.get_json() + body = _read_body() + downlink = body.get("streamDownlink") + if downlink is not None: + return _serve_downlink(downlink) + cancel = body.get("streamCancel") + if cancel is not None: + return _serve_cancel(cancel) # pylint: disable=protected-access cb_ctx = dash_app._initialize_context(body) func = dash_app._prepare_callback(cb_ctx, body) @@ -289,6 +417,11 @@ async def _dispatch_async(): response_data = ctx.run(partial_func) if asyncio.iscoroutine(response_data): response_data = await response_data + if isinstance(response_data, StreamedCallbackResponse): + uplink = _serve_uplink(response_data, body, cb_ctx) + if uplink is not None: + return uplink + return _stream_response(response_data, with_request_ctx=False) return cb_ctx.dash_response.set_response(data=response_data) # Preserve the view function's identity as `dash.dash.dispatch` so that diff --git a/dash/backends/_quart.py b/dash/backends/_quart.py index 34af66102b..242e9f8e5f 100644 --- a/dash/backends/_quart.py +++ b/dash/backends/_quart.py @@ -40,6 +40,25 @@ from dash.exceptions import PreventUpdate, InvalidResourceError from dash.fingerprint import check_fingerprint +from dash._streaming import ( + STREAM_HEADERS, + STREAM_MIMETYPE, + StreamedCallbackResponse, + _shutdown as _streaming_shutdown, + keepalive_seconds, + marker_ndjson_aiter, + to_json, +) +from dash._callback import get_stream_connection_id +from dash._stream_hub import ( + STREAM_ACK, + STREAM_CANCEL_ACK, + async_downlink_marker, + cancel_stream, + install_stream_shutdown_handler, + shutdown_active_streams, + spawn_async_pump, +) from dash._utils import parse_version from dash import _validate from dash._compression import decompress_payload @@ -54,6 +73,7 @@ run_callback_in_executor, run_callback_on_loop, make_callback_done_handler, + make_stream_frame_emitter, shutdown_ws_connection, ) from ._utils import format_traceback_html @@ -92,6 +112,7 @@ def set_response(self, **kwargs): class QuartDashServer(BaseDashServer[Quart]): websocket_capability: bool = True + downlink_mode: str = "stream" def __init__(self, server: Quart) -> None: super().__init__(server) @@ -282,12 +303,14 @@ def run(self, dash_app: Dash, host: str, port: int, debug: bool, **kwargs: _t.An loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) + _streaming_shutdown.clear() # Initialize shutdown event for WebSocket handlers self._ws_shutdown_event = asyncio.Event() def signal_handler(): - """Handle shutdown signal by setting the WebSocket shutdown event.""" + """Handle shutdown signal by tearing down streams and WebSockets.""" + shutdown_active_streams() if self._ws_shutdown_event is not None: self._ws_shutdown_event.set() @@ -384,12 +407,66 @@ def add_redirect_rule(self, app, fullname, path): # pylint: disable=unused-argument def serve_callback(self, dash_app: Dash): # type: ignore[override] # Quart always async - async def _dispatch(): + # Under a bare ASGI server (uvicorn) the app is imported before the + # server installs its signal handlers, which drop the one Dash set at + # import; put it back once serving starts so Ctrl+C still ends active + # streams. Under Quart's own run() the loop handler covers it and this + # wraps a no-op. + @self.server.before_serving + async def _arm_stream_shutdown(): # pylint: disable=unused-variable + _streaming_shutdown.clear() + install_stream_shutdown_handler() + + # Close open downlink subscriptions and cancel stream pumps on shutdown, + # so a long-polling downlink can't block a graceful exit. + @self.server.after_serving + async def _shutdown_streams(): # pylint: disable=unused-variable + shutdown_active_streams() + + def _ndjson_response(marker): + # pylint: disable=protected-access + return Response( # type: ignore[return-value] + marker_ndjson_aiter( + marker, + keepalive_seconds(dash_app._stream_keepalive_interval), + ), + content_type=STREAM_MIMETYPE, + headers=dict(STREAM_HEADERS), + ) + + async def _dispatch(): # pylint: disable=too-many-return-statements adapter = QuartRequestAdapter() if "gzip" in adapter.request.headers.get("Content-Encoding", ""): body = decompress_payload(await adapter.request.get_data()) else: body = await adapter.get_json() + downlink = body.get("streamDownlink") + if downlink is not None: + connection_id = get_stream_connection_id() + if connection_id is None: + return Response(status=403) # type: ignore[return-value] + marker = async_downlink_marker( + dash_app.shared_storage, + connection_id, + downlink.get("from"), + ) + return _ndjson_response(marker) + cancel = body.get("streamCancel") + if cancel is not None: + # A tab closed while the shared downlink stays open for others: + # stop its pump. Same connection-id rule as the downlink: a page + # can only cancel its own streams. + connection_id = ( + get_stream_connection_id() + if dash_app.shared_storage_enabled + else None + ) + if connection_id is None: + return Response(status=403) # type: ignore[return-value] + cancel_stream( + dash_app.shared_storage, connection_id, cancel["requestId"] + ) + return jsonify(STREAM_CANCEL_ACK) # pylint: disable=protected-access cb_ctx = dash_app._initialize_context(body) # pylint: disable=protected-access @@ -404,6 +481,29 @@ async def _dispatch(): response_data = ctx.run(partial_func) if inspect.iscoroutine(response_data): # if user callback is async response_data = await response_data + if isinstance(response_data, StreamedCallbackResponse): + stream_conn = body.get("streamConnection") + if stream_conn is not None: + # A streamConnection asserts "multiplex me": the connection + # id is derived from the signed end_id (never from the + # client), so a page can only publish to its own topic. An + # unverified connection is refused, never run some other way, + # so no frame reaches a topic without a valid token. + connection_id = ( + get_stream_connection_id() + if dash_app.shared_storage_enabled + else None + ) + if connection_id is None: + return Response(status=403) # type: ignore[return-value] + spawn_async_pump( + dash_app.shared_storage, + connection_id, + stream_conn["requestId"], + response_data, + ) + return cb_ctx.dash_response.set_response(data=to_json(STREAM_ACK)) + return _ndjson_response(response_data) return cb_ctx.dash_response.set_response(data=response_data) # type: ignore[arg-type] # Preserve the view function's identity as `dash.dash.dispatch` so that @@ -656,6 +756,12 @@ async def websocket_handler(): # pylint: disable=too-many-branches renderer_id, connection_shutdown_event, ) + stream_emitter = make_stream_frame_emitter( + outbound_queue, + request_id, + renderer_id, + connection_shutdown_event, + ) if is_async: task = asyncio.create_task( @@ -664,6 +770,7 @@ async def websocket_handler(): # pylint: disable=too-many-branches payload, ws_cb, QuartResponseAdapter(), + stream_emitter, ) ) task.add_done_callback(done_handler) @@ -676,6 +783,7 @@ async def websocket_handler(): # pylint: disable=too-many-branches payload, ws_cb, QuartResponseAdapter(), + stream_emitter, ) # Set up done callback to send response future.add_done_callback(done_handler) diff --git a/dash/backends/base_server.py b/dash/backends/base_server.py index 5443662dd2..57034da23d 100644 --- a/dash/backends/base_server.py +++ b/dash/backends/base_server.py @@ -180,6 +180,10 @@ class BaseDashServer(ABC, Generic[ServerType]): request_adapter: Type[RequestAdapter] response_adapter: Type[ResponseAdapter] websocket_capability: bool = False + # How streaming callbacks' downlink is served: "poll" (WSGI: a request + # returns the queued frames and ends, since an open response would hold a + # worker thread) or "stream" (ASGI: one open connection per browser). + downlink_mode: str = "poll" def __init__(self, server: ServerType) -> None: """Initialize the server wrapper. diff --git a/dash/backends/ws.py b/dash/backends/ws.py index d83bde9ed2..2c72bcdf79 100644 --- a/dash/backends/ws.py +++ b/dash/backends/ws.py @@ -21,6 +21,12 @@ from dash.exceptions import PreventUpdate, WebsocketDisconnected from dash.types import CallbackExecutionBody +from dash._streaming import ( + StreamedCallbackResponse, + aiter_stream_frames, + iter_stream_frames, + sync_iter_asyncgen, +) from dash._utils import to_json if TYPE_CHECKING: @@ -542,6 +548,105 @@ def on_done(f: "concurrent.futures.Future | asyncio.Future") -> None: return on_done +def make_stream_frame_emitter( + outbound_queue: janus.Queue[str], + request_id: str, + renderer_id: str, + shutdown_event: threading.Event, +) -> Callable[[dict], None]: + """Create an emitter sending intermediate stream frames for a request. + + Frames ride the ``callback_response`` message type with ``stream: True`` + and no ``done`` flag, so the renderer keeps the request pending until the + final ``callback_response`` (sent by the done handler) resolves it. + """ + + def emit(frame: dict) -> None: + if shutdown_event.is_set(): + return + outbound_queue.sync_q.put_nowait( + cast( + str, + to_json( + { + "type": "callback_response", + "rendererId": renderer_id, + "requestId": request_id, + "payload": {"status": "ok", "stream": True, "data": frame}, + } + ), + ) + ) + # Flush so the frame doesn't sit in the sender's batching window. + outbound_queue.sync_q.put_nowait(FLUSH_SIGNAL) + + return emit + + +_STREAM_DONE_PAYLOAD = {"status": "ok", "stream": True, "done": True} + + +def _stream_frame_result(frame: dict) -> "dict | None": + """Terminal payload for a frame, or None if it is an intermediate frame.""" + if not frame.get("done"): + return None + error = frame.get("error") + if error: + return {"status": "error", "message": error.get("message", "")} + return dict(_STREAM_DONE_PAYLOAD) + + +def consume_stream_frames( + marker: StreamedCallbackResponse, + ws_callback: DashWebsocketCallback, + stream_emitter: "Callable[[dict], None] | None", +) -> dict: + """Drain a streamed callback synchronously (threadpool path). + + Emits intermediate frames over the WebSocket and returns the terminal + payload for the done handler to send as the final callback_response. + """ + emit = stream_emitter or (lambda _frame: None) + if marker.is_async: + # Defensive: an async generator that ended up on the thread path is + # driven on a private event-loop thread. + frames = sync_iter_asyncgen(marker.frames) + else: + frames = iter_stream_frames(marker) + try: + for frame in frames: + if ws_callback.is_shutdown: + return {"status": "prevent_update"} + result = _stream_frame_result(frame) + if result is not None: + return result + emit(frame) + return dict(_STREAM_DONE_PAYLOAD) + finally: + frames.close() + + +async def aconsume_stream_frames( + marker: StreamedCallbackResponse, + ws_callback: DashWebsocketCallback, + stream_emitter: "Callable[[dict], None] | None", +) -> dict: + """Drain a streamed callback on the event loop (async dispatch path).""" + emit = stream_emitter or (lambda _frame: None) + frames = marker.frames if marker.is_async else aiter_stream_frames(marker) + try: + async for frame in frames: + if ws_callback.is_shutdown: + return {"status": "prevent_update"} + result = _stream_frame_result(frame) + if result is not None: + return result + emit(frame) + return dict(_STREAM_DONE_PAYLOAD) + finally: + await frames.aclose() + + def _prepare_ws_partial( dash_app: "dash.Dash", payload: CallbackExecutionBody, @@ -569,6 +674,7 @@ def run_callback_in_executor( payload: CallbackExecutionBody, ws_callback: DashWebsocketCallback, response_adapter: "ResponseAdapter", + stream_emitter: "Callable[[dict], None] | None" = None, ) -> concurrent.futures.Future: """Submit a synchronous callback to the executor for thread pool execution. @@ -607,6 +713,10 @@ def run_callback(): return result response_data = ctx.run(run_callback) + if isinstance(response_data, StreamedCallbackResponse): + # The frame generator carries its own context snapshot, so it + # can be driven outside ctx here. + return consume_stream_frames(response_data, ws_callback, stream_emitter) return {"status": "ok", "data": json.loads(response_data)} except PreventUpdate: @@ -625,6 +735,7 @@ async def run_callback_on_loop( payload: CallbackExecutionBody, ws_callback: DashWebsocketCallback, response_adapter: "ResponseAdapter", + stream_emitter: "Callable[[dict], None] | None" = None, ) -> dict: """Run an async callback as a task on the connection's event loop. @@ -652,6 +763,10 @@ async def run_callback_on_loop( ) result = partial_func() response_data = await result if inspect.iscoroutine(result) else result + if isinstance(response_data, StreamedCallbackResponse): + return await aconsume_stream_frames( + response_data, ws_callback, stream_emitter + ) return {"status": "ok", "data": json.loads(response_data)} except PreventUpdate: diff --git a/dash/dash-renderer/init.template b/dash/dash-renderer/init.template index c50e0bc7c9..2f2fe70553 100644 --- a/dash/dash-renderer/init.template +++ b/dash/dash-renderer/init.template @@ -100,4 +100,9 @@ _js_dist = [ "namespace": "dash", "dynamic": True, }, + { + "relative_package_path": "dash-renderer/build/dash-stream-worker.js", + "namespace": "dash", + "dynamic": True, + }, ] diff --git a/dash/dash-renderer/src/actions/callbacks.ts b/dash/dash-renderer/src/actions/callbacks.ts index 8d065de74c..83d3fe5e93 100644 --- a/dash/dash-renderer/src/actions/callbacks.ts +++ b/dash/dash-renderer/src/actions/callbacks.ts @@ -49,6 +49,7 @@ import {computePaths, getPath} from './paths'; import {requestDependencies} from './requestDependencies'; import {loadLibrary} from '../utils/libraries'; +import {getStreamClient, isStreamMultiplexed} from '../utils/streamClient'; import {parsePMCId} from './patternMatching'; import {replacePMC} from './patternMatching'; @@ -503,6 +504,80 @@ function sideUpdate(outputs: SideUpdateOutput, cb: ICallbackPayload) { }; } +/** + * Apply an intermediate frame from a streaming callback. + * + * The frame's declared outputs are flattened to `id.prop` keys and applied + * through the sideUpdate path (parsePatchProps + updateComponent), so Patch + * values apply incrementally, paths recompute for newly added components, + * and observers are notified — the same semantics as set_props. The + * callback's execution promise stays pending until the terminal frame. + */ +function applyStreamFrame( + dispatch: any, + frame: CallbackResponseData, + payload: ICallbackPayload +) { + if (frame.response) { + const flat: SideUpdateOutput = {}; + toPairs(frame.response).forEach(([id, props]) => { + toPairs(props as Record).forEach(([prop, value]) => { + flat[`${id}.${prop}`] = value; + }); + }); + if (keys(flat).length) { + dispatch(sideUpdate(flat, payload)); + } + } + if (frame.sideUpdate) { + dispatch(sideUpdate(frame.sideUpdate, payload)); + } +} + +/** + * Run a streaming callback over the multiplexed transport: a single downlink + * shared by all of the page's streams -- and, hosted in a SharedWorker, by all + * of the browser's tabs (see utils/streamClient). Frames are + * applied as they arrive, so the resolved value is empty like the WebSocket + * streaming path. + */ +async function handleStreamCallback( + dispatch: any, + config: any, + payload: ICallbackPayload, + running: any +): Promise { + let runningOff: any; + if (running) { + dispatch(sideUpdate(running.running, payload)); + runningOff = running.runningOff; + } + const url = `${urlBase(config)}_dash-update-component`; + const init = mergeDeepRight(config.fetch, { + headers: getCSRFHeader(config) as any + }); + try { + await getStreamClient(config).run( + url, + init, + config.end_id, + payload, + (frame: any) => { + if (frame.dist) { + Promise.all(frame.dist.map(loadLibrary)); + } + applyStreamFrame(dispatch, frame, payload); + } + ); + } finally { + if (runningOff) { + dispatch(sideUpdate(runningOff, payload)); + } + } + // Frames were applied as they arrived; the terminal store update is a no-op. + return {}; +} + function handleServerside( dispatch: any, hooks: any, @@ -683,7 +758,97 @@ function handleServerside( } }; + // Streaming callbacks: read NDJSON frames as they arrive, + // apply each one immediately, and resolve on the terminal frame. + const handleStreamedResponse = async (streamRes: any) => { + const reader = streamRes.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + let lastResponse: CallbackResponse | undefined; + let finished = false; + + const processFrame = async (frame: CallbackResponseData) => { + if (frame.done) { + finished = true; + completeJob(); + if (frame.error) { + recordProfile({}); + reject( + new Error( + frame.error.message || 'Callback error' + ) + ); + return; + } + if (hooks.request_post) { + hooks.request_post(payload, lastResponse); + } + recordProfile(lastResponse || {}); + // Frames were already applied; resolve empty so the + // terminal store update is a no-op (Patch frames must + // not re-apply). + resolve({}); + return; + } + if (frame.dist) { + await Promise.all(frame.dist.map(loadLibrary)); + } + lastResponse = frame.response; + applyStreamFrame(dispatch, frame, payload); + }; + + try { + for (;;) { + const {done, value} = await reader.read(); + if (done) { + break; + } + buffer += decoder.decode(value, {stream: true}); + const lines = buffer.split('\n'); + buffer = lines.pop() as string; + for (const line of lines) { + if (!line.trim()) { + continue; + } + await processFrame(JSON.parse(line)); + if (finished) { + reader.cancel(); + return; + } + } + } + if (buffer.trim() && !finished) { + await processFrame(JSON.parse(buffer)); + } + } catch (err) { + if (!finished) { + finished = true; + completeJob(); + if (lastResponse) { + recordProfile(lastResponse); + resolve({}); + } else { + recordProfile({}); + reject(err); + } + return; + } + } + if (!finished) { + // Stream ended without a terminal frame (connection + // dropped): clear loading, keep applied frames. + completeJob(); + recordProfile(lastResponse || {}); + resolve({}); + } + }; + if (status === STATUS.OK) { + const contentType = res.headers.get('Content-Type') || ''; + if (contentType.includes('application/x-ndjson') && res.body) { + handleStreamedResponse(res); + return; + } res.json().then((data: CallbackResponseData) => { if (!cacheKey && data.cacheKey) { cacheKey = data.cacheKey; @@ -795,7 +960,16 @@ async function handleWebsocketCallback( // Ensure WebSocket connection is established await workerClient.ensureConnected(config); - const response = await workerClient.sendCallback(payload); + // Streaming callbacks deliver intermediate frames before the + // terminal response; apply each one as it arrives. + let lastStreamResponse: CallbackResponse | undefined; + const response = await workerClient.sendCallback( + payload, + (frame: CallbackResponseData) => { + lastStreamResponse = frame?.response; + applyStreamFrame(dispatch, frame, payload); + } + ); // Handle running off state if (runningOff) { @@ -829,6 +1003,34 @@ async function handleWebsocketCallback( throw new Error(response.message || 'Callback error'); } + if (response.stream) { + // Terminal frame of a streamed callback: every frame was already + // applied on arrival, so resolve empty (the terminal store update + // must be a no-op — Patch frames must not re-apply). + if (hooks.request_post) { + hooks.request_post(payload, lastStreamResponse); + } + if (config.ui) { + const totalTime = Date.now() - requestTime; + dispatch( + updateResourceUsage({ + id: payload.output, + usage: { + __dash_server: totalTime, + __dash_client: totalTime, + __dash_upload: 0, + __dash_download: 0 + }, + status: STATUS.OK, + result: lastStreamResponse || {}, + inputs: payload.inputs, + state: payload.state + }) + ); + } + return {}; + } + // Extract the callback data - structure is {multi: boolean, response: {...}} const callbackData = response.data as CallbackResponseData; @@ -1101,6 +1303,15 @@ export function executeCallback( (cb.callback.websocket && isWebSocketAvailable(config))); + // Streaming callbacks ride the single multiplexed downlink when + // the server offers it (shared storage enabled) and they are not + // already on the WebSocket transport or a background job. + const useStream = + !background && + !useWebSocket && + cb.callback.stream && + isStreamMultiplexed(config); + for (let retry = 0; retry <= MAX_AUTH_RETRIES; retry++) { try { let data: CallbackResponse; @@ -1114,6 +1325,34 @@ export function executeCallback( payload, cb.callback.running ); + } else if (useStream) { + try { + data = await handleStreamCallback( + dispatch, + newConfig, + payload, + cb.callback.running + ); + } catch (streamErr: any) { + if (streamErr?.message?.includes('403')) { + data = await handleServerside( + dispatch, + hooks, + newConfig, + payload, + background, + additionalArgs.length + ? additionalArgs + : undefined, + getState, + cb.callback.running, + cb.callback.compress_payload, + cb.callback.compress_threshold + ); + } else { + throw streamErr; + } + } } else { // Use traditional HTTP path data = await handleServerside( diff --git a/dash/dash-renderer/src/config.ts b/dash/dash-renderer/src/config.ts index 42473a4a55..09477715a6 100644 --- a/dash/dash-renderer/src/config.ts +++ b/dash/dash-renderer/src/config.ts @@ -29,6 +29,16 @@ export type DashConfig = { inactivity_timeout?: number; heartbeat_interval?: number; }; + stream?: { + enabled: boolean; + // Served when enabled: the SharedWorker that hosts the browser's + // single streaming downlink, shared across tabs. + worker_url?: string; + // How the server serves the downlink: one open connection (ASGI) or + // polling (WSGI), and the poll interval in ms for the latter. + mode?: 'stream' | 'poll'; + poll_interval?: number; + }; csrf_token_name?: string; csrf_header_name?: string; // Server-issued, server-signed token for this page load. Echoed on every diff --git a/dash/dash-renderer/src/observers/prioritizedCallbacks.ts b/dash/dash-renderer/src/observers/prioritizedCallbacks.ts index 024df36c12..ce5f907ad8 100644 --- a/dash/dash-renderer/src/observers/prioritizedCallbacks.ts +++ b/dash/dash-renderer/src/observers/prioritizedCallbacks.ts @@ -17,6 +17,8 @@ import {combineIdAndProp} from '../actions/dependencies_ts'; import isAppReady from '../actions/isAppReady'; +import {MAX_CONCURRENT_HTTP_CALLBACKS, usesRequestSlot} from './requestSlot'; + import { IBlockedCallback, ICallback, @@ -82,23 +84,38 @@ const observer: IStoreObserverDefinition = { return; } - const available = Math.max(0, 12 - executing.length - watched.length); + // Only callbacks holding an in-flight HTTP request count toward the + // budget; clientside, streaming and websocket-routed ones are exempt. + const countsToward = (cb: ICallback) => usesRequestSlot(cb, config); + + const inFlight = + executing.filter(countsToward).length + + watched.filter(countsToward).length; + const available = Math.max(0, MAX_CONCURRENT_HTTP_CALLBACKS - inFlight); // Order prioritized callbacks based on depth and breadth of callback chain prioritized = sort(sortPriority, prioritized); - // Divide between sync and async - const [syncCallbacks, asyncCallbacks] = partition( - cb => isAppReady(layout, paths, getIds(cb, paths)) === true, - prioritized - ); + // Exempt callbacks always dispatch; only request-slot callbacks are + // limited to the available budget (ready ones first, as before). + const [budgeted, exempt] = partition(countsToward, prioritized); - const pickedSyncCallbacks = syncCallbacks.slice(0, available); - const pickedAsyncCallbacks = asyncCallbacks.slice( + const isReady = (cb: ICallback) => + isAppReady(layout, paths, getIds(cb, paths)) === true; + + const [budgetedSync, budgetedAsync] = partition(isReady, budgeted); + const pickedBudgetedSync = budgetedSync.slice(0, available); + const pickedBudgetedAsync = budgetedAsync.slice( 0, - available - pickedSyncCallbacks.length + available - pickedBudgetedSync.length ); + const [exemptSync, exemptAsync] = partition(isReady, exempt); + + // Divide between sync (components ready) and async (deferred until ready) + const pickedSyncCallbacks = [...exemptSync, ...pickedBudgetedSync]; + const pickedAsyncCallbacks = [...exemptAsync, ...pickedBudgetedAsync]; + if (pickedSyncCallbacks.length) { dispatch( aggregateCallbacks([ diff --git a/dash/dash-renderer/src/observers/requestSlot.ts b/dash/dash-renderer/src/observers/requestSlot.ts new file mode 100644 index 0000000000..ce811ee2fd --- /dev/null +++ b/dash/dash-renderer/src/observers/requestSlot.ts @@ -0,0 +1,36 @@ +import {isWebSocketAvailable, isWebSocketEnabled} from '../utils/workerClient'; + +import type {DashConfig} from '../config'; +import type {ICallback} from '../types/callbacks'; + +// Cap on how many callbacks may have an in-flight HTTP request to the server at +// once. This is NOT a limit on total callbacks -- it only throttles the fan-out +// of concurrent HTTP requests so the browser's ~6-connections-per-host ceiling +// doesn't stall the app under a wide callback graph. Callbacks that don't hold +// an HTTP connection for their lifetime are exempt (see `usesRequestSlot`): +// clientside callbacks (run in-browser), streaming callbacks (long-lived), and +// anything routed over the multiplexed WebSocket transport. +export const MAX_CONCURRENT_HTTP_CALLBACKS = 12; + +// A callback rides the multiplexed WebSocket transport (rather than its own HTTP +// request) when websocket callbacks are enabled globally, or when it opts in +// per-callback and the transport is available. Never for background callbacks, +// which always poll over HTTP. Mirrors the routing decision in handleServerside. +export const routedOverWebSocket = ( + cb: ICallback, + config: DashConfig +): boolean => + !cb.callback.background && + (isWebSocketEnabled(config) || + (Boolean(cb.callback.websocket) && isWebSocketAvailable(config))); + +// True only for callbacks that hold an HTTP connection for their lifetime -- the +// only ones that count against MAX_CONCURRENT_HTTP_CALLBACKS. Clientside +// callbacks make no request, streaming callbacks are long-lived (they must not +// pin a slot for their whole life -- that would starve everything else, +// including clientside callbacks), and websocket-routed callbacks share one +// socket, so none of those count. +export const usesRequestSlot = (cb: ICallback, config: DashConfig): boolean => + !cb.callback.clientside_function && + !cb.callback.stream && + !routedOverWebSocket(cb, config); diff --git a/dash/dash-renderer/src/types/callbacks.ts b/dash/dash-renderer/src/types/callbacks.ts index 0c7c2d63fd..bfe38422c0 100644 --- a/dash/dash-renderer/src/types/callbacks.ts +++ b/dash/dash-renderer/src/types/callbacks.ts @@ -21,6 +21,7 @@ export interface ICallbackDefinition { persistent?: boolean; compress_payload?: boolean; compress_threshold?: number; + stream?: boolean; } export interface ICallbackProperty { @@ -118,6 +119,9 @@ export type CallbackResponseData = { cancel?: ICallbackProperty[]; dist?: any; sideUpdate?: any; + // Streaming callbacks: terminal frame marker and mid-stream error. + done?: boolean; + error?: {message?: string}; }; export type SideUpdateOutput = { diff --git a/dash/dash-renderer/src/utils/streamClient.ts b/dash/dash-renderer/src/utils/streamClient.ts new file mode 100644 index 0000000000..5cc6892e1c --- /dev/null +++ b/dash/dash-renderer/src/utils/streamClient.ts @@ -0,0 +1,744 @@ +/** + * Multiplexed streaming transport for streaming callbacks. + * + * Instead of one long-lived NDJSON connection per streaming callback (which hits + * the browser's ~6-connections-per-host ceiling), every streaming callback shares + * ONE downlink connection. A callback POSTs its request (which returns a fast ack) + * carrying a request id; the server pumps that callback's frames onto the + * connection's shared-storage topic; the single downlink relays them and the + * client routes each frame back to the right callback by request id. + * + * The connection is keyed server-side on the page's signed endId (sent as a query + * parameter on every uplink, downlink and cancel), never on anything the client + * picks, so a page can only ever read or write its own topic. + * + * Two hosts run this transport: + * + * - `StreamClient` owns the HTTP side: the uplink POSTs and the downlink read + * loop. It runs inside the page when no worker is available. + * - `SharedStreamClient` is the page's proxy to a `StreamClient` hosted in a + * SharedWorker (`workers/streamWorker.ts`), so every tab of the browser shares + * one downlink -- the per-host connection cap is shared across tabs, and a + * downlink per tab stalls the sixth tab. The worker pins the endId of the + * first tab that streams for as long as the connection has streams in + * flight, keeps the downlink open while any tab has a stream, and cancels a + * tab's streams server-side when that tab goes away. + * + * Two downlink modes, chosen by the server (config.stream.mode): + * + * - 'stream' (ASGI): one long-lived NDJSON response; frames arrive as produced. + * - 'poll' (WSGI): each downlink request returns the frames queued since the + * cursor and ends at once, so it holds a server worker thread for + * milliseconds rather than for the whole visit. The client re-polls after + * `pollInterval` while frames keep coming, backs off to ten times that while + * quiet, and polls immediately when a new stream starts. + * + * Lifecycle: the downlink opens once the first uplink is acknowledged and closes + * when no acknowledged callback remains in flight ("collect the dones to match + * the runnings"). Waiting for the ack lets a single request slot serve the two + * in turn; frames published before the downlink subscribes are replayed from + * the cursor. If the downlink drops while callbacks are still running it + * reconnects, resuming from the last sequence it saw, backing off until the + * server has been unreachable for the whole reconnect window. + */ + +import {getRendererId} from './rendererId'; + +type Frame = Record; + +/** How the server serves the downlink; see the module comment. */ +export interface StreamTransportOptions { + mode?: 'stream' | 'poll'; + /** Poll mode: delay between polls while frames flow, in ms (default 100). */ + pollInterval?: number; +} + +const DEFAULT_POLL_INTERVAL = 100; +// Idle polls stay at the interval for this many empty polls (a callback's +// next frame is usually moments away), then back off geometrically up to +// MAX_BACKOFF_FACTOR times the interval. Two keeps a stream that yields every +// half second at about four polls per frame instead of ten; the cap bounds +// the latency of a slow stream's next frame so consecutive frames still +// render one at a time instead of bunching into one poll. +const EMPTY_POLLS_BEFORE_BACKOFF = 2; +const MAX_BACKOFF_FACTOR = 5; + +/** What the callbacks action needs from either host. */ +export interface StreamTransport { + /** + * Run one streaming callback. Resolves when the callback's terminal `done` + * frame arrives (its output frames having been delivered to `onFrame` as + * they arrive), or rejects on error. + */ + run( + url: string, + init: RequestInit, + endId: string, + payload: Record, + onFrame: (frame: Frame) => void + ): Promise; +} + +interface PendingStream { + onFrame: (frame: Frame) => void; + resolve: () => void; + reject: (err: Error) => void; + // Whether the uplink POST was acknowledged; only acknowledged callbacks + // keep the downlink open (see the lifecycle note above). + acked: boolean; + // Whether any output frame reached this callback. Decides how a lost + // connection settles it: frames applied -> resolve and keep them (like the + // single-connection NDJSON path does on a drop); nothing applied -> reject, + // so the caller can report it or fall back. + gotFrame: boolean; +} + +interface DownlinkEnvelope { + rid?: string; + frame?: Frame; + seq?: number; + // Set by the server when this connection's buffered frames were lost (its + // owner was re-elected, or the server restarted): the client must reset its + // cursor to the head rather than keep asking to resume from a stale one. + reset?: boolean; +} + +type FetchImpl = typeof fetch; + +const genId = (): string => + `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 10)}`; + +const sleep = (ms: number): Promise => + new Promise(resolve => setTimeout(resolve, ms)); + +export class StreamClient implements StreamTransport { + // Local id, used only to make request ids unique on this host. The server + // does NOT key the topic on it: the connection is keyed on the signed endId + // instead, so a client can't name another page's topic. See streamUrl. + private localId = genId(); + // The endId this connection is keyed on. Pinned by the first stream and + // kept while any stream is in flight, so every tab behind a shared worker + // publishes to and reads from the same topic. + private endId = ''; + private pending = new Map(); + private counter = 0; + // Last sequence applied; the downlink resumes from here on reconnect. Starts + // at 0 so the first connect replays anything published before it subscribed + // (the uplink POST and the downlink open race). + private cursor = 0; + private downlinkOpen = false; + private abort: AbortController | null = null; + // Bumped whenever a read loop starts or the downlink is closed, so a + // retired loop (closed while it was mid-await) notices and exits instead + // of fighting a newer loop for the connection. + private loopGen = 0; + private reconnectDelay: number; + private maxReconnectDelay: number; + private reconnectWindow: number; + private fetchImpl: FetchImpl; + private _transport: StreamTransportOptions = {}; + // Resolves the current idle pause early (a new stream started). + private wakeUp: (() => void) | null = null; + // Set by start(): the read loop resets its idle backoff on the next turn. + private wokenForNewStream = false; + + constructor( + opts: { + fetchImpl?: FetchImpl; + // First retry delay after a downlink drop; doubles up to + // maxReconnectDelay on each further failure. + reconnectDelay?: number; + maxReconnectDelay?: number; + // How long the downlink may stay unreachable before the callbacks + // waiting on it are settled as lost instead of retrying forever. + reconnectWindow?: number; + } & StreamTransportOptions = {} + ) { + // Native fetch must be invoked with `this === window`; calling it as a + // method of this object throws "Illegal invocation", so bind it. + // globalThis is window on a page and self in a worker. + this.fetchImpl = opts.fetchImpl ?? fetch.bind(globalThis); + this.reconnectDelay = opts.reconnectDelay ?? 1000; + this.maxReconnectDelay = opts.maxReconnectDelay ?? 5000; + this.reconnectWindow = opts.reconnectWindow ?? 30000; + this.configure(opts); + } + + /** Adopt the server's downlink mode (the worker host learns it from the page). */ + configure(transport: StreamTransportOptions): void { + if (transport.mode) { + this._transport.mode = transport.mode; + } + if (typeof transport.pollInterval === 'number') { + this._transport.pollInterval = transport.pollInterval; + } + } + + get transport(): StreamTransportOptions { + return {...this._transport}; + } + + get activeCount(): number { + return this.pending.size; + } + + /** The endId this connection is currently keyed on (empty when idle). */ + get connectionEndId(): string { + return this.endId; + } + + /** Append the signed endId so the server can derive (and authorize) the + * connection topic. The server never trusts a client-supplied topic id. */ + private streamUrl(url: string): string { + if (!this.endId) { + return url; + } + const delim = url.includes('?') ? '&' : '?'; + return `${url}${delim}endId=${encodeURIComponent(this.endId)}`; + } + + run( + url: string, + init: RequestInit, + endId: string, + payload: Record, + onFrame: (frame: Frame) => void + ): Promise { + return this.start(url, init, endId, payload, onFrame).settled; + } + + /** + * `run`, also handing back the request id so the caller can `cancel` the + * stream later (the worker host does, for a tab that went away). + */ + start( + url: string, + init: RequestInit, + endId: string, + payload: Record, + onFrame: (frame: Frame) => void + ): {requestId: string; settled: Promise} { + if (this.pending.size === 0 || !this.endId) { + // Nothing in flight: this stream's page keys the connection. While + // streams are in flight the key stays put, so a stream from + // another tab (shared worker) lands on the same topic. + this.endId = endId || ''; + } + const requestId = `${this.localId}-${++this.counter}`; + const settled = new Promise((resolve, reject) => { + this.pending.set(requestId, { + onFrame, + resolve, + reject, + acked: false, + gotFrame: false + }); + }); + // Uplink POST: returns a fast ack; the outputs arrive on the downlink, + // which opens once the ack is in. + this.fetchImpl(this.streamUrl(url), { + ...init, + method: 'POST', + body: JSON.stringify({ + ...payload, + streamConnection: {requestId} + }) + }) + .then(res => this.checkUplink(requestId, res, url, init)) + .catch(err => this.fail(requestId, err)); + return {requestId, settled}; + } + + /** + * Abandon one in-flight callback whose consumer is gone: stop waiting for + * its frames (they are dropped on arrival) and tell the server to cancel + * it, so it does not run to completion for nobody. Its promise rejects. + */ + cancel(requestId: string, url: string, init: RequestInit): void { + const pending = this.pending.get(requestId); + if (!pending) { + return; + } + this.pending.delete(requestId); + pending.reject(new Error('Streaming callback cancelled')); + this.fetchImpl(this.streamUrl(url), { + ...init, + method: 'POST', + body: JSON.stringify({streamCancel: {requestId}}) + }).catch(() => undefined); + this.stopDownlinkIfIdle(); + } + + /** + * The uplink returns a fast ack (200); the frames then arrive on the + * downlink, so make sure one is open. Any non-ok status (e.g. 403 when the + * connection did not verify) means no frames are coming, so fail the + * request loudly rather than leave the callback pending forever. + */ + private checkUplink( + requestId: string, + res: Response, + url: string, + init: RequestInit + ): void { + if (!res.ok) { + this.fail( + requestId, + new Error(`stream uplink responded ${res.status}`) + ); + return; + } + const pending = this.pending.get(requestId); + if (pending) { + pending.acked = true; + this.ensureDownlink(url, init); + // A polling downlink may be pausing between polls: this stream's + // first frame should not wait out that pause. + this.wokenForNewStream = true; + this.wakeUp?.(); + } + } + + /** Whether any acknowledged callback is still waiting on the downlink. */ + private hasAcked(): boolean { + for (const p of this.pending.values()) { + if (p.acked) { + return true; + } + } + return false; + } + + /** Route one downlink envelope to its callback. Public for testing. */ + dispatchEnvelope(envelope: DownlinkEnvelope): void { + if (envelope.reset) { + // Our cursor points into a server incarnation that lost our frames + // (it restarted, or the storage owner changed). Reset to the head so + // a later downlink starts from what the fresh topic actually has, + // and settle the callbacks in flight: the frames they were waiting + // on are gone, and after a restart nothing will ever finish them. + this.cursor = 0; + this.settleAll( + new Error( + 'stream reset: the server lost this connection (restart or owner change)' + ) + ); + return; + } + if (typeof envelope.seq === 'number') { + this.cursor = envelope.seq; + } + if (envelope.rid === undefined || envelope.frame === undefined) { + return; + } + const pending = this.pending.get(envelope.rid); + if (!pending) { + // A frame for a callback we already resolved (e.g. a replayed + // duplicate after reconnect) -- safe to drop. + return; + } + const {frame} = envelope; + if (frame.done) { + this.pending.delete(envelope.rid); + if (frame.error) { + pending.reject( + new Error(frame.error.message || 'Streaming callback error') + ); + } else { + pending.resolve(); + } + this.stopDownlinkIfIdle(); + } else { + pending.gotFrame = true; + pending.onFrame(frame); + } + } + + /** + * The downlink is gone for good (refused by the server, reset, or + * unreachable past the reconnect window). Callbacks that already applied + * frames resolve so those frames stay on the page; ones that never got a + * frame reject with `err`, and the read loop winds down since nothing is + * pending any more. + */ + private settleAll(err: Error): void { + const pending = Array.from(this.pending.values()); + this.pending.clear(); + // Close before settling: a continuation of a settled promise may start + // a new stream right away, and it must get a fresh downlink rather + // than find this one still marked open. + this.closeDownlink(); + pending.forEach(p => (p.gotFrame ? p.resolve() : p.reject(err))); + } + + private fail(requestId: string, err: Error): void { + const pending = this.pending.get(requestId); + if (pending) { + this.pending.delete(requestId); + pending.reject(err); + this.stopDownlinkIfIdle(); + } + } + + private stopDownlinkIfIdle(): void { + if (this.downlinkOpen && !this.hasAcked()) { + this.closeDownlink(); + } + } + + /** Retire the current read loop and end its connection. */ + private closeDownlink(): void { + const abort = this.abort; + this.abort = null; + this.downlinkOpen = false; + this.loopGen++; + this.wakeUp?.(); + if (abort) { + abort.abort(); + } + } + + private ensureDownlink(url: string, init: RequestInit): void { + if (this.downlinkOpen) { + return; + } + this.downlinkOpen = true; + // Fire-and-forget read loop; it exits when no callbacks remain. + this.readLoop(url, init, ++this.loopGen); + } + + private async readLoop( + url: string, + init: RequestInit, + gen: number + ): Promise { + const polling = this._transport.mode === 'poll'; + const base = this._transport.pollInterval ?? DEFAULT_POLL_INTERVAL; + let idlePause = base; + let emptyPolls = 0; + let unreachableSince: number | null = null; + let delay = this.reconnectDelay; + while (this.hasAcked() && this.loopGen === gen) { + if (this.wokenForNewStream) { + this.wokenForNewStream = false; + emptyPolls = 0; + idlePause = base; + } + this.abort = new AbortController(); + let received = 0; + try { + const res = await this.fetchImpl(this.streamUrl(url), { + ...init, + method: 'POST', + signal: this.abort.signal, + body: JSON.stringify({ + streamDownlink: {from: this.cursor} + }) + }); + if (res.status >= 400 && res.status < 500) { + // The server refuses this connection outright, typically + // 403 after a restart minted a new signing secret so our + // endId no longer verifies. Retrying cannot fix that. + this.settleAll( + new Error(`stream downlink responded ${res.status}`) + ); + break; + } + if (!res.ok || !res.body) { + throw new Error(`downlink responded ${res.status}`); + } + received = await this.consume(res.body); + if (received === 0 && !polling) { + // Accepted then closed without a single envelope (a server + // mid-shutdown, a proxy dropping idle connections): back + // off like a failure instead of reconnecting in a burst. + throw new Error('downlink closed without data'); + } + // A productive connection ended (proxy timeout, worker + // recycle) or a poll completed: fresh backoff. + unreachableSince = null; + delay = this.reconnectDelay; + } catch (err) { + if (!this.hasAcked() || this.loopGen !== gen) { + break; // closed on purpose: idle, or settled and retired + } + // Genuine drop with work outstanding: reconnect from the + // cursor, backing off, until the server has been unreachable + // for the whole window. Past that the callbacks are lost. + const now = Date.now(); + unreachableSince = unreachableSince ?? now; + if (now - unreachableSince >= this.reconnectWindow) { + this.settleAll( + new Error( + 'stream downlink lost: could not reconnect to the server' + ) + ); + break; + } + await sleep(delay); + delay = Math.min(delay * 2, this.maxReconnectDelay); + continue; + } + if (polling && this.hasAcked() && this.loopGen === gen) { + // The response was one poll. Frames came: poll again soon. + // Nothing came: keep the pace for a few polls, then back off + // up to a bound, until frames or a new stream wake us. + if (received) { + emptyPolls = 0; + idlePause = base; + } else if (++emptyPolls > EMPTY_POLLS_BEFORE_BACKOFF) { + idlePause = Math.min( + idlePause * 2, + base * MAX_BACKOFF_FACTOR + ); + } + await this.pause(idlePause); + } + } + if (this.loopGen === gen) { + this.downlinkOpen = false; + this.abort = null; + } + } + + /** Sleep for `ms`, or until a new stream starts or the downlink closes. */ + private pause(ms: number): Promise { + return new Promise(resolve => { + const timer = setTimeout(() => { + this.wakeUp = null; + resolve(); + }, ms); + this.wakeUp = () => { + clearTimeout(timer); + this.wakeUp = null; + resolve(); + }; + }); + } + + /** Relay envelopes until the connection ends; returns how many arrived. */ + private async consume(body: ReadableStream): Promise { + const reader = body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + let received = 0; + for (;;) { + const {done, value} = await reader.read(); + if (done) { + return received; // connection ended -> the read loop decides + } + buffer += decoder.decode(value, {stream: true}); + let nl: number; + while ((nl = buffer.indexOf('\n')) >= 0) { + const line = buffer.slice(0, nl); + buffer = buffer.slice(nl + 1); + if (!line.trim()) { + continue; // keepalive blank line + } + received++; + this.dispatchEnvelope(JSON.parse(line)); + } + } + } +} + +// --- SharedWorker host ------------------------------------------------------ +// +// Messages page -> worker: +// {type: 'stream', rendererId, requestId, url, init, endId, payload, transport} +// {type: 'unregister', rendererId} this tab is going away +// Messages worker -> page: +// {type: 'frame', requestId, frame} +// {type: 'done', requestId} +// {type: 'error', requestId, message} + +export type StreamHostMessage = + | { + type: 'stream'; + rendererId: string; + requestId: string; + url: string; + init: RequestInit; + // The page's signed endId; the first tab's keys the connection. + endId: string; + payload: Record; + // The server's downlink mode, from the page's config: the worker + // has no config of its own. + transport?: StreamTransportOptions; + } + | {type: 'unregister'; rendererId: string}; + +export type StreamPortMessage = + | {type: 'frame'; requestId: string; frame: Frame} + | {type: 'done'; requestId: string} + | {type: 'error'; requestId: string; message: string}; + +/** The slice of MessagePort both sides use (fakeable in tests). */ +export interface StreamPort { + postMessage(message: any): void; + onmessage: ((event: MessageEvent) => void) | null; + start?(): void; +} + +/** + * The page's end of the SharedWorker transport: forwards each stream to the + * worker and routes its frames back to the callback that started it. + */ +export class SharedStreamClient implements StreamTransport { + private pending = new Map(); + private counter = 0; + private port: StreamPort; + private rendererId: string; + private transport: StreamTransportOptions; + private _broken = false; + + constructor( + port: StreamPort, + rendererId: string = getRendererId(), + transport: StreamTransportOptions = {} + ) { + this.port = port; + this.rendererId = rendererId; + this.transport = transport; + port.onmessage = event => this.handleMessage(event.data); + port.start?.(); + } + + /** True once the worker failed; callers should fall back to `StreamClient`. */ + get broken(): boolean { + return this._broken; + } + + get activeCount(): number { + return this.pending.size; + } + + run( + url: string, + init: RequestInit, + endId: string, + payload: Record, + onFrame: (frame: Frame) => void + ): Promise { + const requestId = `${this.rendererId}-${++this.counter}`; + return new Promise((resolve, reject) => { + this.pending.set(requestId, { + onFrame, + resolve, + reject, + acked: false, + gotFrame: false + }); + const message: StreamHostMessage = { + type: 'stream', + rendererId: this.rendererId, + requestId, + url, + init, + endId: endId || '', + payload, + transport: this.transport + }; + this.port.postMessage(message); + }); + } + + /** + * This tab is going away (pagehide): its streams have no consumer any + * more, so the worker cancels them server-side. + */ + release(): void { + const message: StreamHostMessage = { + type: 'unregister', + rendererId: this.rendererId + }; + this.port.postMessage(message); + } + + /** The worker died or failed to load: reject everything in flight. */ + fail(err: Error): void { + this._broken = true; + for (const pending of this.pending.values()) { + pending.reject(err); + } + this.pending.clear(); + } + + private handleMessage(message: StreamPortMessage): void { + const pending = this.pending.get(message.requestId); + if (!pending) { + return; + } + if (message.type === 'frame') { + pending.onFrame(message.frame); + return; + } + this.pending.delete(message.requestId); + if (message.type === 'done') { + pending.resolve(); + } else { + pending.reject(new Error(message.message)); + } + } +} + +let singleton: StreamTransport | null = null; + +/** + * The page's streaming transport: the SharedWorker-hosted downlink when the + * server provides the worker script and the browser supports SharedWorker, + * otherwise a downlink of the page's own. + */ +export function getStreamClient( + config: { + stream?: { + worker_url?: string; + mode?: 'stream' | 'poll'; + poll_interval?: number; + }; + } = {} +): StreamTransport { + if (singleton instanceof SharedStreamClient && singleton.broken) { + singleton = null; + } + if (singleton) { + return singleton; + } + const transport: StreamTransportOptions = { + mode: config.stream?.mode, + pollInterval: config.stream?.poll_interval + }; + const workerUrl = config.stream?.worker_url; + if (workerUrl && typeof SharedWorker !== 'undefined') { + try { + const worker = new SharedWorker(workerUrl, { + name: 'dash-stream-worker' + }); + const client = new SharedStreamClient( + worker.port, + getRendererId(), + transport + ); + worker.onerror = () => + client.fail(new Error('Dash stream worker failed')); + window.addEventListener('pagehide', () => client.release()); + singleton = client; + return singleton; + } catch (err) { + // Fall through to the in-page transport. + } + } + singleton = new StreamClient(transport); + return singleton; +} + +/** + * Whether the server offers the multiplexed streaming transport (i.e. it has a + * shared-storage backend). When false, streaming callbacks fall back to one + * NDJSON connection each. + */ +export function isStreamMultiplexed(config: { + stream?: {enabled?: boolean}; +}): boolean { + return !!config.stream?.enabled; +} diff --git a/dash/dash-renderer/src/utils/streamWorkerHost.ts b/dash/dash-renderer/src/utils/streamWorkerHost.ts new file mode 100644 index 0000000000..7e02fc96b2 --- /dev/null +++ b/dash/dash-renderer/src/utils/streamWorkerHost.ts @@ -0,0 +1,100 @@ +/** + * The worker's end of the SharedWorker streaming transport. + * + * One `StreamClient` (one connection id, one downlink) serves every tab of the + * browser. Each tab connects a port and asks for streams; the host runs them on + * the shared client and relays frames back to the asking port. When a tab + * announces it is going away, its streams are cancelled -- dropped here and + * cancelled server-side -- while the downlink stays up for the other tabs. + * + * Kept separate from the worker entry so it can be exercised with plain + * MessageChannel ports in the unit tests. + */ + +import { + StreamClient, + StreamHostMessage, + StreamPort, + StreamPortMessage +} from './streamClient'; + +/** The slice of SharedWorkerGlobalScope the host uses (fakeable in tests). */ +export interface StreamWorkerScope { + onconnect: ((event: MessageEvent) => void) | null; +} + +interface LiveStream { + rendererId: string; + workerRequestId: string; + url: string; + init: RequestInit; +} + +export function attachStreamWorkerHost( + scope: StreamWorkerScope, + client: StreamClient +): void { + // Keyed by the page's request id (unique: it embeds the tab's renderer id). + const live = new Map(); + + const post = (port: StreamPort, message: StreamPortMessage) => + port.postMessage(message); + + const startStream = ( + port: StreamPort, + message: Extract + ) => { + const {rendererId, requestId, url, init, endId, payload, transport} = + message; + if (transport) { + client.configure(transport); + } + const {requestId: workerRequestId, settled} = client.start( + url, + init, + endId, + payload, + frame => post(port, {type: 'frame', requestId, frame}) + ); + live.set(requestId, {rendererId, workerRequestId, url, init}); + settled.then( + () => { + live.delete(requestId); + post(port, {type: 'done', requestId}); + }, + (err: Error) => { + // A stream cancelled on unregister was already forgotten; its + // tab is gone and nobody is listening for the rejection. + if (live.delete(requestId)) { + post(port, { + type: 'error', + requestId, + message: err?.message || String(err) + }); + } + } + ); + }; + + const unregister = (rendererId: string) => { + for (const [requestId, stream] of Array.from(live.entries())) { + if (stream.rendererId === rendererId) { + live.delete(requestId); + client.cancel(stream.workerRequestId, stream.url, stream.init); + } + } + }; + + scope.onconnect = event => { + const port = event.ports[0] as StreamPort; + port.onmessage = e => { + const message = e.data as StreamHostMessage; + if (message.type === 'stream') { + startStream(port, message); + } else if (message.type === 'unregister') { + unregister(message.rendererId); + } + }; + port.start?.(); + }; +} diff --git a/dash/dash-renderer/src/utils/workerClient.ts b/dash/dash-renderer/src/utils/workerClient.ts index 6ed212f041..5724fdf1a9 100644 --- a/dash/dash-renderer/src/utils/workerClient.ts +++ b/dash/dash-renderer/src/utils/workerClient.ts @@ -25,6 +25,10 @@ export interface CallbackResponse { status: 'ok' | 'prevent_update' | 'error'; data?: Record; message?: string; + /** True for responses from a streaming callback */ + stream?: boolean; + /** True on the terminal frame of a streamed callback */ + done?: boolean; } /** Set props message payload */ @@ -45,6 +49,8 @@ export interface GetPropsRequestPayload { interface PendingRequest { resolve: (value: CallbackResponse) => void; reject: (error: Error) => void; + /** Receives intermediate frames from a streaming callback */ + onFrame?: (data: Record) => void; } /** @@ -205,9 +211,15 @@ class WorkerClient { /** * Send a callback request to the server via the worker. * @param payload The callback payload + * @param onFrame Optional handler for intermediate frames from a + * streaming callback; the returned promise still resolves once with + * the terminal response. * @returns Promise that resolves with the callback response */ - public async sendCallback(payload: unknown): Promise { + public async sendCallback( + payload: unknown, + onFrame?: (data: Record) => void + ): Promise { // Wait for initial connection if one is in progress if (this.connectionPromise && !this.isConnected) { await this.connectionPromise; @@ -220,7 +232,7 @@ class WorkerClient { const requestId = `${this.rendererId}-${++this.requestCounter}`; return new Promise((resolve, reject) => { - this.pendingCallbacks.set(requestId, {resolve, reject}); + this.pendingCallbacks.set(requestId, {resolve, reject, onFrame}); this.worker!.port.postMessage({ type: WorkerMessageType.CALLBACK_REQUEST, @@ -303,8 +315,21 @@ class WorkerClient { const requestId = message.requestId; const pending = this.pendingCallbacks.get(requestId); if (pending) { + const payload = message.payload; + if ( + payload?.stream && + !payload.done && + payload.status === 'ok' + ) { + // Intermediate stream frame: deliver it and keep the + // request pending until the terminal frame arrives. + if (pending.onFrame) { + pending.onFrame(payload.data); + } + break; + } this.pendingCallbacks.delete(requestId); - pending.resolve(message.payload); + pending.resolve(payload); } break; } diff --git a/dash/dash-renderer/src/workers/streamWorker.ts b/dash/dash-renderer/src/workers/streamWorker.ts new file mode 100644 index 0000000000..084bd2acfb --- /dev/null +++ b/dash/dash-renderer/src/workers/streamWorker.ts @@ -0,0 +1,11 @@ +/** + * Dash stream worker: a SharedWorker hosting the browser's single streaming + * downlink, shared by every tab (see utils/streamClient and + * utils/streamWorkerHost). Built to build/dash-stream-worker.js and served + * through the component suites like the WebSocket worker. + */ + +import {StreamClient} from '../utils/streamClient'; +import {attachStreamWorkerHost} from '../utils/streamWorkerHost'; + +attachStreamWorkerHost(self as any, new StreamClient()); diff --git a/dash/dash-renderer/tests/helpers/streamMocks.js b/dash/dash-renderer/tests/helpers/streamMocks.js new file mode 100644 index 0000000000..3cdc01c8ea --- /dev/null +++ b/dash/dash-renderer/tests/helpers/streamMocks.js @@ -0,0 +1,105 @@ +/** + * Doubles for the multiplexed streaming transport tests: a controllable NDJSON + * downlink body and a fetch that separates uplink, downlink and cancel POSTs. + */ + +// A controllable downlink body: push NDJSON lines and close it on demand. +export function makeDownlink() { + let controller; + const stream = new ReadableStream({ + start(c) { + controller = c; + } + }); + const enc = new TextEncoder(); + return { + stream, + push: obj => controller.enqueue(enc.encode(JSON.stringify(obj) + '\n')), + pushRaw: text => controller.enqueue(enc.encode(text)), + close: () => controller.close() + }; +} + +// A fetch double that separates uplink POSTs from downlink and cancel POSTs. +export function makeFetch() { + const uplinks = []; + const downlinks = []; + const cancels = []; + const fetchImpl = (url, init) => { + const body = JSON.parse(init.body); + if (body.streamDownlink) { + const dl = makeDownlink(); + downlinks.push({ + from: body.streamDownlink.from, + signal: init.signal, + dl + }); + return Promise.resolve(new Response(dl.stream, {status: 200})); + } + if (body.streamCancel) { + cancels.push({url, ...body.streamCancel}); + return Promise.resolve( + new Response( + JSON.stringify({ + multi: true, + stream: true, + cancelled: true + }), + {status: 200} + ) + ); + } + uplinks.push({url, ...body}); + return Promise.resolve( + new Response(JSON.stringify({multi: true, stream: true}), { + status: 200 + }) + ); + }; + return {fetchImpl, uplinks, downlinks, cancels}; +} + +export const tick = (ms = 5) => new Promise(r => setTimeout(r, ms)); + +export async function waitFor(pred, timeout = 1000) { + const end = Date.now() + timeout; + while (Date.now() < end) { + if (pred()) return; + await tick(5); + } + throw new Error('condition not met in time'); +} + +// A fetch double for POLL mode: each downlink POST returns a complete body +// holding the envelopes queued for that poll (from `queue`, a list of lists; +// empty when exhausted) and records when it was polled. +export function makePollFetch(queue = []) { + const uplinks = []; + const polls = []; + const cancels = []; + const fetchImpl = (url, init) => { + const body = JSON.parse(init.body); + if (body.streamDownlink) { + const envelopes = queue.length ? queue.shift() : []; + polls.push({ + url, + at: Date.now(), + from: body.streamDownlink.from, + n: envelopes.length + }); + const text = envelopes.map(e => JSON.stringify(e) + '\n').join(''); + return Promise.resolve(new Response(text, {status: 200})); + } + if (body.streamCancel) { + cancels.push({url, ...body.streamCancel}); + return Promise.resolve(new Response('{}', {status: 200})); + } + uplinks.push({url, ...body}); + return Promise.resolve( + new Response(JSON.stringify({multi: true, stream: true}), { + status: 200 + }) + ); + }; + return {fetchImpl, uplinks, polls, cancels, queue}; +} diff --git a/dash/dash-renderer/tests/requestSlot.test.js b/dash/dash-renderer/tests/requestSlot.test.js new file mode 100644 index 0000000000..a269ddd02d --- /dev/null +++ b/dash/dash-renderer/tests/requestSlot.test.js @@ -0,0 +1,78 @@ +import {expect} from 'chai'; +import {describe, it} from 'mocha'; + +import { + MAX_CONCURRENT_HTTP_CALLBACKS, + routedOverWebSocket, + usesRequestSlot +} from '../src/observers/requestSlot'; + +// Build a minimal ICallback with only the fields the scheduler inspects. +const cb = definition => ({callback: {websocket: false, ...definition}}); + +// Config flavors. isWebSocketEnabled/isWebSocketAvailable also require +// SharedWorker, which exists in the (Chrome) karma runner. +const HTTP = {}; +const WS_ENABLED = {websocket: {enabled: true, url: '/ws', worker_url: '/w'}}; +const WS_AVAILABLE_NOT_ENABLED = { + websocket: {enabled: false, url: '/ws', worker_url: '/w'} +}; + +describe('prioritizedCallbacks request-slot accounting', () => { + it('the concurrency cap is 12', () => { + expect(MAX_CONCURRENT_HTTP_CALLBACKS).to.equal(12); + }); + + describe('usesRequestSlot', () => { + it('a plain serverside HTTP callback counts against the budget', () => { + expect(usesRequestSlot(cb({}), HTTP)).to.equal(true); + }); + + it('excludes clientside callbacks (they run in-browser)', () => { + const clientside = cb({ + clientside_function: {namespace: 'ns', function_name: 'fn'} + }); + expect(usesRequestSlot(clientside, HTTP)).to.equal(false); + }); + + it('excludes streaming callbacks (they are long-lived)', () => { + expect(usesRequestSlot(cb({stream: true}), HTTP)).to.equal(false); + }); + + it('excludes every non-background callback when websocket is enabled', () => { + expect(usesRequestSlot(cb({}), WS_ENABLED)).to.equal(false); + }); + + it('still counts background callbacks even with websocket enabled', () => { + const background = cb({background: {interval: 1000}}); + expect(usesRequestSlot(background, WS_ENABLED)).to.equal(true); + }); + + it('excludes per-callback websocket routing when the transport is available', () => { + const perCallbackWs = cb({websocket: true}); + expect( + usesRequestSlot(perCallbackWs, WS_AVAILABLE_NOT_ENABLED) + ).to.equal(false); + }); + + it('counts a per-callback websocket that falls back to HTTP (transport unavailable)', () => { + const perCallbackWs = cb({websocket: true}); + expect(usesRequestSlot(perCallbackWs, HTTP)).to.equal(true); + }); + }); + + describe('routedOverWebSocket', () => { + it('routes non-background callbacks over the socket when enabled', () => { + expect(routedOverWebSocket(cb({}), WS_ENABLED)).to.equal(true); + }); + + it('never routes background callbacks over the socket', () => { + const background = cb({background: {interval: 1000}}); + expect(routedOverWebSocket(background, WS_ENABLED)).to.equal(false); + }); + + it('keeps callbacks on HTTP when no websocket transport is configured', () => { + expect(routedOverWebSocket(cb({}), HTTP)).to.equal(false); + }); + }); +}); diff --git a/dash/dash-renderer/tests/streamClient.test.js b/dash/dash-renderer/tests/streamClient.test.js new file mode 100644 index 0000000000..44bd7f143b --- /dev/null +++ b/dash/dash-renderer/tests/streamClient.test.js @@ -0,0 +1,544 @@ +import {expect} from 'chai'; +import {beforeEach, describe, it} from 'mocha'; + +import {StreamClient} from '../src/utils/streamClient'; +import {makePollFetch} from './helpers/streamMocks'; + +// A controllable downlink body: push NDJSON lines and close it on demand. +function makeDownlink() { + let controller; + const stream = new ReadableStream({ + start(c) { + controller = c; + } + }); + const enc = new TextEncoder(); + return { + stream, + push: obj => controller.enqueue(enc.encode(JSON.stringify(obj) + '\n')), + pushRaw: text => controller.enqueue(enc.encode(text)), + close: () => controller.close() + }; +} + +// A fetch double that separates uplink POSTs from downlink POSTs. +function makeFetch() { + const uplinks = []; + const downlinks = []; + // How the next downlink attempts behave: 'ok' streams a body, a number + // answers with that status, 'unreachable' rejects like a network error. + const mode = {downlink: 'ok'}; + const fetchImpl = (url, init) => { + const body = JSON.parse(init.body); + if (body.streamDownlink) { + if (mode.downlink === 'unreachable') { + downlinks.push({url, from: body.streamDownlink.from}); + return Promise.reject(new TypeError('Failed to fetch')); + } + if (typeof mode.downlink === 'number') { + downlinks.push({url, from: body.streamDownlink.from}); + return Promise.resolve( + new Response('', {status: mode.downlink}) + ); + } + const dl = makeDownlink(); + downlinks.push({ + url, + from: body.streamDownlink.from, + signal: init.signal, + dl + }); + return Promise.resolve(new Response(dl.stream, {status: 200})); + } + uplinks.push({url, ...body}); + return Promise.resolve( + new Response(JSON.stringify({multi: true, stream: true}), { + status: 200 + }) + ); + }; + return {fetchImpl, uplinks, downlinks, mode}; +} + +const tick = (ms = 5) => new Promise(r => setTimeout(r, ms)); +async function waitFor(pred, timeout = 1000) { + const end = Date.now() + timeout; + while (Date.now() < end) { + if (pred()) return; + await tick(5); + } + throw new Error('condition not met in time'); +} + +describe('StreamClient', () => { + let mock; + let client; + beforeEach(() => { + mock = makeFetch(); + client = new StreamClient({ + fetchImpl: mock.fetchImpl, + reconnectDelay: 10 + }); + }); + + it('tags the uplink with a request id and the signed endId, not a client topic id', async () => { + client.run('/cb', {}, 'e1', {output: 'a.b'}, () => {}); + await waitFor(() => mock.uplinks.length === 1); + const conn = mock.uplinks[0].streamConnection; + expect(conn.requestId).to.be.a('string'); + // The client never names the topic: only the server-signed endId keys it. + expect(conn.connectionId).to.equal(undefined); + expect(mock.uplinks[0].url).to.contain('endId=e1'); + expect(mock.uplinks[0].output).to.equal('a.b'); // original payload preserved + }); + + it('carries the endId on the downlink too', async () => { + client.run('/cb', {}, 'e1', {output: 'a.b'}, () => {}); + await waitFor(() => mock.downlinks.length === 1); + expect(mock.downlinks[0].url).to.contain('endId=e1'); + }); + + it('routes frames to onFrame and resolves on the done frame', async () => { + const frames = []; + const settled = client.run('/cb', {}, 'e1', {output: 'a.b'}, f => + frames.push(f) + ); + await waitFor(() => mock.downlinks.length === 1); + const {requestId} = mock.uplinks[0].streamConnection; + const dl = mock.downlinks[0].dl; + + dl.push({rid: requestId, frame: {response: {a: 1}}, seq: 1}); + dl.push({rid: requestId, frame: {response: {a: 2}}, seq: 2}); + dl.push({rid: requestId, frame: {done: true}, seq: 3}); + + await settled; + expect(frames).to.deep.equal([{response: {a: 1}}, {response: {a: 2}}]); + // The downlink is aborted once no callbacks remain in flight. + expect(mock.downlinks[0].signal.aborted).to.equal(true); + }); + + it('rejects on an error done frame', async () => { + const settled = client.run('/cb', {}, 'e1', {output: 'a.b'}, () => {}); + await waitFor(() => mock.downlinks.length === 1); + const {requestId} = mock.uplinks[0].streamConnection; + mock.downlinks[0].dl.push({ + rid: requestId, + frame: {done: true, error: {message: 'boom'}}, + seq: 1 + }); + let err; + await settled.catch(e => (err = e)); + expect(err).to.be.an('error'); + expect(err.message).to.contain('boom'); + }); + + it('multiplexes two callbacks over one downlink, routed by request id', async () => { + const aFrames = []; + const bFrames = []; + const a = client.run('/cb', {}, 'e1', {output: 'a'}, f => + aFrames.push(f) + ); + const b = client.run('/cb', {}, 'e1', {output: 'b'}, f => + bFrames.push(f) + ); + await waitFor(() => mock.uplinks.length === 2); + // Both share a single downlink connection. + expect(mock.downlinks.length).to.equal(1); + const ridA = mock.uplinks[0].streamConnection.requestId; + const ridB = mock.uplinks[1].streamConnection.requestId; + const dl = mock.downlinks[0].dl; + + dl.push({rid: ridB, frame: {response: {b: 1}}, seq: 1}); + dl.push({rid: ridA, frame: {response: {a: 1}}, seq: 2}); + dl.push({rid: ridA, frame: {done: true}, seq: 3}); + dl.push({rid: ridB, frame: {done: true}, seq: 4}); + + await Promise.all([a, b]); + expect(aFrames).to.deep.equal([{response: {a: 1}}]); + expect(bFrames).to.deep.equal([{response: {b: 1}}]); + }); + + it('reconnects from the last seen sequence when the downlink drops', async () => { + const frames = []; + const settled = client.run('/cb', {}, 'e1', {output: 'a'}, f => + frames.push(f) + ); + await waitFor(() => mock.downlinks.length === 1); + const {requestId} = mock.uplinks[0].streamConnection; + + mock.downlinks[0].dl.push({ + rid: requestId, + frame: {response: {a: 1}}, + seq: 5 + }); + await waitFor(() => frames.length === 1); + mock.downlinks[0].dl.close(); // drop mid-stream + + // It reconnects, resuming after the last applied sequence. + await waitFor(() => mock.downlinks.length === 2); + expect(mock.downlinks[1].from).to.equal(5); + mock.downlinks[1].dl.push({ + rid: requestId, + frame: {done: true}, + seq: 6 + }); + await settled; + expect(frames).to.deep.equal([{response: {a: 1}}]); + }); + + it('resets its cursor to the head on a reset envelope (server restart)', async () => { + const frames = []; + const settled = client.run('/cb', {}, 'e1', {output: 'a'}, f => + frames.push(f) + ); + await waitFor(() => mock.downlinks.length === 1); + const {requestId} = mock.uplinks[0].streamConnection; + + // Advance the cursor, then the server signals its buffer was lost. + mock.downlinks[0].dl.push({ + rid: requestId, + frame: {response: {a: 1}}, + seq: 5 + }); + await waitFor(() => frames.length === 1); + mock.downlinks[0].dl.push({reset: true}); + mock.downlinks[0].dl.close(); + + // The in-flight callback settles: its buffered frames are gone. + await settled; + expect(frames).to.deep.equal([{response: {a: 1}}]); + + // The next callback's downlink starts from the head (0), not the + // stale cursor (5). + client.run('/cb', {}, 'e1', {output: 'b'}, () => {}); + await waitFor(() => mock.downlinks.length === 2); + expect(mock.downlinks[1].from).to.equal(0); + }); + + it('fails the callback loudly when the uplink is rejected (unverified connection)', async () => { + // The server refuses an unverified multiplexed connection with a 403; no + // frames will arrive on the downlink, so the request must reject rather + // than hang. + const fetchImpl = () => + Promise.resolve(new Response('', {status: 403})); + const c = new StreamClient({fetchImpl, reconnectDelay: 10}); + let err; + await c + .run('/cb', {}, 'e1', {output: 'a'}, () => {}) + .catch(e => { + err = e; + }); + expect(err).to.be.an('error'); + expect(err.message).to.contain('403'); + }); + + it('skips keepalive blank lines', async () => { + const frames = []; + const settled = client.run('/cb', {}, 'e1', {output: 'a'}, f => + frames.push(f) + ); + await waitFor(() => mock.downlinks.length === 1); + const {requestId} = mock.uplinks[0].streamConnection; + const dl = mock.downlinks[0].dl; + dl.pushRaw('\n'); // keepalive + dl.push({rid: requestId, frame: {response: {a: 1}}, seq: 1}); + dl.pushRaw('\n'); + dl.push({rid: requestId, frame: {done: true}, seq: 2}); + await settled; + expect(frames).to.deep.equal([{response: {a: 1}}]); + }); + it('settles pending callbacks when the downlink is refused after a restart', async () => { + const frames = []; + const streamed = client.run('/cb', {}, 'e1', {output: 'a.b'}, f => + frames.push(f) + ); + const fresh = client.run('/cb', {}, 'e1', {output: 'c.d'}, () => {}); + await waitFor(() => mock.downlinks.length === 1); + const rid = mock.uplinks[0].streamConnection.requestId; + mock.downlinks[0].dl.push({rid, frame: {response: {a: 1}}, seq: 1}); + await waitFor(() => frames.length === 1); + + // The server restarts: the downlink drops, and the new process + // refuses our endId with 403 since its signing secret changed. + mock.mode.downlink = 403; + mock.downlinks[0].dl.close(); + + // The one that already applied a frame keeps it and resolves; the one + // that never got a frame rejects so the caller can fall back. + await streamed; + let err; + await fresh.catch(e => (err = e)); + expect(err.message).to.contain('403'); + expect(client.activeCount).to.equal(0); + + // And the downlink is not retried in a loop. + await tick(100); + expect(mock.downlinks.length).to.equal(2); + }); + + it('settles pending callbacks on a reset envelope', async () => { + const frames = []; + const streamed = client.run('/cb', {}, 'e1', {output: 'a.b'}, f => + frames.push(f) + ); + const fresh = client.run('/cb', {}, 'e1', {output: 'c.d'}, () => {}); + await waitFor(() => mock.downlinks.length === 1); + const rid = mock.uplinks[0].streamConnection.requestId; + const dl = mock.downlinks[0].dl; + dl.push({rid, frame: {response: {a: 1}}, seq: 4}); + await waitFor(() => frames.length === 1); + + dl.push({reset: true}); + + await streamed; + let err; + await fresh.catch(e => (err = e)); + expect(err.message).to.contain('reset'); + expect(client.activeCount).to.equal(0); + expect(mock.downlinks[0].signal.aborted).to.equal(true); + }); + + it('gives up after the server stays unreachable for the reconnect window', async () => { + client = new StreamClient({ + fetchImpl: mock.fetchImpl, + reconnectDelay: 10, + maxReconnectDelay: 20, + reconnectWindow: 100 + }); + const settled = client.run('/cb', {}, 'e1', {output: 'a.b'}, () => {}); + await waitFor(() => mock.downlinks.length === 1); + mock.mode.downlink = 'unreachable'; + mock.downlinks[0].dl.close(); + + let err; + await settled.catch(e => (err = e)); + expect(err.message).to.contain('could not reconnect'); + // Bounded: a handful of backed-off attempts, not one per tick forever. + const attempts = mock.downlinks.length; + expect(attempts).to.be.greaterThan(2); + await tick(100); + expect(mock.downlinks.length).to.equal(attempts); + }); + it('opens a fresh downlink for a stream started from a settled continuation', async () => { + const first = client.run('/cb', {}, 'e1', {output: 'a.b'}, () => {}); + await waitFor(() => mock.downlinks.length === 1); + mock.mode.downlink = 403; + mock.downlinks[0].dl.close(); + + // The caller reacts to the failure by starting another stream at once. + let second; + await first.catch(() => { + mock.mode.downlink = 'ok'; + second = client.run('/cb', {}, 'e1', {output: 'c.d'}, () => {}); + }); + await waitFor(() => mock.downlinks.length === 3); + const rid = mock.uplinks[1].streamConnection.requestId; + mock.downlinks[2].dl.push({rid, frame: {done: true}, seq: 1}); + await second; + expect(client.activeCount).to.equal(0); + }); + + it('backs off when the downlink is accepted but closes without data', async () => { + client = new StreamClient({ + fetchImpl: mock.fetchImpl, + reconnectDelay: 30, + maxReconnectDelay: 30, + reconnectWindow: 10000 + }); + client.run('/cb', {}, 'e1', {output: 'a.b'}, () => {}); + await waitFor(() => mock.downlinks.length === 1); + const t0 = Date.now(); + mock.downlinks[0].dl.close(); + await waitFor(() => mock.downlinks.length === 2); + mock.downlinks[1].dl.close(); + await waitFor(() => mock.downlinks.length === 3); + // Two empty closes -> two backoff sleeps, not an immediate burst. + expect(Date.now() - t0).to.be.at.least(55); + expect(client.activeCount).to.equal(1); + }); + it('opens the downlink only after the uplink is acknowledged', async () => { + // A single-threaded WSGI worker can serve one request at a time: the + // downlink must not be opened while the uplink is still in flight. + let ackUplink; + const gated = (url, init) => { + const body = JSON.parse(init.body); + if (body.streamConnection) { + return new Promise(resolve => { + ackUplink = () => + resolve( + new Response(JSON.stringify({stream: true}), { + status: 200 + }) + ); + }); + } + return mock.fetchImpl(url, init); + }; + client = new StreamClient({fetchImpl: gated, reconnectDelay: 10}); + client.run('/cb', {}, 'e1', {output: 'a.b'}, () => {}); + await waitFor(() => !!ackUplink); + await tick(20); + expect(mock.downlinks.length).to.equal(0); + + ackUplink(); + await waitFor(() => mock.downlinks.length === 1); + }); + + it('closes the downlink between acknowledged streams and reopens on the next ack', async () => { + const acks = []; + const gated = (url, init) => { + const body = JSON.parse(init.body); + if (body.streamConnection) { + return new Promise(resolve => { + acks.push({ + rid: body.streamConnection.requestId, + ack: () => resolve(new Response('{}', {status: 200})) + }); + }); + } + return mock.fetchImpl(url, init); + }; + client = new StreamClient({fetchImpl: gated, reconnectDelay: 10}); + const first = client.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + const second = client.run('/cb', {}, 'e1', {output: 'b'}, () => {}); + await waitFor(() => acks.length === 2); + acks[0].ack(); + await waitFor(() => mock.downlinks.length === 1); + + // The first stream finishes while the second is still waiting for its + // ack (queued behind the downlink on a sync worker): the downlink must + // close so that uplink can be served. + mock.downlinks[0].dl.push({ + rid: acks[0].rid, + frame: {done: true}, + seq: 1 + }); + await first; + expect(mock.downlinks[0].signal.aborted).to.equal(true); + expect(client.activeCount).to.equal(1); + + acks[1].ack(); + await waitFor(() => mock.downlinks.length === 2); + expect(mock.downlinks[1].from).to.equal(1); + mock.downlinks[1].dl.push({ + rid: acks[1].rid, + frame: {done: true}, + seq: 2 + }); + await second; + expect(client.activeCount).to.equal(0); + }); +}); + +describe('StreamClient cancellation', () => { + it('cancel drops the request, tells the server (same endId), and closes an idle downlink', async () => { + const mock = makePollFetch(); + const client = new StreamClient({fetchImpl: mock.fetchImpl}); + const frames = []; + const {requestId, settled} = client.start( + '/cb', + {}, + 'e1', + {output: 'a'}, + f => frames.push(f) + ); + await waitFor(() => mock.polls.length >= 1); + expect(mock.uplinks[0].streamConnection.requestId).to.equal(requestId); + + client.cancel(requestId, '/cb', {}); + let err; + await settled.catch(e => (err = e)); + expect(err.message).to.contain('cancelled'); + expect(mock.cancels).to.deep.equal([{url: '/cb?endId=e1', requestId}]); + // A late frame for the cancelled request is dropped, not delivered. + client.dispatchEnvelope({ + rid: requestId, + frame: {response: {a: 1}}, + seq: 1 + }); + expect(frames).to.deep.equal([]); + expect(client.activeCount).to.equal(0); + }); + + it("pins the first stream's endId for the connection while streams are in flight", async () => { + const mock = makePollFetch(); + const client = new StreamClient({fetchImpl: mock.fetchImpl}); + client.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + client.run('/cb', {}, 'e2', {output: 'b'}, () => {}); + await waitFor(() => mock.uplinks.length === 2); + expect(mock.uplinks.map(u => u.url)).to.deep.equal([ + '/cb?endId=e1', + '/cb?endId=e1' + ]); + expect(client.connectionEndId).to.equal('e1'); + }); +}); + +describe('StreamClient in poll mode', () => { + const frameFor = rid => ({rid, frame: {response: {a: 1}}, seq: 1}); + const doneFor = rid => ({rid, frame: {done: true}, seq: 2}); + + it('re-polls promptly while frames flow and resolves on done', async () => { + const mock = makePollFetch(); + const client = new StreamClient({ + fetchImpl: mock.fetchImpl, + mode: 'poll', + pollInterval: 20 + }); + const frames = []; + const settled = client.run('/cb', {}, 'e1', {output: 'a'}, f => + frames.push(f) + ); + await waitFor(() => mock.uplinks.length === 1); + const rid = mock.uplinks[0].streamConnection.requestId; + mock.queue.push([frameFor(rid)], [doneFor(rid)]); + await settled; + expect(frames).to.deep.equal([{response: {a: 1}}]); + const withFrames = mock.polls.filter(p => p.n > 0); + expect(withFrames.length).to.equal(2); + expect(mock.polls[mock.polls.length - 1].from).to.equal(1); + expect(mock.polls[0].url).to.equal('/cb?endId=e1'); + expect(client.activeCount).to.equal(0); + }); + + it('backs off while quiet, bounded, and polls at once for a new stream', async () => { + const mock = makePollFetch(); + const client = new StreamClient({ + fetchImpl: mock.fetchImpl, + mode: 'poll', + pollInterval: 30 + }); + client.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + await waitFor(() => mock.uplinks.length === 1); + // Idle polls hold the 30ms pace for two polls, then back off + // geometrically (60, 120, 150 capped at 5x). + await waitFor(() => mock.polls.length >= 7, 4000); + const gaps = mock.polls.slice(1).map((p, i) => p.at - mock.polls[i].at); + expect(Math.max(gaps[0], gaps[1])).to.be.lessThan(90); + expect(gaps[3]).to.be.greaterThan(gaps[2] * 1.3); + expect(gaps[4]).to.be.greaterThan(gaps[2]); + // The cap holds: no pause grows past five intervals. + expect(Math.max(gaps[4], gaps[5])).to.be.lessThan(30 * 5 + 70); + // A new stream wakes the pause: the next poll comes right away. + const before = mock.polls.length; + client.run('/cb', {}, 'e1', {output: 'b'}, () => {}); + await waitFor(() => mock.uplinks.length === 2); + await tick(15); + expect(mock.polls.length).to.be.greaterThan(before); + }); + + it('in stream mode, an empty clean end backs off like a drop (no hot loop)', async () => { + const mock = makePollFetch(); + const client = new StreamClient({ + fetchImpl: mock.fetchImpl, + reconnectDelay: 30, + maxReconnectDelay: 30 + }); + client.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + await waitFor(() => mock.polls.length >= 4, 2000); + const span = mock.polls[3].at - mock.polls[0].at; + expect(span).to.be.at.least(80); + expect(span).to.be.lessThan(600); + expect(client.transport.mode).to.equal(undefined); + }); +}); diff --git a/dash/dash-renderer/tests/streamWorkerHost.test.js b/dash/dash-renderer/tests/streamWorkerHost.test.js new file mode 100644 index 0000000000..199d5ee1c9 --- /dev/null +++ b/dash/dash-renderer/tests/streamWorkerHost.test.js @@ -0,0 +1,157 @@ +import {expect} from 'chai'; +import {beforeEach, describe, it} from 'mocha'; + +import {SharedStreamClient, StreamClient} from '../src/utils/streamClient'; +import {attachStreamWorkerHost} from '../src/utils/streamWorkerHost'; +import {makeFetch, waitFor} from './helpers/streamMocks'; + +// The worker side: one StreamClient (one downlink) behind a fake worker scope. +// Each "tab" is a MessageChannel: port1 connects to the host, port2 is the +// page's SharedStreamClient. +function makeWorker() { + const mock = makeFetch(); + const client = new StreamClient({ + fetchImpl: mock.fetchImpl, + reconnectDelay: 10 + }); + const scope = {onconnect: null}; + attachStreamWorkerHost(scope, client); + const connectTab = rendererId => { + const channel = new MessageChannel(); + scope.onconnect({ports: [channel.port1]}); + return new SharedStreamClient(channel.port2, rendererId); + }; + return {mock, client, connectTab}; +} + +describe('SharedWorker stream transport', () => { + let worker; + beforeEach(() => { + worker = makeWorker(); + }); + + it('relays frames and the terminal done back to the asking tab', async () => { + const tab = worker.connectTab('tab-a'); + const frames = []; + const settled = tab.run('/cb', {}, 'e1', {output: 'a.b'}, f => + frames.push(f) + ); + await waitFor(() => worker.mock.downlinks.length === 1); + // The uplink carried the tab's signed endId and the payload. + const conn = worker.mock.uplinks[0].streamConnection; + expect(worker.mock.uplinks[0].url).to.equal('/cb?endId=e1'); + expect(worker.client.connectionEndId).to.equal('e1'); + expect(worker.mock.uplinks[0].output).to.equal('a.b'); + + const dl = worker.mock.downlinks[0].dl; + dl.push({rid: conn.requestId, frame: {response: {a: 1}}, seq: 1}); + dl.push({rid: conn.requestId, frame: {done: true}, seq: 2}); + await settled; + expect(frames).to.deep.equal([{response: {a: 1}}]); + expect(tab.activeCount).to.equal(0); + }); + + it('rejects the tab on an error frame', async () => { + const tab = worker.connectTab('tab-a'); + const settled = tab.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + await waitFor(() => worker.mock.downlinks.length === 1); + const {requestId} = worker.mock.uplinks[0].streamConnection; + worker.mock.downlinks[0].dl.push({ + rid: requestId, + frame: {done: true, error: {message: 'boom'}}, + seq: 1 + }); + let err; + await settled.catch(e => (err = e)); + expect(err.message).to.contain('boom'); + }); + + it('serves several tabs over one downlink, routed to the right tab', async () => { + const tabA = worker.connectTab('tab-a'); + const tabB = worker.connectTab('tab-b'); + const aFrames = []; + const bFrames = []; + const a = tabA.run('/cb', {}, 'e1', {output: 'a'}, f => + aFrames.push(f) + ); + const b = tabB.run('/cb', {}, 'e2', {output: 'b'}, f => + bFrames.push(f) + ); + await waitFor(() => worker.mock.uplinks.length === 2); + expect(worker.mock.downlinks.length).to.equal(1); + // Tab B's stream rides tab A's connection: one endId keys the topic. + expect(worker.mock.uplinks[1].url).to.equal('/cb?endId=e1'); + const ridA = worker.mock.uplinks[0].streamConnection.requestId; + const ridB = worker.mock.uplinks[1].streamConnection.requestId; + const dl = worker.mock.downlinks[0].dl; + dl.push({rid: ridB, frame: {response: {b: 1}}, seq: 1}); + dl.push({rid: ridA, frame: {response: {a: 1}}, seq: 2}); + dl.push({rid: ridA, frame: {done: true}, seq: 3}); + dl.push({rid: ridB, frame: {done: true}, seq: 4}); + await Promise.all([a, b]); + expect(aFrames).to.deep.equal([{response: {a: 1}}]); + expect(bFrames).to.deep.equal([{response: {b: 1}}]); + }); + + it('cancels a departing tab’s streams server-side and keeps serving the rest', async () => { + const tabA = worker.connectTab('tab-a'); + const tabB = worker.connectTab('tab-b'); + const bFrames = []; + const a = tabA.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + const b = tabB.run('/cb', {}, 'e2', {output: 'b'}, f => + bFrames.push(f) + ); + await waitFor(() => worker.mock.uplinks.length === 2); + const ridA = worker.mock.uplinks[0].streamConnection.requestId; + const ridB = worker.mock.uplinks[1].streamConnection.requestId; + + tabA.release(); // tab A closed + await waitFor(() => worker.mock.cancels.length === 1); + expect(worker.mock.cancels[0]).to.deep.equal({ + url: '/cb?endId=e1', + requestId: ridA + }); + // Tab B still has a stream in flight: the shared downlink stays open. + expect(worker.mock.downlinks[0].signal.aborted).to.equal(false); + + const dl = worker.mock.downlinks[0].dl; + dl.push({rid: ridA, frame: {response: {a: 1}}, seq: 1}); // dropped + dl.push({rid: ridB, frame: {response: {b: 1}}, seq: 2}); + dl.push({rid: ridB, frame: {done: true}, seq: 3}); + await b; + expect(bFrames).to.deep.equal([{response: {b: 1}}]); + // Nobody is left listening for A; its page-side promise stays pending + // (the tab is gone) rather than surfacing an error anywhere. + expect(tabA.activeCount).to.equal(1); + void a; + // With B done and A cancelled, the downlink is released. + expect(worker.mock.downlinks[0].signal.aborted).to.equal(true); + }); + + it('fail() rejects everything in flight and marks the transport broken', async () => { + const tab = worker.connectTab('tab-a'); + const settled = tab.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + tab.fail(new Error('worker died')); + let err; + await settled.catch(e => (err = e)); + expect(err.message).to.equal('worker died'); + expect(tab.broken).to.equal(true); + }); + + it('adopts the downlink mode the page passes along', async () => { + const channel = new MessageChannel(); + const scope = {onconnect: null}; + attachStreamWorkerHost(scope, worker.client); + scope.onconnect({ports: [channel.port1]}); + const tab = new SharedStreamClient(channel.port2, 'tab-a', { + mode: 'poll', + pollInterval: 250 + }); + tab.run('/cb', {}, 'e1', {output: 'a'}, () => {}); + await waitFor(() => worker.mock.uplinks.length === 1); + expect(worker.client.transport).to.deep.equal({ + mode: 'poll', + pollInterval: 250 + }); + }); +}); diff --git a/dash/dash-renderer/webpack.base.config.js b/dash/dash-renderer/webpack.base.config.js index d2cdf33a6e..7ee8fcf573 100644 --- a/dash/dash-renderer/webpack.base.config.js +++ b/dash/dash-renderer/webpack.base.config.js @@ -88,12 +88,14 @@ const shimOptions = { } }; -// WebSocket Worker configuration -const workerOptions = { +// SharedWorker bundles. Each worker gets its own compilation with an explicit +// tsconfig: ts-loader would otherwise type-check both entries against +// whichever tsconfig it finds first, and the two need different libs +// (WebWorker for the WebSocket worker package, the renderer's DOM lib for the +// stream worker, which shares its transport code with the page). +const workerConfig = (name, entry, configFile) => ({ mode: 'production', - entry: { - 'dash-ws-worker': '../../@plotly/dash-websocket-worker/src/worker.ts', - }, + entry: {[name]: entry}, output: { path: path.resolve(__dirname, "build"), filename: '[name].js', @@ -104,14 +106,28 @@ const workerOptions = { { test: /\.ts$/, exclude: /node_modules/, - use: ['ts-loader'], + use: [{loader: 'ts-loader', options: {configFile}}], }, ] }, resolve: { extensions: ['.ts', '.js'] } -}; +}); + +// WebSocket Worker configuration +const workerOptions = workerConfig( + 'dash-ws-worker', + '../../@plotly/dash-websocket-worker/src/worker.ts', + path.resolve(__dirname, '../../@plotly/dash-websocket-worker/tsconfig.json') +); + +// Streaming downlink worker configuration +const streamWorkerOptions = workerConfig( + 'dash-stream-worker', + './src/workers/streamWorker.ts', + path.resolve(__dirname, 'tsconfig.json') +); module.exports = options => [ R.mergeAll([ @@ -151,8 +167,9 @@ module.exports = options => [ ), } ]), - // WebSocket Worker build + // SharedWorker builds (WebSocket transport, streaming downlink) workerOptions, + streamWorkerOptions, // React compatibility shim build shimOptions ]; diff --git a/dash/dash.py b/dash/dash.py index 563e24b92b..ed1ed5ebc8 100644 --- a/dash/dash.py +++ b/dash/dash.py @@ -448,6 +448,24 @@ class Dash(ObsoleteChecker): :param csrf_header_name: Name of the HTTP header to send the CSRF token in. Default ``'X-CSRFToken'``. :type csrf_header_name: string + + :param stream_keepalive_interval: How long a streaming callback may + go without yielding, in milliseconds, before the response emits a + blank keepalive line. Default 15000. Keeps proxy idle timeouts + (nginx ``proxy_read_timeout`` defaults to 60s) from closing a stream + while the callback is still working. Set to None or 0 to disable. + :type stream_keepalive_interval: int or None + + :param stream_poll_interval: On WSGI (Flask), how often the browser polls + for streamed frames while it has streams running, in milliseconds. + Default 100. It re-polls at this interval while frames are arriving + and backs off to five times it while quiet, so frame latency stays + within a few hundred milliseconds. A WSGI worker thread is held for the life of a + response, so instead of one open connection per browser (which pins a + thread each and exhausts a pool at a few dozen browsers) each poll + takes a thread for milliseconds. ASGI backends (Quart, FastAPI) keep + one open connection per browser instead and ignore this. + :type stream_poll_interval: int """ _plotlyjs_url: str @@ -508,6 +526,8 @@ def __init__( # pylint: disable=too-many-statements, too-many-branches websocket_heartbeat_interval: Optional[int] = 30000, websocket_batch_delay: Optional[float] = 0.005, websocket_max_workers: Optional[int] = 4, + stream_keepalive_interval: Optional[int] = 15000, + stream_poll_interval: int = 100, shared_storage: Optional[ Union[Type[BaseSharedStorage], BaseSharedStorage] ] = LocalSharedStorage, @@ -686,6 +706,8 @@ def __init__( # pylint: disable=too-many-statements, too-many-branches self._websocket_heartbeat_interval = websocket_heartbeat_interval self._websocket_batch_delay = websocket_batch_delay self._websocket_max_workers = websocket_max_workers + self._stream_keepalive_interval = stream_keepalive_interval + self._stream_poll_interval = stream_poll_interval # Shared storage (state manager + pub/sub). Started lazily on first # access so it costs nothing until used and never binds in a gunicorn @@ -955,6 +977,14 @@ def layout(self, value: Any): _validate.validate_layout(value, layout_value) self.validation_layout = layout_value + @property + def shared_storage_enabled(self) -> bool: + """Whether this app has a shared-storage backend (not ``None``). + + Cheap to read and does not start the backend, unlike ``shared_storage``. + """ + return self._shared_storage_arg is not None + @property def shared_storage(self) -> BaseSharedStorage: """The app's shared storage (state manager + pub/sub), backend-agnostic. @@ -1170,6 +1200,24 @@ def _config(self): "heartbeat_interval": self._websocket_heartbeat_interval, } + # Streaming callbacks use the single multiplexed downlink only when a + # shared-storage backend is available to broker frames across workers; + # otherwise the client streams each callback on its own connection. + # The worker script hosts that downlink in a SharedWorker so every tab + # of the browser shares one connection (browsers cap connections per + # host); the renderer falls back to a per-tab downlink without it. + # "mode" is how the downlink is served: one open connection (ASGI) or + # polling (WSGI, where an open response would pin a worker thread). + config["stream"] = {"enabled": self.shared_storage_enabled} + if self.shared_storage_enabled: + config["stream"].update( + worker_url=self._get_worker_url( + "dash-renderer/build/dash-stream-worker.js" + ), + mode=self.backend.downlink_mode, + poll_interval=self._stream_poll_interval, + ) + return config def serve_reload_hash(self): @@ -1197,13 +1245,14 @@ def serve_health(self): """ return self.backend.make_response("OK", status=200, mimetype="text/plain") - def _get_worker_url(self) -> str: - """Get the URL for the WebSocket worker script. + def _get_worker_url( + self, relative_path: str = "dash-renderer/build/dash-ws-worker.js" + ) -> str: + """Get the URL for a renderer worker script (WebSocket or streaming). Returns: The fingerprinted URL for the worker script served via component suites. """ - relative_path = "dash-renderer/build/dash-ws-worker.js" namespace = "dash" # Register the path so it can be served diff --git a/dash/exceptions.py b/dash/exceptions.py index 9366f9359c..9d04caece8 100644 --- a/dash/exceptions.py +++ b/dash/exceptions.py @@ -121,3 +121,7 @@ class WebSocketCallbackError(CallbackException): class WebsocketDisconnected(CallbackException): pass + + +class StreamCallbackError(CallbackException): + pass diff --git a/dash/types.py b/dash/types.py index 9da246b16c..175a1f61ca 100644 --- a/dash/types.py +++ b/dash/types.py @@ -102,3 +102,6 @@ class CallbackExecutionResponse(TypedDict): response: NotRequired[Dict[str, CallbackOutput]] sideUpdate: NotRequired[Dict[str, CallbackSideOutput]] dist: NotRequired[List[Any]] + # Streaming callbacks: terminal frame marker and mid-stream error. + done: NotRequired[bool] + error: NotRequired[Dict[str, str]] diff --git a/requirements/ci.txt b/requirements/ci.txt index ea1ba48fe0..db3c682712 100644 --- a/requirements/ci.txt +++ b/requirements/ci.txt @@ -3,6 +3,7 @@ black==22.3.0 flake8==7.3.0 flaky==3.8.1 flask-talisman==1.1.0 +httpx # fastapi.testclient (tests/websocket/test_ws_stream.py) ipython<9.0.0 mimesis<=11.1.0; python_version < "3.10" mimesis<=21.0.0; python_version >= "3.10" diff --git a/tests/shared_storage/test_engine.py b/tests/shared_storage/test_engine.py index 11eafa4061..26a7e45545 100644 --- a/tests/shared_storage/test_engine.py +++ b/tests/shared_storage/test_engine.py @@ -155,3 +155,40 @@ def test_multiple_subscribers_each_get_every_message(): b = e.poll("t", cursor, timeout=1) assert a.messages == ["a", "b"] assert b.messages == ["a", "b"] # independent cursors, both see all + + +def test_apoll_wakes_on_publish_from_another_thread(): + import asyncio + + e = StoreEngine() + + async def scenario(): + loop = asyncio.get_running_loop() + threading.Timer(0.1, lambda: e.publish("t", "hello")).start() + started = loop.time() + res = await e.apoll("t", 0, timeout=5.0) + assert res.messages == ["hello"] and res.last_seq == 1 and not res.gap + assert loop.time() - started < 2.0 # woken, not timed out + # Nothing new: times out empty without blocking a thread. + res = await e.apoll("t", 1, timeout=0.05) + assert res.messages == [] and res.last_seq == 1 + # A waiter that timed out was removed from the topic. + assert e._topic("t").waiters == [] + + asyncio.run(scenario()) + + +def test_apoll_wakes_on_close_and_serves_many_waiters(): + import asyncio + + e = StoreEngine() + + async def scenario(): + waits = [asyncio.ensure_future(e.apoll(f"t{i}", 0, 5.0)) for i in range(200)] + await asyncio.sleep(0.05) + assert sum(len(e._topic(f"t{i}").waiters) for i in range(200)) == 200 + threading.Timer(0.05, e.close).start() + results = await asyncio.wait_for(asyncio.gather(*waits), 5.0) + assert all(r.messages == [] for r in results) + + asyncio.run(scenario()) diff --git a/tests/shared_storage/test_local_cross_process.py b/tests/shared_storage/test_local_cross_process.py index 7d6f6f136e..cff38d98db 100644 --- a/tests/shared_storage/test_local_cross_process.py +++ b/tests/shared_storage/test_local_cross_process.py @@ -203,6 +203,46 @@ def test_reelection_recovers_persisted_data(tmp_path): o.stop() +def test_patch_frame_published_over_socket(owner): + # Reproduces the multi-process failure: a client publishes a streaming frame + # carrying a dash.Patch to the owner over the socket. The frame must arrive + # reduced to plain JSON (the wire codec can't encode a Patch). + from dash import Patch + from dash._stream_hub import publish_frame, subscribe_envelopes + + ns, _ = owner + client = LocalSharedStorage(namespace=ns) + client.start() + assert not client._coord.is_owner() # publishing over the socket + + out = [] + + def drain(): + gen = subscribe_envelopes(client, "cp", replay_from=0) + for env in gen: + out.append(env) + if env["frame"].get("done"): + break + gen.close() + + th = threading.Thread(target=drain, daemon=True) + th.start() + time.sleep(0.4) + + patch = Patch() + patch["a"] = 1 + publish_frame(client, "cp", "r1", {"response": {"o": {"children": patch}}}) + publish_frame(client, "cp", "r1", {"done": True}) + + th.join(timeout=8) + assert ( + out[0]["frame"]["response"]["o"]["children"]["__dash_patch_update"] + == "__dash_patch_update" + ) + assert out[-1]["frame"] == {"done": True} + client.close() + + def _wait_until(pred, timeout): end = time.monotonic() + timeout while time.monotonic() < end: @@ -210,3 +250,64 @@ def _wait_until(pred, timeout): return time.sleep(0.02) raise AssertionError("condition not met in time") + + +def test_async_client_subscription_across_processes(owner): + """The asyncio subscription path (ASGI servers) long-polls the owner over an + asyncio-streams connection -- no executor thread -- and resumes after a + reconnect from its cursor.""" + import asyncio + + ns, o = owner + client = LocalSharedStorage(namespace=ns) + client.start() + assert not client._coord.is_owner() + sub = client.subscribe("atopic") + + async def consume(): + got = [] + async for msg in sub: + got.append(msg) + if len(got) == 3: + break + return got + + async def scenario(): + task = asyncio.ensure_future(consume()) + await asyncio.sleep(0.3) # the long-poll is established + for i in range(3): + o.do("publish", "atopic", f"a{i}") + return await asyncio.wait_for(task, 10) + + assert asyncio.run(scenario()) == ["a0", "a1", "a2"] + sub.close() + client.close() + + +def test_async_ops_across_processes(owner): + """aget/aset/apublish from a client worker's event loop go over an + asyncio-streams connection of their own: no executor, no blocking.""" + import asyncio + + ns, o = owner + client = LocalSharedStorage(namespace=ns) + client.start() + + async def scenario(): + await client.aset("k", {"v": 1}) + assert await client.aget("k") == {"v": 1} + assert await client.aget("missing", "dflt") == "dflt" + sub = client.subscribe("apub", replay_from=0) + await client.apublish("apub", "m1") + await client.apublish("apub", "m2") + assert sub.poll(0.0) == [(1, "m1"), (2, "m2")] + sub.close() + await client.adelete("k") + assert await client.aget("k") is None + # Many concurrent calls share the one connection safely. + await asyncio.gather(*(client.aset(f"k{i}", i) for i in range(50))) + assert await client.aget("k49") == 49 + + asyncio.run(scenario()) + assert client.get("k7") == 7 # landed in the owner, visible over the sync path + client.close() diff --git a/tests/shared_storage/test_stream_hub.py b/tests/shared_storage/test_stream_hub.py new file mode 100644 index 0000000000..894a425a5a --- /dev/null +++ b/tests/shared_storage/test_stream_hub.py @@ -0,0 +1,199 @@ +"""The stream hub: streaming frames multiplexed over shared-storage pub/sub.""" +import asyncio +import threading +import time +import uuid + +import pytest + +from dash._shared_storage import LocalSharedStorage +from dash._streaming import StreamedCallbackResponse +from dash._stream_hub import ( + _pending_pumps, + apump_to_storage, + publish_frame, + pump_to_storage, + shutdown_active_streams, + spawn_async_pump, + stream_topic, + subscribe_envelopes, +) + + +def _frames_marker(*frames): + async def gen(): + for frame in frames: + yield frame + + return StreamedCallbackResponse(gen(), is_async=True) + + +@pytest.fixture +def storage(): + s = LocalSharedStorage(namespace=f"hub-{uuid.uuid4().hex[:12]}") + s.start() + yield s + s.close() + + +def _drain(storage, conn_id, stop_after, out, replay_from=None): + gen = subscribe_envelopes(storage, conn_id, replay_from) + for envelope in gen: + out.append(envelope) + if len(out) >= stop_after: + break + gen.close() + + +def test_topic_name(): + assert stream_topic("abc") == "_dash_stream:abc" + + +def test_downlink_relays_tagged_frames(storage): + out = [] + th = threading.Thread(target=_drain, args=(storage, "c1", 2, out)) + th.start() + time.sleep(0.3) # let the subscription establish (pub/sub starts at head) + + publish_frame( + storage, "c1", "r1", {"multi": True, "response": {"o": {"children": "a"}}} + ) + publish_frame(storage, "c1", "r1", {"done": True}) + + th.join(timeout=5) + assert [(e["rid"], e["frame"]) for e in out] == [ + ("r1", {"multi": True, "response": {"o": {"children": "a"}}}), + ("r1", {"done": True}), + ] + # Each envelope carries its storage seq, ascending, for reconnect resume. + assert [e["seq"] for e in out] == [1, 2] + + +def test_downlink_multiplexes_multiple_callbacks(storage): + out = [] + th = threading.Thread(target=_drain, args=(storage, "c2", 4, out)) + th.start() + time.sleep(0.3) + + # Two callbacks' frames interleave on one connection; the client demuxes by rid. + publish_frame(storage, "c2", "r1", {"response": {"a": 1}}) + publish_frame(storage, "c2", "r2", {"response": {"b": 1}}) + publish_frame(storage, "c2", "r1", {"done": True}) + publish_frame(storage, "c2", "r2", {"done": True}) + + th.join(timeout=5) + rids = [e["rid"] for e in out] + assert rids == ["r1", "r2", "r1", "r2"] + + +def test_reconnecting_downlink_replays_from_cursor(storage): + topic = stream_topic("c3") + # Publish before anyone subscribes; a reconnecting downlink replays from 0. + publish_frame(storage, "c3", "r1", {"response": {"a": 1}}) + publish_frame(storage, "c3", "r1", {"done": True}) + + out = [] + _drain(storage, "c3", 2, out, replay_from=0) + assert [e["frame"] for e in out] == [{"response": {"a": 1}}, {"done": True}] + assert storage.get(topic) is None # topics are pub/sub, not KV keys + + +def test_async_pump_publishes_frames(storage): + marker = _frames_marker({"response": {"a": 1}}, {"done": True}) + out = [] + th = threading.Thread(target=_drain, args=(storage, "cp", 2, out)) + th.start() + time.sleep(0.3) + asyncio.run(apump_to_storage(storage, "cp", "r9", marker)) + th.join(timeout=5) + assert [(e["rid"], e["frame"]) for e in out] == [ + ("r9", {"response": {"a": 1}}), + ("r9", {"done": True}), + ] + + +def test_publish_frame_reduces_patch_to_plain_json(storage): + # A frame carrying a dash.Patch must be reduced to plain JSON before it hits + # shared storage, or the data-only wire codec (msgspec) cannot encode it -- + # the failure seen in multi-process deployments (the socket path). + from dash import Patch + from dash._shared_storage._codec import encode + + patch = Patch() + patch["x"] = 1 + frame = {"multi": True, "response": {"o": {"children": patch}}} + + out = [] + th = threading.Thread(target=_drain, args=(storage, "cpatch", 1, out), daemon=True) + th.start() + time.sleep(0.3) + publish_frame(storage, "cpatch", "r1", frame) + th.join(timeout=5) + + delivered = out[0]["frame"] + encode(delivered) # the op that raised over the socket; must not raise now + child = delivered["response"]["o"]["children"] + assert child["__dash_patch_update"] == "__dash_patch_update" + + +def test_sync_pump_drives_async_frames(storage): + marker = _frames_marker({"response": {"b": 2}}, {"done": True}) + out = [] + th = threading.Thread(target=_drain, args=(storage, "cs", 2, out)) + th.start() + time.sleep(0.3) + pump_to_storage(storage, "cs", "r10", marker) # sync driver over async gen + th.join(timeout=5) + assert [e["frame"] for e in out] == [{"response": {"b": 2}}, {"done": True}] + + +def test_downlink_resets_when_cursor_is_ahead_of_head(storage): + # A stale cursor -- from a page whose server restarted, or whose storage + # owner was re-elected -- points past everything the fresh topic has + # produced. The downlink must surface a single reset envelope so the client + # resets its cursor, instead of stalling until the fresh sequence climbs + # back past the stale cursor. + envelopes = list(subscribe_envelopes(storage, "c-reset", replay_from=5)) + assert envelopes == [{"reset": True}] + + +def test_shutdown_active_streams_closes_open_downlink(storage): + # An idle downlink sits in a long poll; a graceful shutdown must be able to + # close it (otherwise the server can't exit -- the reported Ctrl+C hang). + done = threading.Event() + + def drain(): + for _envelope in subscribe_envelopes(storage, "c-shutdown"): + pass + done.set() + + th = threading.Thread(target=drain, daemon=True) + th.start() + time.sleep(0.3) # subscription established, now blocked in the poll + assert not done.is_set() + + shutdown_active_streams() + + th.join(timeout=5) + assert done.is_set() + + +def test_shutdown_active_streams_cancels_pump(storage): + # A running stream pump (driving a long-lived callback generator) must be + # cancelled on shutdown so the callback stops producing frames. + async def run(): + async def gen(): + for i in range(1000): + await asyncio.sleep(0.05) + yield {"response": {"n": i}} + + marker = StreamedCallbackResponse(gen(), is_async=True) + spawn_async_pump(storage, "c-pump", "r1", marker) + await asyncio.sleep(0.15) + assert _pending_pumps # the pump is running + + shutdown_active_streams() + await asyncio.sleep(0.15) # let the cancellation propagate + assert not _pending_pumps # cancelled and cleaned up + + asyncio.run(run()) diff --git a/tests/streaming/__init__.py b/tests/streaming/__init__.py new file mode 100644 index 0000000000..255a3fbc17 --- /dev/null +++ b/tests/streaming/__init__.py @@ -0,0 +1,2 @@ +# Streaming (async generator) callback tests. Require the `async` extra +# (flask[async]); run in their own CI job. diff --git a/tests/streaming/conftest.py b/tests/streaming/conftest.py new file mode 100644 index 0000000000..94467a0d00 --- /dev/null +++ b/tests/streaming/conftest.py @@ -0,0 +1,32 @@ +"""Give each streaming test its own shared-storage namespace. + +The default namespace is derived from cwd + argv so that every worker process of +one app shares an owner. In the test suite, though, many apps run in a single +process; without isolation they would contend for one owner and reset each +other's connections at teardown. Patching the default namespace per test keeps +each app its own owner. +""" +import uuid + +import pytest + +import dash._shared_storage.local as _local +from dash._streaming import _shutdown as _streaming_shutdown + + +@pytest.fixture(autouse=True) +def _isolate_shared_storage(monkeypatch): + namespace = f"streamtest-{uuid.uuid4().hex[:12]}" + monkeypatch.setattr(_local, "_default_namespace", lambda: namespace) + yield + + +@pytest.fixture(autouse=True) +def _clear_stream_shutdown_flag(): + """A server shutting down in an earlier test (a dash_duo teardown, a + TestClient lifespan exit) sets the streaming shutdown flag, and only a + server start clears it. Unit tests that drive the frame generators directly + would otherwise exit them immediately.""" + _streaming_shutdown.clear() + yield + _streaming_shutdown.clear() diff --git a/tests/streaming/test_stream_callbacks_integration.py b/tests/streaming/test_stream_callbacks_integration.py new file mode 100644 index 0000000000..86f45bbbf7 --- /dev/null +++ b/tests/streaming/test_stream_callbacks_integration.py @@ -0,0 +1,332 @@ +"""Browser integration tests for streaming callbacks over HTTP (NDJSON).""" +import asyncio +import time + +from dash import ( + Dash, + Input, + Output, + Patch, + html, + no_update, + set_props, +) +from dash.testing.wait import until + + +def test_stst001_stream_progressive_render(dash_duo): + """Intermediate yields render before the stream completes.""" + app = Dash(__name__) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="out", children="idle"), + ] + ) + + @app.callback( + Output("out", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_cb(n): + # Each step stays up well past the driver's 0.5s poll interval, so + # the wait below cannot miss it. + yield "step-1" + await asyncio.sleep(1.0) + yield "step-2" + await asyncio.sleep(1.0) + yield "done" + + dash_duo.start_server(app) + dash_duo.find_element("#btn").click() + # Each yield renders while the callback is still running. + dash_duo.wait_for_text_to_equal("#out", "step-1") + dash_duo.wait_for_text_to_equal("#out", "step-2") + dash_duo.wait_for_text_to_equal("#out", "done") + assert dash_duo.get_logs() == [] + + +def test_stst002_stream_patch_appends_once(dash_duo): + """Patch yields apply exactly once (token streaming).""" + app = Dash(__name__) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="out", children=""), + ] + ) + + @app.callback( + Output("out", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_cb(n): + yield "->" + for token in ["alpha", "beta", "gamma"]: + await asyncio.sleep(0.2) + patch = Patch() + patch += token + yield patch + + dash_duo.start_server(app) + dash_duo.find_element("#btn").click() + # Exact concatenation catches both double-apply and dropped frames. + dash_duo.wait_for_text_to_equal("#out", "->alphabetagamma") + # Give any straggler updates a chance to (incorrectly) re-apply. + time.sleep(0.5) + assert dash_duo.find_element("#out").text == "->alphabetagamma" + assert dash_duo.get_logs() == [] + + +def test_stst003_stream_multi_output_and_set_props(dash_duo): + app = Dash(__name__) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="a", children=""), + html.Div(id="b", children=""), + html.Div(id="side", children=""), + ] + ) + + @app.callback( + Output("a", "children"), + Output("b", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_cb(n): + yield "a1", no_update + await asyncio.sleep(0.3) + set_props("side", {"children": "from-set-props"}) + yield no_update, "b1" + + dash_duo.start_server(app) + dash_duo.find_element("#btn").click() + dash_duo.wait_for_text_to_equal("#a", "a1") + dash_duo.wait_for_text_to_equal("#b", "b1") + dash_duo.wait_for_text_to_equal("#side", "from-set-props") + assert dash_duo.get_logs() == [] + + +def test_stst004_stream_triggers_downstream_callback(dash_duo): + """The final streamed value triggers dependent callbacks.""" + app = Dash(__name__) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="out", children=""), + html.Div(id="downstream", children=""), + ] + ) + + @app.callback( + Output("out", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_cb(n): + yield "one" + await asyncio.sleep(0.2) + yield "two" + + @app.callback( + Output("downstream", "children"), + Input("out", "children"), + prevent_initial_call=True, + ) + def downstream(value): + return f"saw: {value}" + + dash_duo.start_server(app) + dash_duo.find_element("#btn").click() + dash_duo.wait_for_text_to_equal("#downstream", "saw: two") + assert dash_duo.get_logs() == [] + + +def test_stst005_stream_error_shows_in_devtools(dash_duo): + app = Dash(__name__) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="out", children=""), + ] + ) + + @app.callback( + Output("out", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_cb(n): + yield "before-error" + raise ValueError("stream blew up") + + dash_duo.start_server(app, debug=True, use_reloader=False, use_debugger=True) + dash_duo.find_element("#btn").click() + # The frame before the error stays applied. + dash_duo.wait_for_text_to_equal("#out", "before-error") + # And the error surfaces in the devtools error count. + dash_duo.wait_for_text_to_equal(".test-devtools-error-count", "1") + + +def test_stst006_stream_loading_state(dash_duo): + """The callback stays in loading state for the whole stream.""" + app = Dash(__name__) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="out", children="idle"), + ] + ) + + @app.callback( + Output("out", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_cb(n): + yield "working" + await asyncio.sleep(1.5) + yield "finished" + + dash_duo.start_server(app) + dash_duo.find_element("#btn").click() + dash_duo.wait_for_text_to_equal("#out", "working") + # An intermediate frame rendered but the callback is still running: + # the loading state stays on (document title shows "Updating..."). + until(lambda: dash_duo.driver.title == "Updating...", timeout=3) + assert dash_duo.redux_state_is_loading + dash_duo.wait_for_text_to_equal("#out", "finished") + # After the terminal frame the loading state clears. + until(lambda: dash_duo.driver.title != "Updating...", timeout=3) + assert not dash_duo.redux_state_is_loading + assert dash_duo.get_logs() == [] + + +def test_stst020_multiplexed_transport_over_shared_storage(dash_duo): + """Streaming over the multiplexed transport (shared storage enabled). + + Exercises the whole multiplexed path in a real browser: the renderer echoes + the signed endId, the server derives the connection id from it, pumps frames + onto that topic, and the single downlink relays them back. If the endId did + not verify end to end, the uplink would fall back to inline NDJSON (which the + stream client rejects) and the downlink would 403, so nothing would render. + Two callbacks share the one downlink, so this also covers multiplexing. + """ + from dash._shared_storage import LocalSharedStorage + + app = Dash(__name__, shared_storage=LocalSharedStorage()) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="a", children="idle"), + html.Div(id="b", children="idle"), + ] + ) + + @app.callback( + Output("a", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_a(n): + yield "a1" + await asyncio.sleep(0.4) + yield "a2" + + @app.callback( + Output("b", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_b(n): + yield "b1" + await asyncio.sleep(0.4) + yield "b2" + + dash_duo.start_server(app) + dash_duo.find_element("#btn").click() + dash_duo.wait_for_text_to_equal("#a", "a2") + dash_duo.wait_for_text_to_equal("#b", "b2") + + +def _downlink_connections(app): + """Record the endIds the downlink is polled with: the SharedWorker gives one + browser one connection (keyed on the first tab's endId), however many tabs + it has.""" + from flask import request + + seen = set() + + @app.server.before_request + def _record(): + if request.path.endswith("_dash-update-component"): + body = request.get_json(silent=True) or {} + if "streamDownlink" in body: + seen.add(request.args.get("endId")) + + return seen + + +def _open_tab(dash_duo): + dash_duo.driver.switch_to.new_window("tab") + dash_duo.driver.get(dash_duo.server_url) + + +def test_stst010_many_tabs_share_one_downlink(dash_duo): + """Browsers cap connections per host (~6); with the downlink hosted in a + SharedWorker, eight streaming tabs share one and all keep streaming.""" + app = Dash(__name__) + app.layout = html.Div([html.Div(id="out", children="idle")]) + + @app.callback(Output("out", "children"), Input("out", "id")) + async def stream_cb(_): + for i in range(120): + yield f"tick {i}" + await asyncio.sleep(0.5) + + connections = _downlink_connections(app) + dash_duo.start_server(app) + dash_duo.wait_for_contains_text("#out", "tick") + for _ in range(7): + _open_tab(dash_duo) + # Every new tab streams -- the sixth and later would stall with a + # downlink per tab. + dash_duo.wait_for_contains_text("#out", "tick") + assert len(connections) == 1 # one shared downlink for the whole browser + assert dash_duo.get_logs() == [] + + +def test_stst011_closing_a_tab_cancels_its_stream(dash_duo): + """A tab closing while others keep the shared downlink open cancels that + tab's callback server-side instead of running it to completion.""" + app = Dash(__name__) + app.layout = html.Div([html.Div(id="out", children="idle")]) + events = [] + + @app.callback(Output("out", "children"), Input("out", "id")) + async def stream_cb(_): + try: + for i in range(120): + yield f"tick {i}" + await asyncio.sleep(0.5) + except asyncio.CancelledError: + events.append("cancelled") + raise + + dash_duo.start_server(app) + dash_duo.wait_for_contains_text("#out", "tick") + first = dash_duo.driver.current_window_handle + _open_tab(dash_duo) + dash_duo.wait_for_contains_text("#out", "tick") + + dash_duo.driver.close() # the second tab goes away + dash_duo.driver.switch_to.window(first) + until(lambda: events == ["cancelled"], timeout=10) + # The surviving tab keeps streaming on the shared downlink. + before = dash_duo.find_element("#out").text + until(lambda: dash_duo.find_element("#out").text != before, timeout=5) + assert dash_duo.get_logs() == [] diff --git a/tests/streaming/test_stream_callbacks_unit.py b/tests/streaming/test_stream_callbacks_unit.py new file mode 100644 index 0000000000..60d55abe48 --- /dev/null +++ b/tests/streaming/test_stream_callbacks_unit.py @@ -0,0 +1,523 @@ +"""Unit tests for streaming (generator) callbacks - no browser required.""" +import asyncio +import contextvars +import json +import signal +import time + +import pytest + +from dash import Dash, Input, Output, Patch, callback, html, no_update, set_props +from dash._callback import GLOBAL_CALLBACK_LIST, GLOBAL_CALLBACK_MAP +from dash._stream_hub import apump_to_storage, install_stream_shutdown_handler +from dash._streaming import ( + StreamedCallbackResponse, + _keepalive_frames, + _shutdown, + andjson_lines, + keepalive_seconds, + marker_ndjson_aiter, + sync_iter_asyncgen, +) +from dash.exceptions import ( + BackgroundCallbackError, + PreventUpdate, + StreamCallbackError, +) + + +def make_body(output_id, prop, input_id="btn"): + return { + "output": f"{output_id}.{prop}", + "outputs": {"id": output_id, "property": prop}, + "inputs": [{"id": input_id, "property": "n_clicks", "value": 1}], + "changedPropIds": [f"{input_id}.n_clicks"], + } + + +def post_stream_raw(app, body): + """POST a callback request and return the raw NDJSON body.""" + client = app.server.test_client() + resp = client.post("/_dash-update-component", json=body) + assert resp.status_code == 200 + assert resp.headers.get("Content-Type") == "application/x-ndjson" + return resp.get_data(as_text=True) + + +def post_stream(app, body): + """POST a callback request and return the parsed NDJSON frames.""" + data = post_stream_raw(app, body) + return [json.loads(line) for line in data.splitlines() if line.strip()] + + +def test_stcb001_non_generator_is_not_streamed(): + @callback(Output("stcb001", "children"), Input("in", "value")) + def not_a_generator(value): + return value + + assert GLOBAL_CALLBACK_MAP["stcb001.children"]["stream"] is False + + +def test_stcb002_generator_streams_without_a_keyword(): + @callback(Output("stcb002", "children"), Input("in", "value")) + async def a_generator(value): + yield value + + assert GLOBAL_CALLBACK_MAP["stcb002.children"]["stream"] is True + + +def test_stcb003_stream_incompatible_kwargs(): + with pytest.raises(BackgroundCallbackError): + + @callback( + Output("stcb003a", "children"), + Input("in", "value"), + background=True, + ) + async def bg(value): + yield value + + with pytest.raises(StreamCallbackError, match="mcp_enabled"): + + @callback( + Output("stcb003b", "children"), + Input("in", "value"), + mcp_enabled=True, + ) + async def mcp(value): + yield value + + with pytest.raises(StreamCallbackError, match="api_endpoint"): + + @callback( + Output("stcb003c", "children"), + Input("in", "value"), + api_endpoint="/stream", + ) + async def api(value): + yield value + + +def test_stcb005_sync_generator_forbidden(): + with pytest.raises(StreamCallbackError, match="synchronous generator"): + + @callback(Output("stcb005", "children"), Input("in", "value")) + def sync_gen(value): + yield value + + +def test_stcb006_async_generator_allowed(recwarn): + @callback(Output("stcb006", "children"), Input("in", "value")) + async def async_gen(value): + yield value + + assert GLOBAL_CALLBACK_MAP["stcb006.children"]["stream"] is True + assert not [w for w in recwarn.list if issubclass(w.category, RuntimeWarning)] + + +def test_stcb007_stream_wrapper_registered(): + @callback(Output("stcb007", "children"), Input("in", "value")) + async def async_gen(value): + yield value + + assert GLOBAL_CALLBACK_MAP["stcb007.children"]["stream"] is True + + # The client spec carries a server-inferred stream flag. The client still + # detects streaming at runtime from the response (NDJSON content type / + # stream frames); the scheduler reads this flag only to keep long-lived + # streams out of its concurrent-request budget. + spec = [s for s in GLOBAL_CALLBACK_LIST if s["output"] == "stcb007.children"][-1] + assert spec["stream"] is True + + +def test_stcb008_flask_ndjson_frames(): + app = Dash(__name__) + app.layout = html.Div( + [html.Button(id="btn"), html.Div(id="out"), html.Div(id="side")] + ) + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def stream_cb(n): + yield "start" + patch = Patch() + patch += " token" + set_props("side", {"children": "side-value"}) + yield patch + yield "final" + + frames = post_stream(app, make_body("out", "children")) + assert frames[0] == {"multi": True, "response": {"out": {"children": "start"}}} + # Patch value serialized with set_props folded into the same frame, + # and cleared so it is not resent with the next frame. + assert frames[1]["sideUpdate"] == {"side": {"children": "side-value"}} + assert ( + frames[1]["response"]["out"]["children"]["__dash_patch_update"] + == "__dash_patch_update" + ) + assert frames[2] == {"multi": True, "response": {"out": {"children": "final"}}} + assert frames[3] == {"done": True} + + +def test_stcb009_stream_error_frame(): + app = Dash(__name__) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def err_cb(n): + yield "one" + raise ValueError("boom") + + frames = post_stream(app, make_body("out", "children")) + assert frames[0]["response"] == {"out": {"children": "one"}} + assert frames[1]["done"] is True + assert "boom" in frames[1]["error"]["message"] + + +def test_stcb010_stream_on_error_handler(): + app = Dash(__name__) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + + def handle(err): + return f"handled: {err}" + + @app.callback( + Output("out", "children"), + Input("btn", "n_clicks"), + on_error=handle, + ) + async def err_cb(n): + yield "one" + raise ValueError("boom") + + frames = post_stream(app, make_body("out", "children")) + assert frames[0]["response"] == {"out": {"children": "one"}} + assert frames[1]["response"] == {"out": {"children": "handled: boom"}} + assert frames[2] == {"done": True} + + +def test_stcb011_prevent_update_and_no_update_yields(): + app = Dash(__name__) + app.layout = html.Div( + [html.Button(id="btn"), html.Div(id="out"), html.Div(id="out2")] + ) + + @app.callback( + Output("out", "children"), + Output("out2", "children"), + Input("btn", "n_clicks"), + ) + async def stream_cb(n): + yield "a", no_update + yield no_update, no_update # produces no frame + yield no_update, "b" + raise PreventUpdate # ends the stream cleanly + + body = { + "output": "..out.children...out2.children..", + "outputs": [ + {"id": "out", "property": "children"}, + {"id": "out2", "property": "children"}, + ], + "inputs": [{"id": "btn", "property": "n_clicks", "value": 1}], + "changedPropIds": ["btn.n_clicks"], + } + frames = post_stream(app, body) + assert frames[0]["response"] == {"out": {"children": "a"}} + assert frames[1]["response"] == {"out2": {"children": "b"}} + assert frames[2] == {"done": True} + assert len(frames) == 3 + + +def test_stcb012_sync_iter_asyncgen(): + var = contextvars.ContextVar("stcb012") + + async def agen(): + var.set("inside") + for i in range(3): + await asyncio.sleep(0.001) + # The whole generator runs on a single task, so context set + # inside persists across steps. + assert var.get() == "inside" + yield i + + assert list(sync_iter_asyncgen(agen())) == [0, 1, 2] + + +def test_stcb013_sync_iter_asyncgen_error_propagates(): + async def agen(): + yield 1 + raise RuntimeError("kaput") + + gen = sync_iter_asyncgen(agen()) + assert next(gen) == 1 + with pytest.raises(RuntimeError, match="kaput"): + next(gen) + + +def test_stcb014_sync_iter_asyncgen_close_cancels(): + closed = [] + + async def agen(): + try: + for i in range(100): + await asyncio.sleep(0.001) + yield i + finally: + closed.append(True) + + gen = sync_iter_asyncgen(agen()) + assert next(gen) == 0 + gen.close() + # The consumer task is cancelled on a background thread; give it a moment. + for _ in range(100): + if closed: + break + time.sleep(0.01) + assert closed == [True] + + +def test_stcb015_keepalive_seconds_normalization(): + assert keepalive_seconds(15000) == 15.0 + assert keepalive_seconds(None) is None + assert keepalive_seconds(0) is None + assert keepalive_seconds(-1) is None + + +def test_stcb016_flask_keepalive_between_slow_yields(): + app = Dash(__name__, stream_keepalive_interval=50) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def stream_cb(n): + await asyncio.sleep(0.3) + yield "start" + await asyncio.sleep(0.3) + yield "final" + + raw = post_stream_raw(app, make_body("out", "children")) + # Blank keepalive lines while the callback is between yields. + assert len([line for line in raw.splitlines() if not line.strip()]) >= 2 + # The frames themselves are unaffected. + frames = [json.loads(line) for line in raw.splitlines() if line.strip()] + assert frames[0]["response"] == {"out": {"children": "start"}} + assert frames[1]["response"] == {"out": {"children": "final"}} + assert frames[2] == {"done": True} + + +def test_stcb017_flask_keepalive_disabled(): + app = Dash(__name__, stream_keepalive_interval=None) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def stream_cb(n): + await asyncio.sleep(0.2) + yield "only" + + raw = post_stream_raw(app, make_body("out", "children")) + assert [line for line in raw.splitlines() if not line.strip()] == [] + + +def test_stcb018_async_keepalive_does_not_cancel_source(): + async def agen(): + await asyncio.sleep(0.3) + yield {"multi": True} + await asyncio.sleep(0.3) + yield {"done": True} + + async def collect(): + return [line async for line in andjson_lines(agen(), keepalive=0.05)] + + lines = asyncio.run(collect()) + assert lines.count("\n") >= 2 + # Holding the pending __anext__ across keepalives means both frames still + # arrive; a bare wait_for would have cancelled the generator mid-step. + assert [json.loads(line) for line in lines if line.strip()] == [ + {"multi": True}, + {"done": True}, + ] + + +def test_stcb020_async_keepalive_over_sync_generator(): + """marker_ndjson_aiter with is_async=False: sync generator, ASGI backend.""" + + def frames(): + time.sleep(0.3) + yield {"multi": True} + yield {"done": True} + + marker = StreamedCallbackResponse( + frames(), is_async=False, ctx=contextvars.copy_context() + ) + + async def collect(): + return [line async for line in marker_ndjson_aiter(marker, keepalive=0.05)] + + lines = asyncio.run(collect()) + assert lines.count("\n") >= 2 + assert [json.loads(line) for line in lines if line.strip()] == [ + {"multi": True}, + {"done": True}, + ] + + +def test_stcb019_keepalive_frames_closes_generator_when_consumer_leaves(): + closed = [] + + def frames(): + try: + while True: + yield {"multi": True} + finally: + closed.append(True) + + marker = StreamedCallbackResponse( + frames(), is_async=False, ctx=contextvars.copy_context() + ) + gen = _keepalive_frames(marker, 0.05) + assert next(gen) == {"multi": True} + gen.close() + # The pump thread owns the generator, so cleanup happens once it notices + # the stop flag rather than at the consumer's close(). + for _ in range(200): + if closed: + break + time.sleep(0.01) + assert closed == [True] + + +def test_stcb021_shutdown_flag_stops_keepalive_generator(): + """_shutdown event makes _keepalive_frames exit within one poll cycle.""" + from dash._streaming import _shutdown + + _shutdown.clear() + + def frames(): + while True: + yield {"multi": True} + time.sleep(0.05) + + marker = StreamedCallbackResponse( + frames(), is_async=False, ctx=contextvars.copy_context() + ) + gen = _keepalive_frames(marker, keepalive=60) + assert next(gen) == {"multi": True} + + _shutdown.set() + t0 = time.monotonic() + remaining = list(gen) + elapsed = time.monotonic() - t0 + _shutdown.clear() + + assert elapsed < 2, f"generator took {elapsed:.1f}s to stop (expected <2s)" + assert len(remaining) <= 2 + + +def test_stcb022_shutdown_active_streams_sets_flag_and_closes_subs(): + """shutdown_active_streams sets the _shutdown flag and closes subs.""" + from dash._streaming import _shutdown + from dash._stream_hub import ( + _active_subscriptions, + _registry_lock, + shutdown_active_streams, + ) + + _shutdown.clear() + + class FakeSub: + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + sub = FakeSub() + with _registry_lock: + _active_subscriptions.add(sub) + + try: + shutdown_active_streams() + assert _shutdown.is_set() + assert sub.closed + finally: + _shutdown.clear() + with _registry_lock: + _active_subscriptions.discard(sub) + + +def test_stcb023_install_shutdown_handler_wraps_current_handler(): + """Installing over a foreign handler sets the flag, then chains to it.""" + saved = {sig: signal.getsignal(sig) for sig in (signal.SIGINT, signal.SIGTERM)} + calls = [] + + def foreign(sig, _frame): + calls.append((sig, _shutdown.is_set())) + + try: + signal.signal(signal.SIGINT, foreign) + signal.signal(signal.SIGTERM, foreign) + _shutdown.clear() + + install_stream_shutdown_handler() + installed = signal.getsignal(signal.SIGINT) + assert installed is not foreign + + installed(signal.SIGINT, None) + assert _shutdown.is_set() + assert calls == [(signal.SIGINT, True)] + + install_stream_shutdown_handler() + assert signal.getsignal(signal.SIGINT) is installed + assert signal.getsignal(signal.SIGTERM) is not foreign + finally: + _shutdown.clear() + for sig, handler in saved.items(): + signal.signal(sig, handler) + + +def test_stcb024_cancelled_pump_publishes_terminal_error(): + """A pump cancelled by shutdown leaves a terminal error frame in the store.""" + published = [] + + class FakeStorage: + def publish(self, topic, message): + published.append((topic, message)) + + # The pump talks to the store through its loop-native methods. + async def apublish(self, topic, message): + self.publish(topic, message) + + async def aget(self, key, default=None): + return default + + async def adelete(self, key): + pass + + async def frames(): + yield {"multi": True, "response": {"out": {"children": 1}}} + await asyncio.sleep(10) + yield {"done": True} + + async def scenario(): + marker = StreamedCallbackResponse( + frames(), is_async=True, ctx=contextvars.copy_context() + ) + task = asyncio.ensure_future( + apump_to_storage(FakeStorage(), "conn", "rid", marker) + ) + await asyncio.sleep(0.05) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(scenario()) + assert [m["frame"] for _, m in published] == [ + {"multi": True, "response": {"out": {"children": 1}}}, + { + "done": True, + "error": { + "message": "Streaming callback interrupted: " + "the server shut down while it was running" + }, + }, + ] + assert all(m["rid"] == "rid" for _, m in published) diff --git a/tests/streaming/test_stream_shutdown.py b/tests/streaming/test_stream_shutdown.py new file mode 100644 index 0000000000..ac1f5a7b1b --- /dev/null +++ b/tests/streaming/test_stream_shutdown.py @@ -0,0 +1,230 @@ +"""Ctrl+C must stop a real server while a streaming callback is running. + +Drives the same process shape ``app.run()`` spawns for the FastAPI backend +(``python -m uvicorn``): the server's graceful shutdown waits for in-flight +responses, so if the streaming shutdown hook is not wired in, the process +ignores SIGINT until the generator ends. +""" +import os +import signal +import socket +import subprocess +import sys +import textwrap +import threading +import time + +import pytest +import requests + +from dash.testing.wait import until + +pytestmark = pytest.mark.skipif( + sys.platform == "win32", reason="POSIX signals and process groups" +) + +APP = textwrap.dedent( + """ + import asyncio + from dash import Dash, Input, Output, html + + app = Dash(__name__, backend="{backend}") + server = app.server + 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"), + prevent_initial_call=True) + async def stream(n): + for i in range(10000): + yield f"token {{i}}" + await asyncio.sleep(0.2) + """ +) + +CALLBACK_BODY = { + "output": "out.children", + "outputs": {"id": "out", "property": "children"}, + "inputs": [{"id": "btn", "property": "n_clicks", "value": 1}], + "changedPropIds": ["btn.n_clicks"], + "state": [], +} + + +def _free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _wait_ready(url, proc, timeout=30): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if proc.poll() is not None: + raise AssertionError(f"server exited early: {proc.stderr.read()}") + try: + if requests.get(url, timeout=1).status_code == 200: + return + except requests.RequestException: + time.sleep(0.2) + raise AssertionError("server never came up") + + +@pytest.mark.parametrize("backend", ["fastapi", "quart"]) +def test_stsd001_sigint_stops_server_mid_stream(tmp_path, backend): + pytest.importorskip(backend) + (tmp_path / "shutdown_app.py").write_text(APP.format(backend=backend)) + port = _free_port() + proc = subprocess.Popen( # pylint: disable=consider-using-with + [sys.executable, "-m", "uvicorn", "shutdown_app:server", "--port", str(port)], + cwd=tmp_path, + env=dict(os.environ, PYTHONUNBUFFERED="1"), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + start_new_session=True, + ) + frames = [] + try: + base = f"http://127.0.0.1:{port}" + _wait_ready(base, proc) + + def read_stream(): + with requests.post( + f"{base}/_dash-update-component", + json=CALLBACK_BODY, + stream=True, + timeout=30, + ) as resp: + try: + for line in resp.iter_lines(): + frames.append(line) + except requests.RequestException: + pass + + reader = threading.Thread(target=read_stream, daemon=True) + reader.start() + deadline = time.monotonic() + 10 + while len(frames) < 2 and time.monotonic() < deadline: + time.sleep(0.1) + assert len(frames) >= 2, "stream never started" + + proc.send_signal(signal.SIGINT) + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + raise AssertionError( # pylint: disable=raise-missing-from + "server still running 10s after SIGINT with an active stream" + ) + output = proc.stdout.read() + assert "Application shutdown complete" in output, output + finally: + if proc.poll() is None: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + + +MUX_APP = textwrap.dedent( + """ + import asyncio + from dash import Dash, Input, Output, Patch, html + + app = Dash(__name__, backend="fastapi") + server = app.server + 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"), + running=[(Output("btn", "disabled"), True, False)], + prevent_initial_call=True) + async def stream(n): + for i in range(600): + patch = Patch() + patch.append(f"t{i} ") + yield patch + await asyncio.sleep(0.2) + """ +) + +COUNT_DOWNLINKS = """ +if (!window.__downlinks) { + window.__downlinks = 0; + const orig = window.fetch; + window.fetch = function(url, init) { + if (init && init.body && String(init.body).includes('streamDownlink')) { + window.__downlinks += 1; + } + return orig.apply(this, arguments); + }; +} +return window.__downlinks; +""" + + +def _start_server(tmp_path, port): + proc = subprocess.Popen( # pylint: disable=consider-using-with + [sys.executable, "-m", "uvicorn", "restart_app:server", "--port", str(port)], + cwd=tmp_path, + env=dict(os.environ, PYTHONUNBUFFERED="1"), + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + try: + _wait_ready(f"http://127.0.0.1:{port}", proc) + except BaseException: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + raise + return proc + + +def _stop_server(proc): + if proc.poll() is None: + proc.send_signal(signal.SIGINT) + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() + + +def test_stsd002_restart_settles_multiplexed_streams(dash_br, tmp_path): + """A server restart mid-stream must not leave the page looping. + + The default shared storage puts streams on the multiplexed downlink. When + the server restarts, its signing secret changes, so the old page's + downlink is refused: the client must settle the callbacks it was running + (clearing the running state) and stop reconnecting, and a refreshed page + must stream again normally. + """ + (tmp_path / "restart_app.py").write_text(MUX_APP) + port = _free_port() + proc = _start_server(tmp_path, port) + try: + dash_br.server_url = f"http://127.0.0.1:{port}" + dash_br.driver.execute_script(COUNT_DOWNLINKS) + dash_br.find_element("#btn").click() + dash_br.wait_for_contains_text("#out", "t2") + assert dash_br.find_element("#btn").get_attribute("disabled") + + _stop_server(proc) + proc = _start_server(tmp_path, port) + + until( + lambda: not dash_br.find_element("#btn").get_attribute("disabled"), + timeout=15, + ) + # Frames applied before the drop stay on the page. + assert "t2" in dash_br.find_element("#out").text + # No reconnect loop against the refused connection. + count = dash_br.driver.execute_script(COUNT_DOWNLINKS) + time.sleep(3) + assert dash_br.driver.execute_script(COUNT_DOWNLINKS) == count + + # Drain the connection-refused entries logged while the server was down. + dash_br.get_logs() + dash_br.driver.refresh() + dash_br.wait_for_element("#btn").click() + dash_br.wait_for_contains_text("#out", "t2") + assert dash_br.get_logs() == [] + finally: + _stop_server(proc) diff --git a/tests/streaming/test_stream_transport.py b/tests/streaming/test_stream_transport.py new file mode 100644 index 0000000000..458cb93fda --- /dev/null +++ b/tests/streaming/test_stream_transport.py @@ -0,0 +1,600 @@ +"""Multiplexed streaming transport over shared storage (server side). + +The uplink: a streaming callback POST that carries a streamConnection returns a +fast ack and pumps its frames onto the connection's shared-storage topic (from +which the client's single downlink relays them). Exercised over the real HTTP +dispatch on all three backends (Flask WSGI, Quart + FastAPI ASGI). +""" +import asyncio +import json +import threading +import time +import uuid + +import pytest + +from dash import Dash, Input, Output, html +from dash import _callback_signing +from dash._shared_storage import LocalSharedStorage +from dash._stream_hub import subscribe_envelopes + +# The connection topic is keyed on the server-signed end_id, not on anything the +# client sends. A test uplink/downlink must carry a validly signed endId (the +# raw value becomes the connection id / topic) or the server refuses to +# multiplex it. See dash/_callback.get_stream_connection_id. +CONNECTION_ID = "conn-test" + + +def _signed_end_id(app): + secret = app._get_signing_secret() # pylint: disable=protected-access + return _callback_signing.sign(secret, _callback_signing.END_SCOPE, CONNECTION_ID) + + +def _uplink_url(app): + return f"/_dash-update-component?endId={_signed_end_id(app)}" + + +def _uplink_body(request_id): + return { + "output": "out.children", + "outputs": {"id": "out", "property": "children"}, + "inputs": [{"id": "btn", "property": "n_clicks", "value": 1}], + "changedPropIds": ["btn.n_clicks"], + "streamConnection": {"requestId": request_id}, + } + + +def _start_drain(storage, connection_id, out): + """Subscribe to a connection's topic on a daemon thread until 'done'.""" + + def drain(): + gen = subscribe_envelopes(storage, connection_id) + for env in gen: + out.append(env) + if env["frame"].get("done"): + break + gen.close() + + th = threading.Thread(target=drain, daemon=True) + th.start() + time.sleep(0.3) # subscription established before the pump publishes + return th + + +def _streaming_app(server=None): + storage = LocalSharedStorage(namespace=f"tx-{uuid.uuid4().hex[:8]}") + kwargs = {"shared_storage": storage} + if server is not None: + kwargs["server"] = server # default (Flask) server otherwise + app = Dash(__name__, **kwargs) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def cb(n): + yield "a" + yield "b" + + return app, storage + + +def _assert_delivered(out): + assert [e["rid"] for e in out] == ["r1", "r1", "r1"] + frames = [e["frame"] for e in out] + assert frames[0]["response"] == {"out": {"children": "a"}} + assert frames[1]["response"] == {"out": {"children": "b"}} + assert frames[2] == {"done": True} + + +def test_flask_uplink_pumps_callback_frames_to_storage(): + app, storage = _streaming_app() + out = [] + th = _start_drain(storage, CONNECTION_ID, out) + + resp = app.server.test_client().post(_uplink_url(app), json=_uplink_body("r1")) + assert resp.status_code == 200 + assert json.loads(resp.get_data(as_text=True)) == {"multi": True, "stream": True} + + th.join(timeout=5) + _assert_delivered(out) + storage.close() + + +def test_flask_uplink_without_valid_end_id_is_rejected(): + # A multiplexed uplink (carries a streamConnection) whose endId does not + # verify is refused outright: it is never run some other way, so no frame is + # ever published onto a topic without a valid token. + app, storage = _streaming_app() + out = [] + th = _start_drain(storage, CONNECTION_ID, out) + + resp = app.server.test_client().post( + "/_dash-update-component?endId=forged~deadbeef", json=_uplink_body("r1") + ) + assert resp.status_code == 403 + th.join(timeout=1) + assert out == [] + storage.close() + + +def test_flask_downlink_rejects_missing_end_id(): + # A downlink with no valid signed endId cannot name a topic at all: the + # server refuses it (403) rather than serving an attacker-named connection. + app, storage = _streaming_app() + resp = app.server.test_client().post( + "/_dash-update-component", json={"streamDownlink": {"from": 0}} + ) + assert resp.status_code == 403 + storage.close() + + +def test_flask_downlink_resets_a_stale_cursor(): + # A downlink resuming from a cursor the fresh topic never reached (the page's + # server restarted, so the topic is back at seq 0) gets a reset line, not a + # silent stall until the new sequence climbs past the stale cursor. + app, storage = _streaming_app() + resp = app.server.test_client().post( + _uplink_url(app), json={"streamDownlink": {"from": 99}} + ) + assert resp.status_code == 200 + lines = [line for line in resp.get_data(as_text=True).splitlines() if line.strip()] + assert json.loads(lines[0]) == {"reset": True} + storage.close() + + +def test_fastapi_uplink_pumps_callback_frames_to_storage(): + pytest.importorskip("httpx", reason="fastapi.testclient requires httpx") + from fastapi import FastAPI + from fastapi.testclient import TestClient + + server = FastAPI() + app, storage = _streaming_app(server=server) + app._setup_server() # pylint: disable=protected-access + out = [] + th = _start_drain(storage, CONNECTION_ID, out) + + with TestClient(server) as client: + resp = client.post(_uplink_url(app), json=_uplink_body("r1")) + assert resp.status_code == 200 + assert resp.json() == {"multi": True, "stream": True} + th.join(timeout=8) + + _assert_delivered(out) + storage.close() + + +def test_quart_uplink_pumps_callback_frames_to_storage(): + quart = pytest.importorskip("quart") + + server = quart.Quart(__name__) + app, storage = _streaming_app(server=server) + app._setup_server() # pylint: disable=protected-access + out = [] + th = _start_drain(storage, CONNECTION_ID, out) + + async def run(): + client = server.test_client() + resp = await client.post(_uplink_url(app), json=_uplink_body("r1")) + assert resp.status_code == 200 + assert await resp.get_json() == {"multi": True, "stream": True} + # Keep the loop alive so the fire-and-forget pump task delivers. + for _ in range(100): + if len(out) >= 3: + break + await asyncio.sleep(0.05) + + asyncio.run(run()) + th.join(timeout=5) + _assert_delivered(out) + storage.close() + + +# --- downlink lifecycle: tab close stops the relay and cancels the pumps ------ + + +def _storage(tag): + storage = LocalSharedStorage(namespace=f"tx-{tag}-{uuid.uuid4().hex[:8]}") + storage.start() + return storage + + +def _wait_for(pred, timeout=5.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return True + time.sleep(0.02) + return pred() + + +def test_downlink_gone_semantics(monkeypatch): + from dash import _stream_hub as hub + + storage = _storage("gone") + cid = "c1" + started = time.time() + # No record yet (downlink open racing the uplink): not gone until grace. + assert not hub.downlink_gone(storage, cid, started, grace=1.0) + assert hub.downlink_gone(storage, cid, started - 2.0, grace=1.0) + + downlink = hub.Downlink(storage, cid) + 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 + # fresh downlink for its new stream, so the pump gets the full grace. + monkeypatch.setattr(hub.time, "time", lambda: started + 0.5) + assert not hub.downlink_gone(storage, cid, started, grace=1.0) + monkeypatch.setattr(hub.time, "time", lambda: started + 1.5) + assert hub.downlink_gone(storage, cid, started, grace=1.0) + + +def test_replaced_downlink_does_not_mark_the_new_one_closed(): + from dash import _stream_hub as hub + + storage = _storage("replace") + cid = "c2" + old = hub.Downlink(storage, cid) + new = hub.Downlink(storage, cid, replay_from=0) # the client reconnected + 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 _cancellation_probe(): + """An endless async frame generator that records its cancellation.""" + state = {"cancelled": False, "frames": 0} + + async def frames(): + try: + while True: + state["frames"] += 1 + yield {"multi": True, "response": {"out": {"children": "x"}}} + await asyncio.sleep(0.05) + except asyncio.CancelledError: + state["cancelled"] = True + raise + + return state, frames() + + +def test_sync_pump_cancels_callback_when_downlink_gone(monkeypatch): + from dash import _stream_hub as hub + from dash._streaming import StreamedCallbackResponse + + monkeypatch.setattr(hub, "DOWNLINK_GRACE", 0.3) + monkeypatch.setattr(hub, "DOWNLINK_CHECK_INTERVAL", 0.1) + storage = _storage("pump") + cid = "c4" + out = [] + downlink = hub.Downlink(storage, cid) + + def drain(): + for env in downlink.envelopes(): + out.append(env) + if env["frame"].get("done"): + break + + drain_th = threading.Thread(target=drain, daemon=True) + drain_th.start() + + state, frames = _cancellation_probe() + marker = StreamedCallbackResponse(frames, is_async=True) + pump = hub.pump_to_storage(storage, cid, "r1", marker) + assert _wait_for(lambda: len(out) >= 3) + assert not state["cancelled"] + + downlink.close() # the tab closed + pump.result(timeout=5) + assert state["cancelled"] + # Once the frames stop, the pump published a terminal frame so a client + # that reconnects late resolves the request instead of waiting forever. + with storage.subscribe(hub.stream_topic(cid), replay_from=0) as sub: + got = [] + for _seq, message in sub.iter_with_seq(): + got.append(message["frame"]) + if message["frame"].get("done"): + break + assert got[-1] == {"done": True} + assert all(f.get("multi") for f in got[:-1]) + + +def test_async_pump_cancels_callback_when_downlink_gone(monkeypatch): + from dash import _stream_hub as hub + from dash._streaming import StreamedCallbackResponse + + monkeypatch.setattr(hub, "DOWNLINK_GRACE", 0.3) + monkeypatch.setattr(hub, "DOWNLINK_CHECK_INTERVAL", 0.1) + storage = _storage("apump") + cid = "c5" + downlink = hub.Downlink(storage, cid) + + async def scenario(): + state, frames = _cancellation_probe() + marker = StreamedCallbackResponse(frames, is_async=True) + task = asyncio.ensure_future(hub.apump_to_storage(storage, cid, "r1", marker)) + while state["frames"] < 3: + await asyncio.sleep(0.02) + assert not state["cancelled"] + downlink.close() # the tab closed + await asyncio.wait_for(task, timeout=5) + return state + + state = asyncio.run(scenario()) + assert state["cancelled"] + with storage.subscribe(hub.stream_topic(cid), replay_from=0) as sub: + got = [] + for _seq, message in sub.iter_with_seq(): + got.append(message["frame"]) + if message["frame"].get("done"): + break + assert got[-1] == {"done": True} + + +# --- explicit cancellation (a tab closed while the downlink stays shared) ----- + + +def test_stream_cancel_stops_sync_pump_while_downlink_stays_open(monkeypatch): + from dash import _stream_hub as hub + from dash._streaming import StreamedCallbackResponse + + monkeypatch.setattr(hub, "DOWNLINK_CHECK_INTERVAL", 0.1) + storage = _storage("cancel-pump") + cid = "c6" + downlink = hub.Downlink(storage, cid) # stays open: other tabs still stream + state, frames = _cancellation_probe() + marker = StreamedCallbackResponse(frames, is_async=True) + pump = hub.pump_to_storage(storage, cid, "r1", marker) + assert _wait_for(lambda: state["frames"] >= 2) + + hub.cancel_stream(storage, cid, "r1") + pump.result(timeout=5) + assert state["cancelled"] + assert storage.get(hub.connection_key(cid))["open"] is True + # The pump cleans up its cancel record once it has acted on it. + assert storage.get(hub.cancel_key(cid, "r1")) is None + downlink.close() + + +_CANCEL_BODY = {"streamCancel": {"requestId": "r9"}} + + +def _assert_cancel_recorded(storage, status, data): + from dash import _stream_hub as hub + + assert status == 200 + assert data == hub.STREAM_CANCEL_ACK + assert hub.stream_cancelled(storage, CONNECTION_ID, "r9") + + +def test_flask_stream_cancel_endpoint_records_the_request(): + app, storage = _streaming_app() + client = app.server.test_client() + response = client.post(_uplink_url(app), json=_CANCEL_BODY) + _assert_cancel_recorded(storage, response.status_code, response.get_json()) + # Without a valid signed endId a cancel cannot name a connection at all. + assert client.post("/_dash-update-component", json=_CANCEL_BODY).status_code == 403 + storage.close() + + +def test_fastapi_stream_cancel_endpoint_records_the_request(): + pytest.importorskip("httpx", reason="fastapi.testclient requires httpx") + from fastapi import FastAPI + from fastapi.testclient import TestClient + + server = FastAPI() + app, storage = _streaming_app(server=server) + app._setup_server() # pylint: disable=protected-access + with TestClient(server) as client: + response = client.post(_uplink_url(app), json=_CANCEL_BODY) + _assert_cancel_recorded(storage, response.status_code, response.json()) + storage.close() + + +def test_quart_stream_cancel_endpoint_records_the_request(): + quart = pytest.importorskip("quart") + + server = quart.Quart(__name__) + app, storage = _streaming_app(server=server) + app._setup_server() # pylint: disable=protected-access + + async def run(): + resp = await server.test_client().post(_uplink_url(app), json=_CANCEL_BODY) + return resp.status_code, await resp.get_json() + + status, data = asyncio.run(run()) + _assert_cancel_recorded(storage, status, data) + storage.close() + + +# --- process shutdown: Ctrl+C ends live downlinks and pumps ------------------- + + +def test_shutdown_active_streams_ends_live_downlinks_and_pumps(): + from dash import _stream_hub as hub + from dash._streaming import StreamedCallbackResponse + + storage = _storage("shutdown") + downlink = hub.Downlink(storage, "c8") + relayed = [] + + def relay(): + for env in downlink.envelopes(): + relayed.append(env) + + relay_th = threading.Thread(target=relay, daemon=True) + relay_th.start() + + async def scenario(): + state, frames = _cancellation_probe() + marker = StreamedCallbackResponse(frames, is_async=True) + hub.spawn_async_pump(storage, "c8", "r1", marker) + while state["frames"] < 2: + await asyncio.sleep(0.02) + hub.shutdown_active_streams() # what the SIGINT hook does + await asyncio.sleep(0.3) + return state + + state = asyncio.run(scenario()) + assert state["cancelled"] + relay_th.join(timeout=3) + assert not relay_th.is_alive() # the downlink response ended + assert storage.get(hub.connection_key("c8"))["open"] is False + + +# --- polling downlink (WSGI) -------------------------------------------------- + + +def test_subscription_poll_is_non_blocking_and_advances_the_cursor(): + storage = _storage("poll") + sub = storage.subscribe("t", replay_from=0) + started = time.monotonic() + assert sub.poll(0.0) == [] + assert time.monotonic() - started < 0.2 + storage.publish("t", {"n": 1}) + storage.publish("t", {"n": 2}) + assert sub.poll(0.0) == [(1, {"n": 1}), (2, {"n": 2})] + assert sub.poll(0.0) == [] # cursor advanced past what was delivered + sub.close() + + +def test_poll_downlink_returns_queued_frames_and_heartbeats(monkeypatch): + from dash import _stream_hub as hub + + storage = _storage("polldl") + cid = "c9" + assert hub.poll_downlink(storage, cid, 0) == [] + record = storage.get(hub.connection_key(cid)) + assert record["mode"] == "poll" and record["open"] is True + first_beat = record["at"] + + hub.publish_frame(storage, cid, "r1", {"multi": True, "response": {"a": 1}}) + hub.publish_frame(storage, cid, "r1", {"multi": True, "response": {"a": 2}}) + envelopes = hub.poll_downlink(storage, cid, 0) + assert [e["frame"]["response"] for e in envelopes] == [{"a": 1}, {"a": 2}] + assert [e["seq"] for e in envelopes] == [1, 2] + # Resuming from the cursor yields only what came after it. + hub.publish_frame(storage, cid, "r1", {"multi": True, "response": {"a": 3}}) + assert [e["seq"] for e in hub.poll_downlink(storage, cid, 2)] == [3] + # Heartbeats are throttled: many polls a second, one record write. + assert storage.get(hub.connection_key(cid))["at"] == first_beat + monkeypatch.setattr(hub.time, "time", lambda: first_beat + 2.0) + hub.poll_downlink(storage, cid, 3) + assert storage.get(hub.connection_key(cid))["at"] == first_beat + 2.0 + monkeypatch.undo() + + # A polling browser counts as present while its heartbeat is fresh, and as + # gone once it stops polling for longer than POLL_GRACE -- wider than the + # closed-downlink grace, since a loaded pool can delay polls for seconds. + beat = storage.get(hub.connection_key(cid))["at"] + assert not hub.downlink_gone(storage, cid, beat - 100) + monkeypatch.setattr(hub.time, "time", lambda: beat + hub.DOWNLINK_GRACE + 5.0) + assert not hub.downlink_gone(storage, cid, beat - 100) + monkeypatch.setattr(hub.time, "time", lambda: beat + hub.POLL_GRACE + 1.0) + assert hub.downlink_gone(storage, cid, beat - 100) + # ...but a pump that just started gives it the grace to resume polling. + assert not hub.downlink_gone(storage, cid, beat + hub.POLL_GRACE) + + +def test_flask_downlink_poll_returns_at_once_with_the_queued_frames(): + from dash import _stream_hub as hub + + app, storage = _streaming_app() + client = app.server.test_client() + hub.publish_frame( + storage, CONNECTION_ID, "r1", {"multi": True, "response": {"a": 1}} + ) + + started = time.monotonic() + response = client.post(_uplink_url(app), json={"streamDownlink": {"from": 0}}) + assert time.monotonic() - started < 1.0 # no waiting: a poll, not a stream + assert response.status_code == 200 + assert response.content_type.startswith("application/x-ndjson") + lines = [line for line in response.get_data(as_text=True).split("\n") if line] + envelopes = [json.loads(line) for line in lines] + assert [e["frame"]["response"] for e in envelopes] == [{"a": 1}] + cursor = envelopes[-1]["seq"] + + # Nothing new: an empty body, still at once. + response = client.post(_uplink_url(app), json={"streamDownlink": {"from": cursor}}) + assert response.get_data(as_text=True) == "" + assert storage.get(hub.connection_key(CONNECTION_ID))["mode"] == "poll" + storage.close() + + +def test_wsgi_pumps_share_one_loop_thread(): + """Many streams on a WSGI worker cost one pump thread, not one each.""" + from dash import _stream_hub as hub + from dash._streaming import StreamedCallbackResponse + + storage = _storage("pumploop") + hub.Downlink(storage, "c10") # a browser is present + pumps = [] + for i in range(25): + _state, frames = _cancellation_probe() + marker = StreamedCallbackResponse(frames, is_async=True) + pumps.append(hub.pump_to_storage(storage, "c10", f"r{i}", marker)) + assert _wait_for( + lambda: storage.subscribe(hub.stream_topic("c10"), 0).poll(0.0) != [] + ) + names = [t.name for t in threading.enumerate() if t.name.startswith("dash-stream")] + assert names.count("dash-stream-pumps") == 1 + assert "dash-stream-bridge" not in names and "dash-stream-pump" not in names + for i in range(25): + hub.cancel_stream(storage, "c10", f"r{i}") + for pump in pumps: + pump.result(timeout=5) + + +def test_wsgi_pump_loop_survives_an_exception_raised_into_it(): + """The pump loop carries every stream in the process: an exception raised + into its thread (a harness that stops every thread an app started) must + not end it -- and if the thread is gone anyway, the next pump gets a + fresh loop instead of being scheduled onto a dead one.""" + import ctypes + + from dash import _stream_hub as hub + from dash._streaming import StreamedCallbackResponse + + storage = _storage("pumprestart") + hub.Downlink(storage, "c11") + hub._shared_pump_loop() # pylint: disable=protected-access + thread = hub._pump_thread # pylint: disable=protected-access + _state, frames = _cancellation_probe() + pump = hub.pump_to_storage( + storage, "c11", "r1", StreamedCallbackResponse(frames, is_async=True) + ) + assert _wait_for( + lambda: storage.subscribe(hub.stream_topic("c11"), 0).poll(0.0) != [] + ) + # What dash.testing's KillerThread does to "new" threads at teardown. + assert ( + ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_long(thread.ident), ctypes.py_object(SystemExit) + ) + == 1 + ) + time.sleep(0.3) + assert thread.is_alive() # shrugged it off + _state2, frames2 = _cancellation_probe() + pump2 = hub.pump_to_storage( + storage, "c11", "r2", StreamedCallbackResponse(frames2, is_async=True) + ) + assert _wait_for(lambda: _state2["frames"] >= 2) + for rid in ("r1", "r2"): + hub.cancel_stream(storage, "c11", rid) + pump2.result(timeout=5) + pump.result(timeout=5) + + # And if the thread is truly gone, the loop is replaced. + loop = hub._pump_loop # pylint: disable=protected-access + loop.call_soon_threadsafe(loop.stop) + thread.join(timeout=5) + assert not thread.is_alive() + _state3, frames3 = _cancellation_probe() + pump3 = hub.pump_to_storage( + storage, "c11", "r3", StreamedCallbackResponse(frames3, is_async=True) + ) + assert _wait_for(lambda: _state3["frames"] >= 1) + assert hub._pump_thread is not thread # pylint: disable=protected-access + hub.cancel_stream(storage, "c11", "r3") + pump3.result(timeout=5) diff --git a/tests/streaming/test_stream_wsgi.py b/tests/streaming/test_stream_wsgi.py new file mode 100644 index 0000000000..bf5b616a6c --- /dev/null +++ b/tests/streaming/test_stream_wsgi.py @@ -0,0 +1,96 @@ +"""Streaming over the multiplexed transport on a single-threaded WSGI worker. + +gunicorn's default sync worker serves one request at a time. The client must +open its long-lived downlink only once the uplink is acknowledged, otherwise +the downlink holds the only worker and the stream never starts. +""" +import os +import signal +import socket +import subprocess +import sys +import textwrap +import time + +import pytest +import requests + +pytestmark = pytest.mark.skipif( + sys.platform == "win32", reason="gunicorn is POSIX only" +) +pytest.importorskip("gunicorn") + +APP = textwrap.dedent( + """ + import asyncio + from dash import Dash, Input, Output, html + + app = Dash(__name__) + server = app.server + 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.2) + """ +) + + +def _free_port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def test_stwg001_single_sync_worker_streams_promptly(dash_br, tmp_path): + (tmp_path / "wsgi_app.py").write_text(APP) + port = _free_port() + proc = subprocess.Popen( # pylint: disable=consider-using-with + [ + sys.executable, + "-m", + "gunicorn", + "wsgi_app:server", + "--bind", + f"127.0.0.1:{port}", + "--workers", + "1", + "--timeout", + "30", + ], + cwd=tmp_path, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + try: + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + try: + if requests.get(f"http://127.0.0.1:{port}", timeout=1).ok: + break + except requests.RequestException: + time.sleep(0.2) + else: + raise AssertionError("gunicorn never came up") + + dash_br.server_url = f"http://127.0.0.1:{port}" + dash_br.find_element("#btn").click() + started = time.monotonic() + # Well under gunicorn's worker timeout: the stream must not need the + # worker to be killed and respawned before it starts. + dash_br.wait_for_contains_text("#out", "token 1", timeout=10) + assert time.monotonic() - started < 10 + dash_br.wait_for_contains_text("#out", "token 19", timeout=15) + assert dash_br.get_logs() == [] + finally: + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait() diff --git a/tests/websocket/test_ws_stream.py b/tests/websocket/test_ws_stream.py new file mode 100644 index 0000000000..262cdd3d79 --- /dev/null +++ b/tests/websocket/test_ws_stream.py @@ -0,0 +1,151 @@ +"""WebSocket streaming callback tests. + +Protocol-level tests (FastAPI TestClient, no browser) verifying that +streaming callbacks emit intermediate callback_response frames with +``stream: true`` followed by a terminal done frame, plus browser tests for the +full renderer round-trip. +""" +import asyncio +import json + +import pytest + +from dash import Dash, Input, Output, Patch, html + + +def _collect_stream_messages(ws): + """Read ws messages, flattening batched arrays, until the terminal frame.""" + out = [] + while True: + parsed = json.loads(ws.receive_text()) + msgs = parsed if isinstance(parsed, list) else [parsed] + for msg in msgs: + if msg.get("type") != "callback_response": + continue + out.append(msg) + payload = msg.get("payload") or {} + if payload.get("done") or not payload.get("stream"): + return out + + +def _make_ws_app(): + from fastapi import FastAPI + + server = FastAPI() + app = Dash(__name__, server=server, websocket_callbacks=True) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + return app, server + + +def _callback_request(request_id, output_id="out", prop="children"): + return { + "type": "callback_request", + "requestId": request_id, + "rendererId": "rend1", + "payload": { + "output": f"{output_id}.{prop}", + "outputs": {"id": output_id, "property": prop}, + "inputs": [{"id": "btn", "property": "n_clicks", "value": 1}], + "changedPropIds": ["btn.n_clicks"], + }, + } + + +def test_wsst001_async_stream_frames_over_ws(): + pytest.importorskip("httpx", reason="fastapi.testclient requires httpx") + from fastapi.testclient import TestClient + + app, server = _make_ws_app() + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def stream_cb(n): + yield "start" + await asyncio.sleep(0.01) + yield "final" + + app._setup_server() + + client = TestClient(server) + with client.websocket_connect( + "/_dash-ws-callback", headers={"origin": "http://testserver"} + ) as ws: + ws.send_text(json.dumps(_callback_request("r1"))) + msgs = _collect_stream_messages(ws) + + assert [m["requestId"] for m in msgs] == ["r1"] * 3 + assert msgs[0]["payload"]["stream"] is True + assert msgs[0]["payload"]["data"]["response"] == {"out": {"children": "start"}} + assert msgs[1]["payload"]["data"]["response"] == {"out": {"children": "final"}} + assert msgs[2]["payload"] == {"status": "ok", "stream": True, "done": True} + + +def test_wsst002_sync_stream_generator_forbidden(): + """Sync generator streaming callbacks are rejected at registration.""" + from dash.exceptions import StreamCallbackError + + app, _ = _make_ws_app() + + with pytest.raises(StreamCallbackError, match="synchronous generator"): + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + def stream_cb(n): + yield "s1" + yield "s2" + + +def test_wsst003_stream_error_over_ws(): + pytest.importorskip("httpx", reason="fastapi.testclient requires httpx") + from fastapi.testclient import TestClient + + app, server = _make_ws_app() + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + async def stream_cb(n): + yield "one" + raise ValueError("boom") + + app._setup_server() + + client = TestClient(server) + with client.websocket_connect( + "/_dash-ws-callback", headers={"origin": "http://testserver"} + ) as ws: + ws.send_text(json.dumps(_callback_request("r1"))) + msgs = _collect_stream_messages(ws) + + assert msgs[0]["payload"]["data"]["response"] == {"out": {"children": "one"}} + assert msgs[1]["payload"]["status"] == "error" + assert "boom" in msgs[1]["payload"]["message"] + + +def test_wsst004_browser_stream_over_websocket(dash_duo): + """Full round-trip: streamed frames render progressively over WS.""" + app = Dash(__name__, backend="fastapi", websocket_callbacks=True) + app.layout = html.Div( + [ + html.Button("Start", id="btn", n_clicks=0), + html.Div(id="out", children="idle"), + ] + ) + + @app.callback( + Output("out", "children"), + Input("btn", "n_clicks"), + prevent_initial_call=True, + ) + async def stream_cb(n): + yield "streaming" + for token in ["a", "b", "c"]: + await asyncio.sleep(0.2) + patch = Patch() + patch += token + yield patch + + dash_duo.start_server(app) + dash_duo.wait_for_text_to_equal("#out", "idle") + dash_duo.find_element("#btn").click() + # Intermediate frame renders before the stream finishes. + dash_duo.wait_for_text_to_equal("#out", "streaming") + # Patch frames appended exactly once each. + dash_duo.wait_for_text_to_equal("#out", "streamingabc") + assert dash_duo.get_logs() == []