diff --git a/.ai/ARCHITECTURE.md b/.ai/ARCHITECTURE.md index d90f670f9b..7f40ee47a9 100644 --- a/.ai/ARCHITECTURE.md +++ b/.ai/ARCHITECTURE.md @@ -1057,6 +1057,109 @@ 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. +- Works with sync generators (Flask/Quart/FastAPI) and async generators. + Sync generators warn at registration: they occupy a server worker (or WS + executor thread) for the whole stream — prefer `async def`. +- 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. + +### 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 +- `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 + ## Security ### XSS Protection diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 8153a00876..e8e27739bf 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 @@ -61,6 +62,14 @@ jobs: - 'dash/dash-renderer/src/**' - 'tests/websocket/**' - 'requirements/**' + streaming_paths: + - 'dash/_callback.py' + - 'dash/_streaming.py' + - 'dash/_callback_context.py' + - 'dash/backends/**' + - 'dash/dash-renderer/src/**' + - 'tests/streaming/**' + - 'requirements/**' lint-unit: name: Lint & Unit Tests (Python ${{ matrix.python-version }}) @@ -129,6 +138,9 @@ jobs: - name: Run unit tests run: npm run citest.unit + - name: Run shared-storage tests + run: pytest tests/shared_storage -v + build: name: Build Dash Package runs-on: ubuntu-latest @@ -607,6 +619,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 5c4cff63c9..a9e4d226aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ This project adheres to [Semantic Versioning](https://semver.org/). ## [Unreleased] ### Added +- [#3888](https://github.com/plotly/dash/pull/3888/) Shared storage: a backend-agnostic state manager available on every app via `dash.ctx.shared_storage` (and `app.shared_storage`). It offers a cross-process key/value store (`get`/`set`/`delete`) and ordered publish/subscribe (`publish`/`subscribe`) for sharing state between callbacks or across worker processes without an external service. The default `LocalSharedStorage` backend elects a single owner process per machine (binding an `AF_UNIX` socket on POSIX, TCP loopback on Windows) and serves the rest; a single-process deployment is its own owner and pays no socket overhead. Subscriptions are ordered and replayable — a consumer that reconnects resumes from its last-seen message from a bounded buffer, and a buffer overrun surfaces as an explicit gap rather than a silent loss. Values must be JSON-compatible (like `dcc.Store`). Started lazily on first use, so it costs nothing until touched; pass `shared_storage=None` to `Dash(...)` to disable, or a `BaseSharedStorage` subclass/instance to swap the backend. The wire codec is msgspec (never pickle). +- [#3888](https://github.com/plotly/dash/pull/3888/) Streaming callbacks: a callback defined as a generator (or async generator) function streams — its yields are pushed to the browser as they are produced, no keyword needed. Each yielded value has the same shape as a regular return value and replaces the outputs; yielding `dash.Patch` objects gives incremental updates (e.g. LLM token streaming). Streams ride the WebSocket callback transport when active, otherwise the HTTP response streams NDJSON frames. Works on Flask, Quart, and FastAPI; the callback must be an `async def` generator — synchronous generators are rejected at registration since they occupy a server worker for the whole stream. When the app has a shared-storage backend (the default), all of a page's HTTP streams are multiplexed over a single downlink connection instead of one per callback — so they no longer count against the browser's ~6-connections-per-host limit — and the connection resumes from its last sequence on a reconnect without dropping frames. HTTP streams emit a blank keepalive line every `stream_keepalive_interval` milliseconds (default 15000) that the callback spends between yields, so proxy idle timeouts (nginx's `proxy_read_timeout` defaults to 60s) don't close a stream while the callback is still working; set `stream_keepalive_interval=None` on the app to disable. - [#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. ### Removed diff --git a/dash/__init__.py b/dash/__init__.py index 6f16a068aa..cb4572b539 100644 --- a/dash/__init__.py +++ b/dash/__init__.py @@ -43,6 +43,14 @@ from ._patch import Patch # noqa: F401,E402 from ._jupyter import jupyter_dash # noqa: F401,E402 +from ._shared_storage import ( # noqa: F401,E402 + BaseSharedStorage, + LocalSharedStorage, + SharedStorageError, + SharedStorageGap, + Subscription, +) + from ._hooks import hooks # noqa: F401,E402 ctx = callback_context @@ -94,4 +102,9 @@ def _jupyter_nbextension_paths(): "ctx", "hooks", "stringify_id", + "BaseSharedStorage", + "LocalSharedStorage", + "SharedStorageError", + "SharedStorageGap", + "Subscription", ] diff --git a/dash/_callback.py b/dash/_callback.py index 61908942ef..b74a7879fd 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] = [] @@ -109,6 +114,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 @@ -256,6 +270,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(): @@ -347,6 +387,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, } @@ -774,9 +817,23 @@ def register_callback( mcp_enabled=_kwargs.get("mcp_enabled", None), mcp_expose_docstring=_kwargs.get("mcp_expose_docstring"), ) + # 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 @@ -935,7 +992,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/_callback_context.py b/dash/_callback_context.py index b4285105f5..430d04591a 100644 --- a/dash/_callback_context.py +++ b/dash/_callback_context.py @@ -210,6 +210,16 @@ def states_list(self): def response(self): return getattr(_get_context_value(), "dash_response") + @property + def shared_storage(self): + """The app's shared storage (state manager + pub/sub). + + A backend-agnostic key/value + publish/subscribe store shared across + worker processes -- for cross-callback or session state without an + external service. Requires ``Dash(shared_storage=...)`` (on by default). + """ + return get_app().shared_storage + @staticmethod @has_context def record_timing(name, duration, description=None): diff --git a/dash/_shared_storage/__init__.py b/dash/_shared_storage/__init__.py new file mode 100644 index 0000000000..1e23a67ece --- /dev/null +++ b/dash/_shared_storage/__init__.py @@ -0,0 +1,22 @@ +"""Backend-agnostic shared state + pub/sub (state manager). + +Public surface grows as backends land. Today: the interface and the in-memory +engine; ``LocalSharedStorage`` (owner-elected, cross-process) and the Dash +wiring come next. +""" + +from .base import ( + BaseSharedStorage, + SharedStorageError, + SharedStorageGap, + Subscription, +) +from .local import LocalSharedStorage + +__all__ = [ + "BaseSharedStorage", + "LocalSharedStorage", + "SharedStorageError", + "SharedStorageGap", + "Subscription", +] diff --git a/dash/_shared_storage/_codec.py b/dash/_shared_storage/_codec.py new file mode 100644 index 0000000000..120f46f21a --- /dev/null +++ b/dash/_shared_storage/_codec.py @@ -0,0 +1,37 @@ +"""Wire codec for the shared-storage socket transport. + +Prefers ``msgspec`` (msgpack: fast, compact) and falls back to stdlib ``json``. +Both are **data-only -- never pickle** -- so bytes read off the socket cannot +execute code. All workers in one deployment share a single install, hence a +single codec, so the two ends always agree on the format. + +Payloads must be JSON-compatible (dict / list / str / int / float / bool / +None) under either codec -- the same constraint as ``dcc.Store`` and callback +outputs. +""" + +from typing import Any + +try: + import msgspec + + _encoder = msgspec.msgpack.Encoder() + _decoder = msgspec.msgpack.Decoder() + + def encode(obj: Any) -> bytes: + return _encoder.encode(obj) + + def decode(data: bytes) -> Any: + return _decoder.decode(data) + + CODEC = "msgpack" +except ImportError: # pragma: no cover - exercised via the fallback install + import json + + def encode(obj: Any) -> bytes: + return json.dumps(obj).encode("utf-8") + + def decode(data: bytes) -> Any: + return json.loads(data.decode("utf-8")) + + CODEC = "json" diff --git a/dash/_shared_storage/_engine.py b/dash/_shared_storage/_engine.py new file mode 100644 index 0000000000..35880685d4 --- /dev/null +++ b/dash/_shared_storage/_engine.py @@ -0,0 +1,119 @@ +"""In-memory store engine shared by the owner process and the single-process +fast path. + +Holds the authoritative key/value map and, per topic, an ordered log with a +bounded replay buffer. Sequence numbers are monotonic per topic starting at 1 +(``0`` means "before the first message"), so a subscriber tracks a cursor and a +reconnecting one resumes from its last-seen sequence. A consumer that fell +farther behind than the buffer holds gets an explicit gap signal instead of a +silent hole. + +The engine is thread-safe and transport-agnostic; sockets live one layer up. +""" + +import threading +import time +from collections import deque +from typing import Any, Deque, Dict, List, NamedTuple, Tuple + +# Per-topic replay buffer size. Sized to cover a reconnect window comfortably; +# a producer that outruns a disconnected consumer past this raises a gap rather +# than dropping messages silently. +DEFAULT_BUFFER = 2048 + + +class PollResult(NamedTuple): + messages: List[Any] + last_seq: int + gap: bool + + +class _Topic: # pylint: disable=too-few-public-methods + __slots__ = ("seq", "buf", "cond") + + def __init__(self, maxlen: int): + self.seq = 0 + self.buf: Deque[Tuple[int, Any]] = deque(maxlen=maxlen) + self.cond = threading.Condition() + + +class StoreEngine: + def __init__(self, buffer_size: int = DEFAULT_BUFFER): + self._buffer_size = buffer_size + self._data: Dict[str, Any] = {} + self._data_lock = threading.Lock() + self._topics: Dict[str, _Topic] = {} + self._topics_lock = threading.Lock() + self._closed = False + + @property + def closed(self) -> bool: + return self._closed + + # --- key/value --------------------------------------------------------- + def get(self, key: str, default: Any = None) -> Any: + with self._data_lock: + return self._data.get(key, default) + + def set(self, key: str, value: Any) -> None: + with self._data_lock: + self._data[key] = value + + def delete(self, key: str) -> None: + with self._data_lock: + self._data.pop(key, None) + + # --- pub/sub ----------------------------------------------------------- + def _topic(self, name: str) -> _Topic: + with self._topics_lock: + topic = self._topics.get(name) + if topic is None: + topic = self._topics[name] = _Topic(self._buffer_size) + return topic + + def publish(self, topic: str, message: Any) -> int: + t = self._topic(topic) + with t.cond: + t.seq += 1 + t.buf.append((t.seq, message)) + t.cond.notify_all() + return t.seq + + def head_seq(self, topic: str) -> int: + """Current highest sequence -- where a fresh subscription starts.""" + t = self._topic(topic) + with t.cond: + return t.seq + + 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 + elapsed (caller re-polls) or the engine closed. ``gap`` is True when the + next expected message was already evicted from the buffer. + """ + t = self._topic(topic) + 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) + remaining = deadline - time.monotonic() + if remaining <= 0: + return PollResult([], after_seq, False) + t.cond.wait(remaining) + + def close(self) -> None: + self._closed = True + with self._topics_lock: + topics = list(self._topics.values()) + for t in topics: + with t.cond: + t.cond.notify_all() diff --git a/dash/_shared_storage/_transport.py b/dash/_shared_storage/_transport.py new file mode 100644 index 0000000000..47d7f5e774 --- /dev/null +++ b/dash/_shared_storage/_transport.py @@ -0,0 +1,142 @@ +"""Socket transport for the owner-elected shared store. + +A single owner process serves one :class:`StoreEngine` over a stream socket +(AF_UNIX where available, TCP loopback otherwise). Messages are length-prefixed: +a 4-byte big-endian length followed by a body encoded by ``_codec`` (msgspec +msgpack, or stdlib JSON as a fallback). The codec is **data-only, never pickle** +-- the socket is localhost-only, but deserializing pickle off a socket is a +remote-code-execution hazard, so the wire carries only data. The channel is +additionally gated by a per-owner random token; only same-host clients that read +the owner's advertisement file (written 0600) can attach. + +The consequence is that stored values and published messages must be +JSON-serializable -- the same constraint as ``dcc.Store`` and callback outputs. + +Requests are ``[op, *args]`` lists; responses are ``["ok", value]`` or +``["err", repr]``. Ops: get/set/delete/publish/head/poll -- a thin passthrough +to the engine. ``poll`` blocks server-side up to its timeout (one thread per +connection), which is what makes long-poll subscriptions cheap. +""" + +import socket +import struct +import threading +from typing import Any, Optional + +from ._codec import decode, encode + +EOF = object() # returned by recv_frame on a clean close +OK = ["ok", None] + + +def send_frame(sock: socket.socket, obj: Any) -> None: + data = encode(obj) + sock.sendall(struct.pack("!I", len(data)) + data) + + +def _recv_exactly(sock: socket.socket, n: int) -> Optional[bytes]: + chunks = [] + remaining = n + while remaining: + chunk = sock.recv(remaining) + if not chunk: + return None + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def recv_frame(sock: socket.socket) -> Any: + header = _recv_exactly(sock, 4) + if header is None: + return EOF + (length,) = struct.unpack("!I", header) + body = _recv_exactly(sock, length) + if body is None: + return EOF + return decode(body) + + +class OwnerServer: + """Serves one StoreEngine to client workers over ``listen_sock``.""" + + def __init__(self, engine, listen_sock: socket.socket, token: str): + self._engine = engine + self._sock = listen_sock + self._token = token + self._closed = threading.Event() + self._accept_thread = threading.Thread( + target=self._accept_loop, name="dash-shared-owner", daemon=True + ) + + def start(self) -> None: + self._sock.listen(128) + self._accept_thread.start() + + def _accept_loop(self) -> None: + while not self._closed.is_set(): + try: + conn, _ = self._sock.accept() + except OSError: + break + threading.Thread(target=self._handle, args=(conn,), daemon=True).start() + + def _handle(self, conn: socket.socket) -> None: + try: + if recv_frame(conn) != self._token: + return + send_frame(conn, OK) + while not self._closed.is_set(): + req = recv_frame(conn) + if req is EOF: + break + send_frame(conn, self._dispatch(req)) + except (OSError, EOFError, ValueError): + pass + finally: + try: + conn.close() + except OSError: + pass + + def _dispatch(self, req): # pylint: disable=too-many-return-statements + op = req[0] + try: + if op == "get": + return ("ok", self._engine.get(req[1], req[2])) + if op == "set": + self._engine.set(req[1], req[2]) + return ("ok", None) + if op == "delete": + self._engine.delete(req[1]) + return ("ok", None) + if op == "publish": + return ("ok", self._engine.publish(req[1], req[2])) + if op == "head": + return ("ok", self._engine.head_seq(req[1])) + if op == "poll": + return ("ok", self._engine.poll(req[1], req[2], req[3])) + return ("err", f"unknown op {op!r}") + except Exception as err: # pylint: disable=broad-exception-caught + return ("err", repr(err)) + + def close(self) -> None: + self._closed.set() + try: + self._sock.close() + except OSError: + pass + self._engine.close() + + +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) + sock.settimeout(timeout) + sock.connect(address) + send_frame(sock, token) + if recv_frame(sock) != OK: + sock.close() + raise ConnectionError("shared-storage owner rejected the handshake") + sock.settimeout(None) + return sock diff --git a/dash/_shared_storage/base.py b/dash/_shared_storage/base.py new file mode 100644 index 0000000000..b1feda860e --- /dev/null +++ b/dash/_shared_storage/base.py @@ -0,0 +1,130 @@ +"""Backend-agnostic shared state + pub/sub for Dash. + +A ``BaseSharedStorage`` gives every worker process a common key/value store and +an ordered publish/subscribe channel. Dash uses it internally (e.g. to route +streaming-callback frames between the worker that runs a callback and the +worker holding the client's streaming connection), and apps can use it directly +for cross-callback or session state without reaching for an external service. + +Semantics every backend must honor: + +- **Values are picklable.** They may cross a process boundary. +- **Pub/sub is ordered and replayable.** Each ``publish`` to a topic gets a + monotonically increasing sequence number; a subscriber receives every message + published after it subscribed, in order. A consumer that drops and resubscribes + replays what it missed from a bounded buffer, so no message is silently lost + within that window -- a buffer overrun surfaces as an explicit gap error rather + than a missing frame. +""" + +import abc +from typing import Any, AsyncIterator, Iterator, Optional + + +class SharedStorageError(Exception): + """An operation failed on the shared-storage backend.""" + + +class SharedStorageGap(SharedStorageError): + """Raised by a subscription when messages were evicted before delivery. + + Signals that the replay buffer overran (the consumer fell too far behind), + or that the owner was re-elected and its buffer was lost; the caller must + treat it as lost data rather than a clean end of stream. + """ + + +class Subscription(abc.ABC): + """A live, ordered view of a topic. + + Iterate it synchronously (``for msg in sub``) or asynchronously + (``async for msg in sub``) to receive messages as they are published; + 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. + """ + + @abc.abstractmethod + def iter_with_seq(self) -> Iterator[Any]: + ... + + @abc.abstractmethod + 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 + + def __exit__(self, *exc) -> None: + self.close() + + async def __aenter__(self) -> "Subscription": + return self + + async def __aexit__(self, *exc) -> None: + self.close() + + +class BaseSharedStorage(abc.ABC): + """Shared key/value store + ordered pub/sub, usable across worker processes. + + A backend must be safe to construct in every worker. However the authoritative + state is held (a single elected owner, an external service, ...), all workers + that share a backend see the same keys and topics. + """ + + def start(self) -> None: + """Prepare the backend for use (elect/attach to the owner, connect, ...). + + Called once per worker before first use. Idempotent. + """ + + def close(self) -> None: + """Release this worker's handle on the backend. Idempotent.""" + + @abc.abstractmethod + def get(self, key: str, default: Any = None) -> Any: + ... + + @abc.abstractmethod + def set(self, key: str, value: Any) -> None: + ... + + @abc.abstractmethod + def delete(self, key: str) -> None: + ... + + @abc.abstractmethod + def publish(self, topic: str, message: Any) -> None: + """Append ``message`` to ``topic``; delivered to every current subscriber.""" + + @abc.abstractmethod + def subscribe(self, topic: str, replay_from: Optional[int] = None) -> Subscription: + """Subscribe to ``topic``. + + By default the subscription starts at the topic's current head, so it + only receives messages published from now on. ``replay_from`` (a sequence + number a previous subscription last saw) resumes after that point, + replaying buffered messages -- this is how a reconnecting consumer avoids + losing frames. + """ diff --git a/dash/_shared_storage/local.py b/dash/_shared_storage/local.py new file mode 100644 index 0000000000..da0b2defbb --- /dev/null +++ b/dash/_shared_storage/local.py @@ -0,0 +1,410 @@ +"""Owner-elected, cross-process shared storage. + +Every worker constructs a :class:`LocalSharedStorage`; on first use they race to +become the single *owner* by binding a stable address (AF_UNIX socket on POSIX, +TCP loopback on Windows). The winner hosts the authoritative :class:`StoreEngine` +and serves it; the losers become clients that connect to it. The bind itself is +the lease -- when the owner dies the address frees, and a client that finds it +gone re-elects (coming up cold, by design for the in-memory backend). + +Because the owner holds the engine directly, a single-process deployment (its +own owner) pays no socket overhead at all; only extra worker processes proxy. +""" + +import asyncio +import atexit +import hashlib +import json +import os +import secrets +import socket +import sys +import tempfile +import threading +import time +from typing import Any, Optional + +from ._engine import DEFAULT_BUFFER, PollResult, StoreEngine +from ._transport import EOF, OwnerServer, connect_to_owner, recv_frame, send_frame +from .base import BaseSharedStorage, SharedStorageError, SharedStorageGap, Subscription + +_HAS_AF_UNIX = hasattr(socket, "AF_UNIX") +_CLIENT_POLL_TIMEOUT = 20.0 # long-poll cycle for remote subscribers +_OWNER_POLL_TIMEOUT = 1.0 # local poll cycle; short so close() stays responsive + + +def _default_namespace() -> str: + raw = f"{os.getcwd()}|{sys.argv[0] if sys.argv else ''}".encode() + return hashlib.sha1(raw).hexdigest()[:16] + + +def _paths(namespace: str): + base = os.path.join(tempfile.gettempdir(), f"dash-shared-{namespace}") + return base + ".sock", base + ".addr" + + +def _tcp_port(namespace: str) -> int: + h = int(hashlib.sha1(namespace.encode()).hexdigest(), 16) + return 49152 + (h % (65535 - 49152)) + + +def _write_advertisement(addr_path, family, address, token) -> None: + payload = json.dumps( + {"family": int(family), "address": address, "token": token} + ).encode("utf-8") + tmp = f"{addr_path}.{os.getpid()}.tmp" + fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "wb") as f: + f.write(payload) + os.replace(tmp, addr_path) + + +def _read_advertisement(addr_path, retries=100, delay=0.05): + for _ in range(retries): + try: + with open(addr_path, "rb") as f: + data = json.loads(f.read().decode("utf-8")) + family = data["family"] + address = data["address"] + if family == int(socket.AF_INET): + address = tuple(address) # json list -> (host, port) + return family, address, data["token"] + except (FileNotFoundError, ValueError, KeyError): + time.sleep(delay) + raise ConnectionError("shared-storage owner advertisement not found") + + +def _unix_is_stale(path: str) -> bool: + probe = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + probe.settimeout(0.5) + try: + probe.connect(path) + return False # a live owner answered + except OSError: + return True # refused / no listener -> stale socket file + finally: + probe.close() + + +def _try_bind(namespace: str, sock_path: str): + """Try to become the owner. Returns (listen_sock, family, address) or None.""" + if _HAS_AF_UNIX: + s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + s.bind(sock_path) + except OSError: + if _unix_is_stale(sock_path): + try: + os.unlink(sock_path) + s.bind(sock_path) + except OSError: + s.close() + return None + else: + s.close() + return None + os.chmod(sock_path, 0o600) + return s, socket.AF_UNIX, sock_path + + port = _tcp_port(namespace) + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + s.bind(("127.0.0.1", port)) + except OSError: + s.close() + return None + return s, socket.AF_INET, ("127.0.0.1", port) + + +class _Coordinator: + """Owns this worker's role (owner or client) and re-elects on owner loss.""" + + def __init__(self, namespace: str, buffer_size: int): + self._namespace = namespace + self._buffer_size = buffer_size + self._sock_path, self._addr_path = _paths(namespace) + self._lock = threading.RLock() + self._role: Optional[str] = None + self._engine: Optional[StoreEngine] = None + self._server: Optional[OwnerServer] = None + self._family = None + self._address = None + self._token: Optional[str] = None + + @property + def engine(self) -> Optional[StoreEngine]: + return self._engine + + @property + def token(self) -> Optional[str]: + return self._token + + def ensure(self) -> None: + if self._role is None: + with self._lock: + if self._role is None: + self._elect() + + def is_owner(self) -> bool: + self.ensure() + return self._role == "owner" + + def _elect(self) -> None: + bound = _try_bind(self._namespace, self._sock_path) + if bound: + listen_sock, family, address = bound + token = secrets.token_hex(16) + _write_advertisement(self._addr_path, family, address, token) + engine = StoreEngine(self._buffer_size) + server = OwnerServer(engine, listen_sock, token) + server.start() + self._engine, self._server = engine, server + self._family, self._address, self._token = family, address, token + self._role = "owner" + # Best-effort cleanup of the socket/advertisement on clean exit; a + # crash leaves them stale, which the next election self-heals. + atexit.register(self.close) + else: + self._family, self._address, self._token = _read_advertisement( + self._addr_path + ) + self._role = "client" + + def connect(self) -> socket.socket: + self.ensure() + return connect_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: + if self._role == "owner": + return + self._role = None + self._elect() + + def close(self) -> None: + with self._lock: + if self._server is not None: + self._server.close() + for path in (self._sock_path, self._addr_path): + try: + os.unlink(path) + except OSError: + pass + self._role = None + + +class _LocalSubscription(Subscription): + """A topic view that reads from the owner engine directly (owner role) or + long-polls it over a dedicated connection (client role), resuming from its + cursor on reconnect so no buffered message is missed. + """ + + def __init__(self, coord: _Coordinator, topic: str, start_seq: int): + self._coord = coord + self._topic = topic + self._cursor = start_seq + self._conn: Optional[socket.socket] = None + self._closed = threading.Event() + + def close(self) -> None: + self._closed.set() + conn, self._conn = self._conn, None + if conn is not None: + try: + conn.close() + except OSError: + pass + + def _poll_once(self) -> PollResult: + if self._coord.is_owner(): + engine = self._coord.engine + if engine is None or engine.closed: + self._closed.set() # owner gone -> end iteration, don't busy-loop + return PollResult([], self._cursor, False) + return engine.poll(self._topic, self._cursor, _OWNER_POLL_TIMEOUT) + return self._client_poll() + + def _client_poll(self) -> PollResult: + for attempt in range(4): + if self._closed.is_set(): + return PollResult([], self._cursor, False) + if self._conn is None: + prev_token = self._coord.token + try: + self._conn = self._coord.connect() + except OSError as exc: + # Owner unreachable: brief retries to the same owner, then + # re-elect. A changed owner means its buffer is gone -> gap. + 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 + time.sleep(0.1 * (attempt + 1)) + continue + try: + send_frame( + self._conn, + ["poll", self._topic, self._cursor, _CLIENT_POLL_TIMEOUT], + ) + resp = recv_frame(self._conn) + 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): + if self._conn is not None: + try: + self._conn.close() + except OSError: + pass + self._conn = None # reconnect on the next attempt + return PollResult([], self._cursor, False) + + @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 iter_with_seq(self): + try: + while not self._closed.is_set(): + res = self._poll_once() + if res.gap: + raise SharedStorageGap( + f"replay buffer overran on topic {self._topic!r}" + ) + yield from self._with_seq(res) + self._cursor = res.last_seq + finally: + self.close() + + def aiter_with_seq(self): + return self._aiter_with_seq() + + async def _aiter_with_seq(self): + loop = asyncio.get_running_loop() + 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 + if res.gap: + raise SharedStorageGap( + f"replay buffer overran on topic {self._topic!r}" + ) + for pair in self._with_seq(res): + yield pair + self._cursor = res.last_seq + finally: + self.close() + + +class LocalSharedStorage(BaseSharedStorage): + """In-memory shared storage, elected to a single owner process per machine. + + Values and published messages must be JSON-compatible. State is not durable: + if the owning process dies, a survivor re-elects with an empty store. + """ + + def __init__( + self, namespace: Optional[str] = None, buffer_size: int = DEFAULT_BUFFER + ): + self._coord = _Coordinator(namespace or _default_namespace(), buffer_size) + self._conn: Optional[socket.socket] = None + self._conn_lock = threading.Lock() + + def start(self) -> None: + self._coord.ensure() + + def close(self) -> None: + with self._conn_lock: + if self._conn is not None: + try: + self._conn.close() + except OSError: + pass + self._conn = None + self._coord.close() + + # --- key/value + publish (short request/response) ---------------------- + def _call(self, req): + last_err: Optional[Exception] = None + for _ in range(3): + if self._coord.is_owner(): + return self._local(req) + try: + return self._remote(req) + except (OSError, EOFError, ConnectionError) as err: + last_err = err + self._coord.on_owner_lost() + raise SharedStorageError(f"shared-storage owner unreachable: {last_err}") + + def _local(self, req): + engine = self._coord.engine + assert engine is not None + op = req[0] + if op == "get": + return engine.get(req[1], req[2]) + if op == "set": + return engine.set(req[1], req[2]) + if op == "delete": + return engine.delete(req[1]) + if op == "publish": + return engine.publish(req[1], req[2]) + if op == "head": + return engine.head_seq(req[1]) + raise ValueError(f"unknown op {op!r}") + + def _remote(self, req): + with self._conn_lock: + if self._conn is None: + self._conn = self._coord.connect() + try: + send_frame(self._conn, req) + resp = recv_frame(self._conn) + if resp is EOF: + raise ConnectionError("owner closed the connection") + except (OSError, EOFError, ConnectionError): + try: + self._conn.close() + except OSError: + pass + self._conn = None + raise + status, val = resp + if status == "err": + raise SharedStorageError(val) + return val + + def get(self, key: str, default: Any = None) -> Any: + return self._call(["get", key, default]) + + def set(self, key: str, value: Any) -> None: + self._call(["set", key, value]) + + def delete(self, key: str) -> None: + self._call(["delete", key]) + + def publish(self, topic: str, message: Any) -> None: + self._call(["publish", topic, message]) + + def _head(self, topic: str) -> int: + return self._call(["head", topic]) + + 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) diff --git a/dash/_stream_hub.py b/dash/_stream_hub.py new file mode 100644 index 0000000000..242af780f1 --- /dev/null +++ b/dash/_stream_hub.py @@ -0,0 +1,151 @@ +"""Multiplexed streaming over shared storage. + +A browser holds a single downlink NDJSON connection 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 subscribes to 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. + +Downlink line shape (one JSON object per NDJSON line):: + + {"rid": "", "frame": {}} + +where ``frame`` is a ``CallbackExecutionResponse`` frame or a ``{"done": true}`` +terminal, exactly as the single-callback NDJSON transport emits today. +""" + +import asyncio +import json +from typing import Any, AsyncIterator, Iterator, Optional + +from ._shared_storage.base import BaseSharedStorage +from ._streaming import StreamedCallbackResponse, sync_iter_asyncgen, to_json + +_TOPIC_PREFIX = "_dash_stream:" + +# 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} + + +def stream_topic(connection_id: str) -> str: + return f"{_TOPIC_PREFIX}{connection_id}" + + +def publish_frame( + storage: BaseSharedStorage, + connection_id: str, + request_id: str, + frame: Any, +) -> None: + """Publish one streaming frame onto a connection's downlink topic. + + 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. + """ + plain = json.loads(to_json(frame)) + storage.publish(stream_topic(connection_id), {"rid": request_id, "frame": plain}) + + +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. + + These frames feed a ``StreamedCallbackResponse`` so the existing NDJSON + response path serializes and keep-alives them -- no bespoke endpoint. It is + long-lived: it carries frames for every callback on the connection, not one + stream, and ends when the client hangs up. ``replay_from`` resumes a + reconnecting downlink from its last cursor. Each envelope carries its ``seq`` + so the client can resume from it after a reconnect without losing frames. + """ + with storage.subscribe(stream_topic(connection_id), replay_from) as sub: + for seq, message in sub.iter_with_seq(): + yield {**message, "seq": seq} + + +async def asubscribe_envelopes( + storage: BaseSharedStorage, + connection_id: str, + replay_from: Optional[int] = None, +) -> AsyncIterator[Any]: + """Async counterpart of :func:`subscribe_envelopes` for ASGI backends.""" + sub = storage.subscribe(stream_topic(connection_id), replay_from) + try: + async for seq, message in sub.aiter_with_seq(): + yield {**message, "seq": seq} + finally: + sub.close() + + +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 background task on the uplink worker 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. + """ + async for frame in marker.frames: + publish_frame(storage, connection_id, request_id, frame) + + +def pump_to_storage( + storage: BaseSharedStorage, + connection_id: str, + request_id: str, + marker: StreamedCallbackResponse, +) -> None: + """Sync driver for WSGI workers: drive the async frame generator on a + private event loop (via ``sync_iter_asyncgen``) and publish each frame. + Runs on a background thread so the callback's POST returns immediately. + """ + for frame in sync_iter_asyncgen(marker.frames): + publish_frame(storage, connection_id, request_id, frame) + + +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.""" + return StreamedCallbackResponse( + asubscribe_envelopes(storage, connection_id, replay_from), is_async=True + ) + + +# Keep a reference to in-flight pump tasks so the loop doesn't GC them mid-stream. +_pending_pumps: "set[asyncio.Task]" = set() + + +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) diff --git a/dash/_streaming.py b/dash/_streaming.py new file mode 100644 index 0000000000..bf9ec39a49 --- /dev/null +++ b/dash/_streaming.py @@ -0,0 +1,311 @@ +"""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 +from typing import cast + +from ._utils import to_json as _to_json + +logger = logging.getLogger(__name__) + +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.""" + while True: + try: + yield marker.ctx.run(next, marker.frames) + except StopIteration: + return + + +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() + try: + while True: + try: + kind, value = frames.get(timeout=keepalive) + except queue.Empty: + yield _KEEPALIVE + continue + 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. + """ + iterator = frames.__aiter__() + pending = None + try: + while True: + pending = asyncio.ensure_future(iterator.__anext__()) + while True: + done, _ = await asyncio.wait({pending}, timeout=keepalive) + if done: + break + yield _KEEPALIVE + try: + frame = pending.result() + except StopAsyncIteration: + return + finally: + pending = None + yield frame + finally: + # Reached on client disconnect too: cancel the in-flight step so the + # user generator sees CancelledError at its current await. + 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) + for frame in frames: + line, fatal = _line(frame) + yield line + if fatal: + return + + +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. + """ + 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) diff --git a/dash/backends/_fastapi.py b/dash/backends/_fastapi.py index 6a23c4873b..dc868e626e 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,15 @@ from dash.fingerprint import check_fingerprint from dash import _validate, get_app +from dash._streaming import ( + STREAM_HEADERS, + STREAM_MIMETYPE, + StreamedCallbackResponse, + keepalive_seconds, + marker_ndjson_aiter, + to_json, +) +from dash._stream_hub import STREAM_ACK, async_downlink_marker, spawn_async_pump from dash.exceptions import PreventUpdate from .base_server import ( BaseDashServer, @@ -48,6 +57,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 @@ -547,9 +557,30 @@ def add_redirect_rule(self, app, fullname, path): ) def serve_callback(self, dash_app: Dash): + 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 # pylint: disable=protected-access body = self.request_adapter().get_json() + downlink = body.get("streamDownlink") + if downlink is not None: + # The client's single multiplexed streaming connection. + return _ndjson_response( + async_downlink_marker( + dash_app.shared_storage, + downlink["connectionId"], + downlink.get("from"), + ) + ) cb_ctx = dash_app._initialize_context( body ) # pylint: disable=protected-access @@ -566,6 +597,19 @@ 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 and dash_app.shared_storage_enabled: + # Multiplexed uplink: pump frames onto the connection's topic + # and return immediately instead of holding this connection. + spawn_async_pump( + dash_app.shared_storage, + stream_conn["connectionId"], + 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 @@ -812,6 +856,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( @@ -820,6 +870,7 @@ async def websocket_handler(websocket: WebSocket): payload, ws_cb, FastAPIResponseAdapter(), + stream_emitter, ) ) task.add_done_callback(done_handler) @@ -832,6 +883,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 43f24634a2..8694d775ae 100644 --- a/dash/backends/_flask.py +++ b/dash/backends/_flask.py @@ -4,6 +4,7 @@ import pkgutil import sys import mimetypes +import threading import time import inspect import traceback @@ -22,6 +23,7 @@ g as flask_g, has_request_context, redirect, + stream_with_context, ) from werkzeug.debug import tbtools @@ -29,6 +31,17 @@ from dash import _validate from dash.exceptions import PreventUpdate, InvalidResourceError from dash._callback import _invoke_callback, _async_invoke_callback +from dash._streaming import ( + STREAM_HEADERS, + STREAM_MIMETYPE, + StreamedCallbackResponse, + andjson_lines, + keepalive_seconds, + ndjson_lines, + sync_iter_asyncgen, + to_json, +) +from dash._stream_hub import pump_to_storage, subscribe_envelopes from dash._utils import parse_version from .base_server import BaseDashServer, RequestAdapter, ResponseAdapter @@ -251,8 +264,74 @@ def add_redirect_rule(self, app, fullname, path): # pylint: disable=unused-argument def serve_callback(self, dash_app: Dash): + 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 = 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, with_request_ctx): + # The client's single multiplexed streaming connection: relay this + # connection's frames (published by streaming callbacks, possibly on + # other workers, via shared storage) as an ordinary NDJSON stream. + storage = dash_app.shared_storage + frames = subscribe_envelopes( + storage, downlink["connectionId"], downlink.get("from") + ) + marker = StreamedCallbackResponse( + frames, is_async=False, ctx=copy_context() + ) + return _stream_response(marker, with_request_ctx=with_request_ctx) + + def _serve_uplink(marker, body, cb_ctx): + # Multiplexed uplink: if the streaming callback carries a connection, + # pump its frames onto that connection's topic on a background thread + # and return immediately, so this request does not hold a connection + # for the stream's life. Returns None to fall back to inline NDJSON. + stream_conn = body.get("streamConnection") + if not (stream_conn and dash_app.shared_storage_enabled): + return None + threading.Thread( + target=pump_to_storage, + args=( + dash_app.shared_storage, + stream_conn["connectionId"], + stream_conn["requestId"], + marker, + ), + daemon=True, + name="dash-stream-pump", + ).start() + return cb_ctx.dash_response.set_response( + data=to_json({"multi": True, "stream": True}) + ) + def _dispatch(): body = request.get_json() + downlink = body.get("streamDownlink") + if downlink is not None: + return _serve_downlink(downlink, with_request_ctx=True) # pylint: disable=protected-access cb_ctx = dash_app._initialize_context(body) func = dash_app._prepare_callback(cb_ctx, body) @@ -262,6 +341,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]. " @@ -272,6 +356,11 @@ def _dispatch(): async def _dispatch_async(): body = request.get_json() + downlink = body.get("streamDownlink") + if downlink is not None: + # Async view: request context lives in a different contextvars + # context, so stream_with_context must not wrap the body. + return _serve_downlink(downlink, with_request_ctx=False) # pylint: disable=protected-access cb_ctx = dash_app._initialize_context(body) func = dash_app._prepare_callback(cb_ctx, body) @@ -283,6 +372,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 9487cac7d8..fce24f65de 100644 --- a/dash/backends/_quart.py +++ b/dash/backends/_quart.py @@ -40,6 +40,15 @@ from dash.exceptions import PreventUpdate, InvalidResourceError from dash.fingerprint import check_fingerprint +from dash._streaming import ( + STREAM_HEADERS, + STREAM_MIMETYPE, + StreamedCallbackResponse, + keepalive_seconds, + marker_ndjson_aiter, + to_json, +) +from dash._stream_hub import STREAM_ACK, async_downlink_marker, spawn_async_pump from dash._utils import parse_version from dash import _validate from .base_server import ( @@ -53,6 +62,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 @@ -383,9 +393,29 @@ 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 + 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(): adapter = QuartRequestAdapter() body = await adapter.get_json() + downlink = body.get("streamDownlink") + if downlink is not None: + # The client's single multiplexed streaming connection. + marker = async_downlink_marker( + dash_app.shared_storage, + downlink["connectionId"], + downlink.get("from"), + ) + return _ndjson_response(marker) # pylint: disable=protected-access cb_ctx = dash_app._initialize_context(body) # pylint: disable=protected-access @@ -400,6 +430,19 @@ 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 and dash_app.shared_storage_enabled: + # Multiplexed uplink: pump frames onto the connection's topic + # and return immediately instead of holding this connection. + spawn_async_pump( + dash_app.shared_storage, + stream_conn["connectionId"], + 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 @@ -652,6 +695,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( @@ -660,6 +709,7 @@ async def websocket_handler(): # pylint: disable=too-many-branches payload, ws_cb, QuartResponseAdapter(), + stream_emitter, ) ) task.add_done_callback(done_handler) @@ -672,6 +722,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/ws.py b/dash/backends/ws.py index 6dc2c0e926..2f74851e34 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: @@ -506,6 +512,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, @@ -533,6 +638,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. @@ -571,6 +677,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: @@ -589,6 +699,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. @@ -616,6 +727,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/src/actions/callbacks.ts b/dash/dash-renderer/src/actions/callbacks.ts index 6dec01996c..b63dcbff0e 100644 --- a/dash/dash-renderer/src/actions/callbacks.ts +++ b/dash/dash-renderer/src/actions/callbacks.ts @@ -47,6 +47,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'; @@ -460,6 +461,73 @@ 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 (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().run(url, init, 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, @@ -614,7 +682,92 @@ 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(); + 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; @@ -726,7 +879,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) { @@ -760,6 +922,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; @@ -1044,6 +1234,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; @@ -1057,6 +1256,14 @@ export function executeCallback( payload, cb.callback.running ); + } else if (useStream) { + // Multiplexed HTTP streaming (single downlink) + data = await handleStreamCallback( + dispatch, + newConfig, + payload, + cb.callback.running + ); } 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..d1403e80c4 100644 --- a/dash/dash-renderer/src/config.ts +++ b/dash/dash-renderer/src/config.ts @@ -29,6 +29,9 @@ export type DashConfig = { inactivity_timeout?: number; heartbeat_interval?: number; }; + stream?: { + enabled: boolean; + }; 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 5f963463d2..51da885df8 100644 --- a/dash/dash-renderer/src/types/callbacks.ts +++ b/dash/dash-renderer/src/types/callbacks.ts @@ -17,6 +17,11 @@ export interface ICallbackDefinition { no_output?: boolean; websocket?: boolean; persistent?: boolean; + // Server-inferred: true when the callback is a (streaming) generator. Read + // only by the scheduler, to keep long-lived streams out of the concurrent + // HTTP-request budget. The client still detects streaming at runtime from + // the response (NDJSON content type / stream+done markers), not from this. + stream?: boolean; } export interface ICallbackProperty { @@ -109,6 +114,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..86d1dafc10 --- /dev/null +++ b/dash/dash-renderer/src/utils/streamClient.ts @@ -0,0 +1,219 @@ +/** + * Single multiplexed streaming transport for the page. + * + * 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 connection id + request id; the server pumps that callback's frames + * onto the connection's shared-storage topic; the single downlink relays them and + * this client routes each frame back to the right callback by request id. + * + * Lifecycle: the downlink opens on the first streaming callback and closes once no + * callbacks remain in flight ("collect the dones to match the runnings"). If it + * drops while callbacks are still running it reconnects, resuming from the last + * sequence it saw so buffered frames are replayed rather than lost. + * + * This currently runs on the page; it is written to be host-agnostic so it can + * move into a SharedWorker (one connection per browser, shared across tabs) later. + */ + +type Frame = Record; + +interface PendingStream { + onFrame: (frame: Frame) => void; + resolve: () => void; + reject: (err: Error) => void; +} + +interface DownlinkEnvelope { + rid: string; + frame: Frame; + seq?: number; +} + +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 { + private connectionId = genId(); + 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; + private reconnectDelay: number; + private fetchImpl: FetchImpl; + + constructor(opts: {fetchImpl?: FetchImpl; reconnectDelay?: number} = {}) { + // 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; + } + + get activeCount(): number { + return this.pending.size; + } + + /** + * Run one streaming callback over the multiplexed transport. 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, + payload: Record, + onFrame: (frame: Frame) => void + ): Promise { + const requestId = `${this.connectionId}-${++this.counter}`; + const settled = new Promise((resolve, reject) => { + this.pending.set(requestId, {onFrame, resolve, reject}); + }); + this.ensureDownlink(url, init); + // Uplink POST: returns a fast ack; the outputs arrive on the downlink. + this.fetchImpl(url, { + ...init, + method: 'POST', + body: JSON.stringify({ + ...payload, + streamConnection: {connectionId: this.connectionId, requestId} + }) + }).catch(err => this.fail(requestId, err)); + return settled; + } + + /** Route one downlink envelope to its callback. Public for testing. */ + dispatchEnvelope(envelope: DownlinkEnvelope): void { + if (typeof envelope.seq === 'number') { + this.cursor = envelope.seq; + } + 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.onFrame(frame); + } + } + + 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.pending.size === 0 && this.abort) { + this.abort.abort(); // ends the read loop; downlink closes + } + } + + 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).finally(() => { + this.downlinkOpen = false; + this.abort = null; + }); + } + + private async readLoop(url: string, init: RequestInit): Promise { + while (this.pending.size > 0) { + this.abort = new AbortController(); + try { + const res = await this.fetchImpl(url, { + ...init, + method: 'POST', + signal: this.abort.signal, + body: JSON.stringify({ + streamDownlink: { + connectionId: this.connectionId, + from: this.cursor + } + }) + }); + if (!res.ok || !res.body) { + throw new Error(`downlink responded ${res.status}`); + } + await this.consume(res.body); + } catch (err) { + if (this.pending.size === 0) { + break; // deliberately aborted because we went idle + } + // Genuine drop with work outstanding: reconnect from the cursor. + await sleep(this.reconnectDelay); + } + } + } + + private async consume(body: ReadableStream): Promise { + const reader = body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + for (;;) { + const {done, value} = await reader.read(); + if (done) { + return; // connection ended -> reconnect via the read loop + } + 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 + } + this.dispatchEnvelope(JSON.parse(line)); + } + } + } +} + +let singleton: StreamClient | null = null; + +export function getStreamClient(): StreamClient { + if (!singleton) { + singleton = new StreamClient(); + } + 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/workerClient.ts b/dash/dash-renderer/src/utils/workerClient.ts index 1c07a4c1c3..b09433d85c 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 */ @@ -43,6 +47,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; } /** @@ -203,9 +209,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; @@ -218,7 +230,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, @@ -301,8 +313,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/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..0ce6e214d3 --- /dev/null +++ b/dash/dash-renderer/tests/streamClient.test.js @@ -0,0 +1,177 @@ +import {expect} from 'chai'; +import {beforeEach, describe, it} from 'mocha'; + +import {StreamClient} from '../src/utils/streamClient'; + +// 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 = []; + 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})); + } + uplinks.push(body); + return Promise.resolve( + new Response(JSON.stringify({multi: true, stream: true}), { + status: 200 + }) + ); + }; + return {fetchImpl, uplinks, downlinks}; +} + +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('sends an uplink tagging the callback with a connection + request id', async () => { + client.run('/cb', {}, {output: 'a.b'}, () => {}); + await waitFor(() => mock.uplinks.length === 1); + const conn = mock.uplinks[0].streamConnection; + expect(conn.connectionId).to.be.a('string'); + expect(conn.requestId).to.be.a('string'); + expect(mock.uplinks[0].output).to.equal('a.b'); // original payload preserved + }); + + it('routes frames to onFrame and resolves on the done frame', async () => { + const frames = []; + const settled = client.run('/cb', {}, {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', {}, {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', {}, {output: 'a'}, f => aFrames.push(f)); + const b = client.run('/cb', {}, {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', {}, {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('skips keepalive blank lines', async () => { + const frames = []; + const settled = client.run('/cb', {}, {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}}]); + }); +}); diff --git a/dash/dash.py b/dash/dash.py index fd6a8d017d..a5bc1d1278 100644 --- a/dash/dash.py +++ b/dash/dash.py @@ -16,7 +16,17 @@ import hashlib import base64 from urllib.parse import urlparse -from typing import Any, Callable, Dict, Optional, Union, Sequence, Literal, List +from typing import ( + Any, + Callable, + Dict, + Optional, + Type, + Union, + Sequence, + Literal, + List, +) import traceback @@ -65,6 +75,11 @@ from . import backends from ._get_app import with_app_context, with_app_context_factory +from ._shared_storage import ( + BaseSharedStorage, + LocalSharedStorage, + SharedStorageError, +) from ._grouping import map_grouping, grouping_len, update_args_group from ._obsolete import ObsoleteChecker from ._callback_context import callback_context @@ -433,6 +448,13 @@ 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 """ _plotlyjs_url: str @@ -493,6 +515,10 @@ 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, + shared_storage: Optional[ + Union[Type[BaseSharedStorage], BaseSharedStorage] + ] = LocalSharedStorage, enable_mcp: Optional[bool] = None, mcp_path: Optional[str] = None, **obsolete, @@ -668,6 +694,15 @@ 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 + + # Shared storage (state manager + pub/sub). Started lazily on first + # access so it costs nothing until used and never binds in a gunicorn + # preload master or the Flask reloader parent -- only in the worker that + # actually touches it. + self._shared_storage_arg = shared_storage + self._shared_storage_instance = None + self._shared_storage_lock = threading.Lock() self.logger = logging.getLogger(__name__) @@ -929,6 +964,37 @@ 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. + + Started on first access in the worker that uses it. Access it from a + callback via ``dash.ctx.shared_storage``. Raises if the app was created + with ``shared_storage=None``. + """ + if self._shared_storage_arg is None: + raise SharedStorageError( + "Shared storage is disabled for this app; construct it with " + "Dash(shared_storage=...) to use dash.ctx.shared_storage." + ) + if self._shared_storage_instance is None: + with self._shared_storage_lock: + if self._shared_storage_instance is None: + storage = self._shared_storage_arg + if isinstance(storage, type): + storage = storage() + storage.start() + self._shared_storage_instance = storage + return self._shared_storage_instance + def _layout_value(self): if self._layout_is_function: layout = self._layout() # type: ignore[reportOptionalCall] @@ -1121,6 +1187,11 @@ 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. + config["stream"] = {"enabled": self.shared_storage_enabled} + return config def serve_reload_hash(self): 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 8e18280d04..178082dc99 100644 --- a/requirements/ci.txt +++ b/requirements/ci.txt @@ -3,6 +3,7 @@ black==22.3.0 flake8==7.0.0 flaky==3.8.1 flask-talisman==1.0.0 +httpx # fastapi.testclient (tests/websocket/test_ws_stream.py) ipython<9.0.0 mimesis<=11.1.0 mock==4.0.3 diff --git a/requirements/install.txt b/requirements/install.txt index 176d9c8f5e..795f0c85c8 100644 --- a/requirements/install.txt +++ b/requirements/install.txt @@ -10,3 +10,4 @@ setuptools janus>=1.0.0 pydantic>=2.10 comm +msgspec diff --git a/tests/shared_storage/__init__.py b/tests/shared_storage/__init__.py new file mode 100644 index 0000000000..4c3c33707f --- /dev/null +++ b/tests/shared_storage/__init__.py @@ -0,0 +1 @@ +# Shared-storage (state manager + pub/sub) tests. diff --git a/tests/shared_storage/test_dash_integration.py b/tests/shared_storage/test_dash_integration.py new file mode 100644 index 0000000000..7b319d7375 --- /dev/null +++ b/tests/shared_storage/test_dash_integration.py @@ -0,0 +1,62 @@ +"""Shared storage wired into a Dash app and reachable from a callback.""" +import json +import uuid + +import pytest + +from dash import Dash, Input, Output, ctx, html +from dash._shared_storage import LocalSharedStorage, SharedStorageError + + +def _dispatch(app, output_id="out", prop="children", input_id="btn"): + body = { + "output": f"{output_id}.{prop}", + "outputs": {"id": output_id, "property": prop}, + "inputs": [{"id": input_id, "property": "n_clicks", "value": 3}], + "changedPropIds": [f"{input_id}.n_clicks"], + } + resp = app.server.test_client().post("/_dash-update-component", json=body) + assert resp.status_code == 200 + return json.loads(resp.get_data(as_text=True)) + + +def _app(): + # Unique namespace per app so tests don't share one owner/store. + storage = LocalSharedStorage(namespace=f"itest-{uuid.uuid4().hex[:12]}") + app = Dash(__name__, shared_storage=storage) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + return app, storage + + +def test_ctx_shared_storage_roundtrip_in_callback(): + app, storage = _app() + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + def cb(n): + ctx.shared_storage.set("hits", n) + return f"stored {ctx.shared_storage.get('hits')}" + + out = _dispatch(app) + assert out["response"]["out"]["children"] == "stored 3" + # Value is visible on the app's storage handle outside the callback too. + assert storage.get("hits") == 3 + storage.close() + + +def test_shared_storage_disabled_raises_in_callback(): + app = Dash(__name__, shared_storage=None) + app.layout = html.Div([html.Button(id="btn"), html.Div(id="out")]) + + @app.callback(Output("out", "children"), Input("btn", "n_clicks")) + def cb(n): + return ctx.shared_storage.get("x") + + # The callback error surfaces; storage access without a backend is rejected. + with pytest.raises(SharedStorageError): + app.shared_storage # noqa: B018 + + +def test_default_shared_storage_is_enabled(): + app = Dash(__name__) + assert isinstance(app.shared_storage, LocalSharedStorage) + app.shared_storage.close() diff --git a/tests/shared_storage/test_engine.py b/tests/shared_storage/test_engine.py new file mode 100644 index 0000000000..b1f849d4d0 --- /dev/null +++ b/tests/shared_storage/test_engine.py @@ -0,0 +1,119 @@ +"""Unit tests for the in-memory StoreEngine (KV + sequenced pub/sub).""" +import threading +import time + +from dash._shared_storage._engine import StoreEngine + + +def test_kv_get_set_delete(): + e = StoreEngine() + assert e.get("missing") is None + assert e.get("missing", 42) == 42 + e.set("a", {"x": 1}) + assert e.get("a") == {"x": 1} + e.delete("a") + assert e.get("a") is None + e.delete("a") # idempotent + + +def test_publish_assigns_monotonic_seq(): + e = StoreEngine() + assert e.head_seq("t") == 0 + assert e.publish("t", "a") == 1 + assert e.publish("t", "b") == 2 + assert e.head_seq("t") == 2 + + +def test_fresh_subscriber_only_sees_future_messages(): + e = StoreEngine() + e.publish("t", "old") + cursor = e.head_seq("t") # subscribe "now" + e.publish("t", "new1") + e.publish("t", "new2") + res = e.poll("t", cursor, timeout=1) + assert res.messages == ["new1", "new2"] + assert res.last_seq == 3 # "old" took seq 1, so new2 is seq 3 + assert res.gap is False + + +def test_replay_from_cursor_after_reconnect(): + e = StoreEngine() + e.publish("t", "m1") + e.publish("t", "m2") + e.publish("t", "m3") + # A consumer that saw up to seq 1 reconnects and replays 2 and 3. + res = e.poll("t", 1, timeout=1) + assert res.messages == ["m2", "m3"] + assert res.last_seq == 3 + assert res.gap is False + + +def test_gap_when_buffer_overruns(): + e = StoreEngine(buffer_size=2) + for i in range(5): + e.publish("t", f"m{i}") # seqs 1..5, buffer holds only seqs 4,5 + # A consumer stuck at seq 1 wanted seq 2, which was evicted -> gap. + res = e.poll("t", 1, timeout=1) + assert res.gap is True + assert res.messages == [] + + +def test_no_gap_at_buffer_edge(): + e = StoreEngine(buffer_size=2) + for i in range(4): + e.publish("t", f"m{i}") # seqs 1..4, buffer holds 3,4 + # Consumer at seq 2 wants seq 3, which is still buffered -> no gap. + res = e.poll("t", 2, timeout=1) + assert res.gap is False + assert res.messages == ["m2", "m3"] + + +def test_poll_times_out_empty_when_no_messages(): + e = StoreEngine() + start = time.monotonic() + res = e.poll("t", 0, timeout=0.2) + assert res.messages == [] + assert res.gap is False + assert time.monotonic() - start >= 0.2 + + +def test_poll_wakes_on_publish_from_another_thread(): + e = StoreEngine() + received = [] + + def consumer(): + res = e.poll("t", 0, timeout=2) + received.extend(res.messages) + + th = threading.Thread(target=consumer) + th.start() + time.sleep(0.1) # ensure the poll is waiting + e.publish("t", "live") + th.join(timeout=2) + assert received == ["live"] + + +def test_close_unblocks_waiting_pollers(): + e = StoreEngine() + done = threading.Event() + + def consumer(): + e.poll("t", 0, timeout=10) + done.set() + + th = threading.Thread(target=consumer) + th.start() + time.sleep(0.1) + e.close() + assert done.wait(timeout=2) + + +def test_multiple_subscribers_each_get_every_message(): + e = StoreEngine() + cursor = e.head_seq("t") + e.publish("t", "a") + e.publish("t", "b") + a = e.poll("t", cursor, timeout=1) + b = e.poll("t", cursor, timeout=1) + assert a.messages == ["a", "b"] + assert b.messages == ["a", "b"] # independent cursors, both see all diff --git a/tests/shared_storage/test_local_cross_process.py b/tests/shared_storage/test_local_cross_process.py new file mode 100644 index 0000000000..8e376ec8d9 --- /dev/null +++ b/tests/shared_storage/test_local_cross_process.py @@ -0,0 +1,225 @@ +"""Cross-process tests for LocalSharedStorage. + +An owner runs in a spawned child process; the pytest process attaches as a +client (the child wins election first). Exercises KV visibility, pub/sub, and +client reconnect-with-replay across a forced socket drop -- all over the real +socket transport, not the in-process fast path. +""" +import multiprocessing as mp +import threading +import time +import uuid + +import pytest + +from dash._shared_storage import LocalSharedStorage + +CTX = mp.get_context("spawn") + + +def _owner_main(namespace, cmd_q, done_q, ready_ev): + """A controllable owner: wins election, then runs commands on demand.""" + store = LocalSharedStorage(namespace=namespace) + store.start() + if not store._coord.is_owner(): # pragma: no cover - defensive + done_q.put(("error", "child did not win election")) + return + ready_ev.set() + while True: + cmd = cmd_q.get() + if cmd is None: + break + op, args = cmd + if op == "set": + store.set(*args) + elif op == "publish": + store.publish(*args) + done_q.put(("done", op)) + store.close() + + +class _Owner: + """Test helper: spawn a controllable owner and drive it synchronously.""" + + def __init__(self, namespace): + self.cmd_q = CTX.Queue() + self.done_q = CTX.Queue() + self.ready = CTX.Event() + self.proc = CTX.Process( + target=_owner_main, + args=(namespace, self.cmd_q, self.done_q, self.ready), + daemon=True, + ) + + def start(self): + self.proc.start() + assert self.ready.wait(timeout=10), "owner failed to start" + + def do(self, op, *args): + self.cmd_q.put((op, args)) + kind, _ = self.done_q.get(timeout=10) + assert kind == "done" + + def stop(self): + self.cmd_q.put(None) + self.proc.join(timeout=10) + + +@pytest.fixture +def owner(): + ns = f"xproc-{uuid.uuid4().hex[:12]}" + o = _Owner(ns) + o.start() + try: + yield ns, o + finally: + o.stop() + + +def test_kv_visible_across_processes(owner): + ns, o = owner + client = LocalSharedStorage(namespace=ns) + client.start() + assert not client._coord.is_owner() # child owns; we are the client + + o.do("set", "shared", {"n": 42}) + assert client.get("shared") == {"n": 42} + assert client.get("nope", "fallback") == "fallback" + client.close() + + +def test_pubsub_across_processes(owner): + ns, o = owner + client = LocalSharedStorage(namespace=ns) + client.start() + + sub = client.subscribe("topic") # cursor at current head + received = [] + + def consume(): + for msg in sub: + received.append(msg) + if len(received) == 3: + break + + th = threading.Thread(target=consume) + th.start() + time.sleep(0.3) # ensure the long-poll is established + + for i in range(3): + o.do("publish", "topic", f"m{i}") + + th.join(timeout=10) + sub.close() + client.close() + assert received == ["m0", "m1", "m2"] + + +def test_client_reconnect_replays_missed_messages(owner): + ns, o = owner + client = LocalSharedStorage(namespace=ns) + client.start() + + sub = client.subscribe("stream") + received = [] + ready = threading.Event() + + def consume(): + ready.set() + for msg in sub: + received.append(msg) + if len(received) == 6: + break + + th = threading.Thread(target=consume) + th.start() + ready.wait() + time.sleep(0.3) + + o.do("publish", "stream", "a") + o.do("publish", "stream", "b") + + # Wait until the first two land, then forcibly drop the client's socket. + _wait_until(lambda: len(received) >= 2, timeout=5) + conn = sub._conn + if conn is not None: + conn.close() # simulate a proxy idle-timeout / network blip + + # Messages published during/after the drop must still arrive (buffer replay). + for msg in ("c", "d", "e", "f"): + o.do("publish", "stream", msg) + + th.join(timeout=10) + sub.close() + client.close() + assert received == ["a", "b", "c", "d", "e", "f"] + + +def test_reelection_after_owner_killed(owner): + ns, o = owner + client = LocalSharedStorage(namespace=ns) + client.start() + assert not client._coord.is_owner() + + o.do("set", "before", "value") + assert client.get("before") == "value" + + # Kill the owner without a clean shutdown (leaves a stale socket). + o.proc.terminate() + o.proc.join(timeout=10) + + # The survivor re-elects: it becomes the new (cold) owner and keeps serving. + assert client.get("before") is None # cold store, prior state gone + client.set("after", "fresh") + assert client.get("after") == "fresh" + assert client._coord.is_owner() # we are the new owner + client.close() + + +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: + if pred(): + return + time.sleep(0.02) + raise AssertionError("condition not met in time") diff --git a/tests/shared_storage/test_stream_hub.py b/tests/shared_storage/test_stream_hub.py new file mode 100644 index 0000000000..69d0d82665 --- /dev/null +++ b/tests/shared_storage/test_stream_hub.py @@ -0,0 +1,144 @@ +"""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 ( + apump_to_storage, + publish_frame, + pump_to_storage, + 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}] 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..0abf821210 --- /dev/null +++ b/tests/streaming/conftest.py @@ -0,0 +1,20 @@ +"""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 + + +@pytest.fixture(autouse=True) +def _isolate_shared_storage(monkeypatch): + namespace = f"streamtest-{uuid.uuid4().hex[:12]}" + monkeypatch.setattr(_local, "_default_namespace", lambda: namespace) + yield diff --git a/tests/streaming/test_stream_callbacks_integration.py b/tests/streaming/test_stream_callbacks_integration.py new file mode 100644 index 0000000000..b32802e077 --- /dev/null +++ b/tests/streaming/test_stream_callbacks_integration.py @@ -0,0 +1,204 @@ +"""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): + yield "step-1" + await asyncio.sleep(0.5) + yield "step-2" + await asyncio.sleep(0.5) + 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() == [] diff --git a/tests/streaming/test_stream_callbacks_unit.py b/tests/streaming/test_stream_callbacks_unit.py new file mode 100644 index 0000000000..2f8b65d4eb --- /dev/null +++ b/tests/streaming/test_stream_callbacks_unit.py @@ -0,0 +1,382 @@ +"""Unit tests for streaming (generator) callbacks - no browser required.""" +import asyncio +import contextvars +import json +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._streaming import ( + StreamedCallbackResponse, + _keepalive_frames, + 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] diff --git a/tests/streaming/test_stream_transport.py b/tests/streaming/test_stream_transport.py new file mode 100644 index 0000000000..a85c3e956f --- /dev/null +++ b/tests/streaming/test_stream_transport.py @@ -0,0 +1,137 @@ +"""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._shared_storage import LocalSharedStorage +from dash._stream_hub import subscribe_envelopes + + +def _uplink_body(connection_id, request_id): + return { + "output": "out.children", + "outputs": {"id": "out", "property": "children"}, + "inputs": [{"id": "btn", "property": "n_clicks", "value": 1}], + "changedPropIds": ["btn.n_clicks"], + "streamConnection": { + "connectionId": connection_id, + "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, "c1", out) + + resp = app.server.test_client().post( + "/_dash-update-component", json=_uplink_body("c1", "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_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, "c1", out) + + with TestClient(server) as client: + resp = client.post("/_dash-update-component", json=_uplink_body("c1", "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, "c1", out) + + async def run(): + client = server.test_client() + resp = await client.post( + "/_dash-update-component", json=_uplink_body("c1", "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() 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() == []