Streaming: verify connections, stop on shutdown, reset stale cursor · plotly/dash@7ae8d52 · GitHub
Skip to content

Commit 7ae8d52

Browse files
T4rk1nT4rk1n
committed
Streaming: verify connections, stop on shutdown, reset stale cursor
Addresses the review comments on the multiplexed streaming transport. - Authorize stream connections. The connection topic is now keyed on the server-signed end_id, derived server-side, never from a client-supplied id, so a page can only read or write its own streams. A stream request whose token does not verify is refused with 403 and never run some other way, so no frame can reach a topic without a valid token. Multiplexing across worker processes needs a secret_key so every worker verifies the same token (the same requirement background callbacks have); the callback fails visibly in the browser otherwise. - Stop streams on server shutdown. In-flight pump tasks are cancelled and open downlink subscriptions closed on shutdown (FastAPI shutdown event, Quart after_serving, Flask atexit), so a long-polling downlink no longer keeps a streaming app from exiting on Ctrl+C. - Reset a stale cursor. A cursor ahead of the topic head (server restart or owner re-election left the topic at seq 0) is surfaced as a gap; the downlink emits a reset envelope and the client resets its cursor to the head instead of stalling until the fresh sequence climbs past it.
1 parent 07f334d commit 7ae8d52

15 files changed

Lines changed: 501 additions & 71 deletions

File tree

‎CHANGELOG.md‎

Lines changed: 1 addition & 1 deletion

‎dash/_callback.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,24 @@ def get_request_end_id(secret: bytes):
445445
return _callback_signing.unsign(secret, _callback_signing.END_SCOPE, token)
446446

447447

448+
def get_stream_connection_id() -> "str | None":
449+
"""Return the verified streaming connection id for the request, or ``None``.
450+
451+
The multiplexed streaming transport keys each page's downlink topic on this
452+
id, so it must be unforgeable: a client that could name an arbitrary topic
453+
could read another page's stream or inject frames into it. It is derived from
454+
the server-signed ``end_id`` (the same per-page-load token background-callback
455+
handles are bound to), never from anything the client picks. A missing or
456+
forged token yields ``None``, and the backend refuses the request (403).
457+
458+
``end_id`` is signed with the server secret, so across worker processes every
459+
worker must resolve the same secret: set a ``secret_key`` on the server, or
460+
cross-worker stream requests will not verify. Single-process apps are fine
461+
with no configuration.
462+
"""
463+
return get_request_end_id(_get_signing_secret())
464+
465+
448466
def _get_signing_secret() -> bytes:
449467
return get_app()._get_signing_secret() # pylint: disable=protected-access
450468

‎dash/_shared_storage/_engine.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,14 @@ def poll(self, topic: str, after_seq: int, timeout: float) -> PollResult:
163163
while True:
164164
if self._closed:
165165
return PollResult([], after_seq, False)
166+
# The cursor points past every sequence this topic has ever
167+
# produced. That can only happen when the cursor was minted by a
168+
# previous incarnation of the topic (the owner was re-elected, or
169+
# the server restarted with an empty store) -- treat it as a gap
170+
# so the consumer resets instead of stalling until the fresh
171+
# sequence climbs back past the stale cursor.
172+
if after_seq > t.seq:
173+
return PollResult([], after_seq, True)
166174
# The next message we want is after_seq + 1; if the buffer's
167175
# oldest is newer than that, it was evicted -> gap.
168176
if t.buf and after_seq + 1 < t.buf[0][0]:

‎dash/_shared_storage/diskcache.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,12 @@ def _poll(self, topic: str, after_seq: int, timeout: float) -> PollResult:
113113
deadline = time.monotonic() + timeout
114114
while True:
115115
head = self._head(topic)
116+
# Cursor past the head: it was minted before the store was reset
117+
# (the cache was cleared, or the counter evicted under the cache's
118+
# size limit). Gap so the consumer resets rather than blocking until
119+
# the sequence climbs back past the cursor.
120+
if after_seq > head:
121+
return PollResult([], after_seq, True)
116122
if head > after_seq:
117123
floor = max(1, head - self._buffer_size + 1)
118124
if after_seq + 1 < floor:

‎dash/_shared_storage/redis.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,12 @@ def _head(self, topic: str) -> int:
137137

138138
def _poll(self, topic: str, after_seq: int, timeout: float) -> PollResult:
139139
stream = self._stream(topic)
140+
# Cursor past the head: it was minted before the stream was reset (the
141+
# key was flushed, or evicted under a maxmemory policy). Gap so the
142+
# consumer resets rather than blocking on XREAD until the sequence climbs
143+
# back past the cursor.
144+
if after_seq > self._head(topic):
145+
return PollResult([], after_seq, True)
140146
# Gap: the next wanted sequence sits below the trimmed floor. Checked
141147
# before XREAD, which would otherwise silently resume at the floor.
142148
first = self._redis.xrange(stream, count=1)

‎dash/_stream_hub.py‎

Lines changed: 52 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,10 @@
2323

2424
import asyncio
2525
import json
26+
import threading
2627
from typing import Any, AsyncIterator, Iterator, Optional
2728

28-
from ._shared_storage.base import BaseSharedStorage
29+
from ._shared_storage.base import BaseSharedStorage, SharedStorageGap, Subscription
2930
from ._streaming import StreamedCallbackResponse, sync_iter_asyncgen, to_json
3031

3132
_TOPIC_PREFIX = "_dash_stream:"
@@ -34,6 +35,11 @@
3435
# immediately; its outputs arrive on the downlink, not this response.
3536
STREAM_ACK = {"multi": True, "stream": True}
3637

38+
# Control envelope telling the client its cursor is stale (its frames were lost
39+
# to an owner re-election or a server restart) and it must reset to the head and
40+
# resubscribe, rather than stall waiting for the fresh sequence to pass it.
41+
RESET_ENVELOPE = {"reset": True}
42+
3743

3844
def stream_topic(connection_id: str) -> str:
3945
return f"{_TOPIC_PREFIX}{connection_id}"
@@ -57,6 +63,14 @@ def publish_frame(
5763
storage.publish(stream_topic(connection_id), {"rid": request_id, "frame": plain})
5864

5965

66+
# Open downlink subscriptions, so a server shutdown can close them (each one
67+
# otherwise blocks its worker in a long poll, stalling a graceful shutdown).
68+
# Guarded by a lock: subscriptions open/close on worker threads while a shutdown
69+
# hook iterates the set, and a plain set is not safe against that.
70+
_active_subscriptions: "set[Subscription]" = set()
71+
_registry_lock = threading.Lock()
72+
73+
6074
def subscribe_envelopes(
6175
storage: BaseSharedStorage,
6276
connection_id: str,
@@ -69,11 +83,22 @@ def subscribe_envelopes(
6983
long-lived: it carries frames for every callback on the connection, not one
7084
stream, and ends when the client hangs up. ``replay_from`` resumes a
7185
reconnecting downlink from its last cursor. Each envelope carries its ``seq``
72-
so the client can resume from it after a reconnect without losing frames.
86+
so the client can resume from it after a reconnect without losing frames. A
87+
lost buffer (owner re-election, server restart) surfaces as a single reset
88+
envelope so the client resets its cursor instead of stalling.
7389
"""
74-
with storage.subscribe(stream_topic(connection_id), replay_from) as sub:
90+
sub = storage.subscribe(stream_topic(connection_id), replay_from)
91+
with _registry_lock:
92+
_active_subscriptions.add(sub)
93+
try:
7594
for seq, message in sub.iter_with_seq():
7695
yield {**message, "seq": seq}
96+
except SharedStorageGap:
97+
yield dict(RESET_ENVELOPE)
98+
finally:
99+
with _registry_lock:
100+
_active_subscriptions.discard(sub)
101+
sub.close()
77102

78103

79104
async def asubscribe_envelopes(
@@ -83,10 +108,16 @@ async def asubscribe_envelopes(
83108
) -> AsyncIterator[Any]:
84109
"""Async counterpart of :func:`subscribe_envelopes` for ASGI backends."""
85110
sub = storage.subscribe(stream_topic(connection_id), replay_from)
111+
with _registry_lock:
112+
_active_subscriptions.add(sub)
86113
try:
87114
async for seq, message in sub.aiter_with_seq():
88115
yield {**message, "seq": seq}
116+
except SharedStorageGap:
117+
yield dict(RESET_ENVELOPE)
89118
finally:
119+
with _registry_lock:
120+
_active_subscriptions.discard(sub)
90121
sub.close()
91122

92123

@@ -149,3 +180,21 @@ def spawn_async_pump(
149180
)
150181
_pending_pumps.add(task)
151182
task.add_done_callback(_pending_pumps.discard)
183+
184+
185+
def shutdown_active_streams() -> None:
186+
"""Stop every in-flight stream so the server can shut down.
187+
188+
Cancels the background pump tasks that drive streaming callbacks and closes
189+
the open downlink subscriptions. Each downlink otherwise sits in a long poll
190+
that a graceful shutdown would wait on forever (the reason a streaming app
191+
can ignore Ctrl+C). Backends call this from their shutdown hook, mirroring
192+
how ``ws.py`` tears down WebSocket connections. Idempotent and safe to call
193+
when nothing is streaming.
194+
"""
195+
for task in list(_pending_pumps):
196+
task.cancel()
197+
with _registry_lock:
198+
subscriptions = list(_active_subscriptions)
199+
for sub in subscriptions:
200+
sub.close()

‎dash/backends/_fastapi.py‎

Lines changed: 32 additions & 7 deletions

0 commit comments

Comments
 (0)