diff --git a/docs/events/page_load_events.md b/docs/events/page_load_events.md index 851e09ccce0..72899f41ae9 100644 --- a/docs/events/page_load_events.md +++ b/docs/events/page_load_events.md @@ -41,6 +41,17 @@ def index(): return rx.text("A Beautiful App") ``` +## Cancellation + +Unfinished `on_load` event chains are cancelled when the client disconnects, +navigates to another page or reconnects. This includes events returned or yielded +by the page-load handler. A reconnect starts the new page's load events again. +Independent background events started outside the page-load chain continue +running after a disconnect. + +Use `try`/`finally` to release resources when an async handler is cancelled. Avoid +suppressing `asyncio.CancelledError`, which prevents the handler from stopping. + ## Handling Loading and Errors Avoid heavy synchronous work directly in an `on_load` handler — the page renders before the handler finishes, so a slow handler leaves the user staring at stale or empty data with no feedback. Instead: diff --git a/news/+cancel-disconnected-page-loads.bugfix.md b/news/+cancel-disconnected-page-loads.bugfix.md new file mode 100644 index 00000000000..816760396be --- /dev/null +++ b/news/+cancel-disconnected-page-loads.bugfix.md @@ -0,0 +1 @@ +Cancel unfinished page-load event chains when the client disconnects, including hydration that has not yet started. Independent background events continue running. diff --git a/packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md b/packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md new file mode 100644 index 00000000000..f5bbec43a14 --- /dev/null +++ b/packages/reflex-base/news/+cancel-disconnected-page-loads.bugfix.md @@ -0,0 +1 @@ +Cancel pending hydration and page-load event chains when their client disconnects. diff --git a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py index 876f35e2d2f..389fc054ece 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py +++ b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py @@ -31,7 +31,7 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: - from reflex.state import BaseState + from reflex.state import BaseState, State # Resolved once at import: find_spec on a missing package scans sys.path (~90us), # far too slow for the per-event-argument path below. @@ -394,23 +394,17 @@ class BaseStateEventProcessor(EventProcessor): """ @staticmethod - async def _rehydrate(root_state: BaseState): - """Rehydrate the state by calling the hydrate event handler. + async def _rehydrate_snapshot(root_state: State) -> None: + """Send hydration data, finishing hydration when no page is available. Args: root_state: The root state to rehydrate. """ - from reflex.state import OnLoadInternalState, State - - if type(root_state) is not State: - return + from reflex.state import State - # A backend-initiated event carries no route, so nothing sets - # router_data and every one of them would rehydrate again. routeless = not root_state.router_data if routeless and root_state.is_hydrated: return - await process_event( handler=State.event_handlers["hydrate"], payload={}, @@ -418,20 +412,70 @@ async def _rehydrate(root_state: BaseState): root_state=root_state, ) if routeless: - # No page to load, but hydration still has to finish. await process_event( handler=State.event_handlers["set_is_hydrated"], payload={"value": True}, state=root_state, root_state=root_state, ) + + @functools.cached_property + def _pending_hydrates(self) -> dict[str, dict[str, EventFuture]]: + """Hydrations cancellable on disconnect, but never by a navigation. + + Returns: + Pending hydration futures indexed by client token and transaction id. + """ + return {} + + async def _rehydrate(self, root_state: BaseState): + """Rehydrate the state by calling the hydrate event handler. + + Args: + root_state: The root state to rehydrate. + """ + from reflex.event import Event + from reflex.state import OnLoadInternalState, State + + if type(root_state) is not State: + return + + # A backend-initiated event carries no route, so nothing sets + # router_data and every one of them would rehydrate again. + if not root_state.router_data: + await self._rehydrate_snapshot(root_state) return - await process_event( - handler=OnLoadInternalState.event_handlers["on_load_internal"], - payload={}, - state=await root_state.get_state(OnLoadInternalState), - root_state=root_state, - ) + + ctx = EventContext.get() + # Disconnect cancels this marker without cancelling the ordinary event. + pending = EventFuture(txid=ctx.txid) + self._pending_hydrates.setdefault(ctx.token, {})[pending.txid] = pending + try: + await self._rehydrate_snapshot(root_state) + if pending.cancelled(): + return + parent = self._futures[ctx.txid] + first_child = len(parent.children) + handler = OnLoadInternalState.event_handlers["on_load_internal"] + await process_event( + handler=handler, + payload={}, + state=await root_state.get_state(OnLoadInternalState), + root_state=root_state, + ) + # Inline rehydration bypasses enqueue's page-load registration. Track + # only its children so the ordinary event remains independent. + load_event = Event(name=format_event_handler(handler)) + for index in range(first_child, len(parent.children)): + child = parent.children[index] + if pending.cancelled(): + child.cancel() + else: + self._supersede_previous( + token=ctx.token, event=load_event, tracked=child + ) + finally: + self._forget_hydrate(ctx.token, pending) def _supersede_previous( self, *, token: str, event: Event, tracked: EventFuture @@ -455,8 +499,36 @@ def _supersede_previous( boot_name, on_load_name = _connect_supersedes() if event.name == boot_name: self._cancel_older_chains((on_load_name, token), tracked.root_gen) + self._pending_hydrates.setdefault(token, {})[tracked.txid] = tracked + tracked.add_done_callback(functools.partial(self._forget_hydrate, token)) return super()._supersede_previous(token=token, event=event, tracked=tracked) + def _forget_hydrate(self, token: str, future: EventFuture) -> None: + """Stop tracking a completed hydration invocation. + + Its page-load descendants have their own supersession registration. + + Args: + token: The client token. + future: The completed hydration future. + """ + if (pending := self._pending_hydrates.get(token)) is not None: + pending.pop(future.txid, None) + if not pending: + del self._pending_hydrates[token] + + def _on_disconnect(self, token: str) -> None: + """Cancel hydration and unfinished page-load chains for a disconnected client. + + Args: + token: The disconnected client token. + """ + _, on_load_name = _connect_supersedes() + for future in self._pending_hydrates.pop(token, {}).values(): + future.cancel() + for future in self._superseded.pop((on_load_name, token), {}).values(): + future.cancel() + async def _execute_event( self, *, entry: EventQueueEntry, registered_handler: RegisteredEventHandler ) -> None: diff --git a/packages/reflex-base/src/reflex_base/event/processor/event_processor.py b/packages/reflex-base/src/reflex_base/event/processor/event_processor.py index 071655da946..51f7e1231be 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/event_processor.py +++ b/packages/reflex-base/src/reflex_base/event/processor/event_processor.py @@ -460,6 +460,15 @@ async def enqueue( await queue.put(EventQueueEntry(event=event, ctx=ev_ctx)) return tracked + def _on_disconnect(self, token: str) -> None: + """Handle a client disconnect without cancelling independent event work. + + Subclasses may cancel work tied to the disconnected page's lifecycle. + + Args: + token: The disconnected client token. + """ + async def enqueue_many(self, token: str, *events: Event) -> Sequence[EventFuture]: """Enqueue multiple events to be processed. @@ -585,15 +594,15 @@ def _try_clean_future(self, future: EventFuture) -> None: # type: ignore[overri return parent = future.parent self._futures.pop(future.txid, None) - if ( - (key := future.supersede_key) is not None - and (slot := self._superseded.get(key)) is not None - and slot.get(future.txid) is future - and future.all_done() - ): - del slot[future.txid] - if not slot: - del self._superseded[key] + if future.supersede_key is not None and future.all_done(): + # Inline page loads can also have their own supersession key. + for key in future.covered_supersede_keys: + if (slot := self._superseded.get(key)) is not None and slot.get( + future.txid + ) is future: + del slot[future.txid] + if not slot: + del self._superseded[key] if parent is None or not parent.txid: return future = parent diff --git a/reflex/app.py b/reflex/app.py index 588e1daf835..97d7d7472e5 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -2319,6 +2319,14 @@ def on_disconnect(self, sid: str) -> asyncio.Task | None: # Get token before cleaning up disconnect_token = self.sid_to_token.get(sid) if disconnect_token: + socket_record = self._token_manager.token_to_socket.get(disconnect_token) + if ( + socket_record is not None + and socket_record.sid == sid + and socket_record.instance_id == self._token_manager.instance_id + and (processor := self.app._event_processor) is not None + ): + processor._on_disconnect(disconnect_token) # Use async cleanup through token manager task = asyncio.create_task( self._token_manager.disconnect_token(disconnect_token, sid), diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 67abb9e73b9..118ae02d922 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -281,7 +281,7 @@ async def __aenter__(self) -> Self: # next event loads its current page. from reflex_base.event.processor import BaseStateEventProcessor - await BaseStateEventProcessor._rehydrate(mutable_state) + await BaseStateEventProcessor._rehydrate_snapshot(mutable_state) except (Exception, asyncio.CancelledError): # Restore the proxy to a consistent state since __aexit__ will not be called when __aenter__ raises. await self.__aexit__(*sys.exc_info()) diff --git a/tests/units/conftest.py b/tests/units/conftest.py index 87c64245b0b..e4e83140c64 100644 --- a/tests/units/conftest.py +++ b/tests/units/conftest.py @@ -350,6 +350,9 @@ def mock_base_state_event_processor_obj( A fresh BaseState event processor. """ monkeypatch.setattr(BaseStateEventProcessor, "_rehydrate", mock.AsyncMock()) + monkeypatch.setattr( + BaseStateEventProcessor, "_rehydrate_snapshot", mock.AsyncMock() + ) def handle_backend_exception(ex: Exception) -> None: formatted_exc = "\n".join(traceback.format_exception(ex)) diff --git a/tests/units/istate/test_proxy.py b/tests/units/istate/test_proxy.py index 7be7a6744e8..d21edf95bd2 100644 --- a/tests/units/istate/test_proxy.py +++ b/tests/units/istate/test_proxy.py @@ -198,7 +198,9 @@ async def failing_rehydrate(root_state: BaseState): # noqa: RUF029 msg = "Simulated rehydrate failure" raise RuntimeError(msg) - monkeypatch.setattr(BaseStateEventProcessor, "_rehydrate", failing_rehydrate) + monkeypatch.setattr( + BaseStateEventProcessor, "_rehydrate_snapshot", failing_rehydrate + ) with pytest.raises(RuntimeError, match="Simulated rehydrate failure"): async with state_proxy: pass diff --git a/tests/units/reflex_base/event/processor/test_base_state_processor.py b/tests/units/reflex_base/event/processor/test_base_state_processor.py index b1613881257..7cc5b9be3d3 100644 --- a/tests/units/reflex_base/event/processor/test_base_state_processor.py +++ b/tests/units/reflex_base/event/processor/test_base_state_processor.py @@ -31,7 +31,7 @@ import reflex as rx from reflex import event -from reflex.app import App +from reflex.app import App, EventNamespace from reflex.compiler.utils import compile_state from reflex.event import Event, EventHandler, EventSpec from reflex.istate.delta import Delta @@ -2740,3 +2740,286 @@ async def slow_load(self): async with _read_back(real_base_state_processor, token) as root: assert (await root.get_state(SlowLoadState)).seen != "stale-finished" assert (await root.get_state(State)).is_hydrated is True + + +@pytest.mark.parametrize( + "processor_state_manager", ["in_process", "redis"], indirect=True +) +@pytest.mark.parametrize("self_chain", [False, True]) +@pytest.mark.parametrize("rehydrate", [False, True]) +@pytest.mark.parametrize("supersedes", [False, True]) +async def test_disconnect_cancels_page_load_chain( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, + self_chain: bool, + rehydrate: bool, + supersedes: bool, +): + """Disconnect cancels a running load even after its root handler has returned. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + self_chain: Whether to reach the suspended handler through a self-chain. + rehydrate: Whether an ordinary routed event triggers the page load. + supersedes: Whether the user's load handler also supersedes earlier calls. + """ + started = asyncio.Event() + cancelled = asyncio.Event() + invocations = 0 + + class DisconnectState(State): + @event + def ordinary(self): + """Trigger compatibility hydration on a fresh client state.""" + + @event(supersedes=supersedes) + async def load(self): + """Suspend a load or its self-chained successor until cancellation. + + Yields: + A successor event or an empty update before suspending. + """ + nonlocal invocations + invocations += 1 + if self_chain and invocations == 1: + yield DisconnectState.load + return + yield + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + wired_app.add_page(lambda: rx.text("load"), route="/", on_load=DisconnectState.load) + wired_app._event_processor = real_base_state_processor + namespace = EventNamespace("/_event", wired_app) + await namespace._token_manager.link_token_to_sid(token, "sid") + + async with real_base_state_processor as processor: + stale = await processor.enqueue( + token, + _client_event( + DisconnectState.ordinary() + if rehydrate + else OnLoadInternalState.on_load_internal(), + _view("/"), + ), + ) + await asyncio.wait_for(started.wait(), timeout=5) + cleanup = namespace.on_disconnect("sid") + assert cleanup is not None + await cleanup + await asyncio.wait_for(cancelled.wait(), timeout=1) + await asyncio.wait_for(stale.wait_all(), timeout=5) + assert invocations == (2 if self_chain else 1) + # Cancellation releases the lock so a later connection can use the state. + async with _read_back(processor, token) as root: + assert await root.get_state(DisconnectState) is not None + await asyncio.sleep(0) + assert not processor._superseded + + +@pytest.mark.parametrize( + "processor_state_manager", ["in_process", "redis"], indirect=True +) +@pytest.mark.parametrize("disconnect", [False, True]) +async def test_disconnect_during_inline_hydration( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, + disconnect: bool, +): + """A disconnect during the hydration emit prevents late page-load work. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + disconnect: Whether to disconnect while sending the hydration update. + """ + sending = asyncio.Event() + release = asyncio.Event() + invocations = [] + + class InlineHydrationState(State): + @event + def ordinary(self): + """Record that independent work still runs after hydration.""" + invocations.append("ordinary") + + @event + def load(self): + """Record whether the page load ran.""" + invocations.append("load") + + async def pause_emit(token: str, delta: Mapping[str, Mapping[str, Any]]) -> None: + """Pause sending the hydration snapshot. + + Args: + token: The client token. + delta: The state update. + """ + sending.set() + await release.wait() + + processor = real_base_state_processor + assert processor._root_context is not None + processor._root_context = dataclasses.replace( + processor._root_context, emit_delta_impl=pause_emit + ) + wired_app.add_page( + lambda: rx.text("load"), route="/", on_load=InlineHydrationState.load + ) + wired_app._event_processor = processor + namespace = EventNamespace("/_event", wired_app) + await namespace._token_manager.link_token_to_sid(token, "sid") + + try: + async with processor: + ordinary = await processor.enqueue( + token, _client_event(InlineHydrationState.ordinary(), _view("/")) + ) + try: + await asyncio.wait_for(sending.wait(), timeout=5) + if disconnect: + cleanup = namespace.on_disconnect("sid") + assert cleanup is not None + await cleanup + await namespace._token_manager.link_token_to_sid(token, "new-sid") + finally: + release.set() + await asyncio.wait_for(ordinary.wait_all(), timeout=5) + await asyncio.sleep(0) + assert not processor._pending_hydrates + assert not processor._superseded + finally: + await namespace._token_manager.disconnect_all() + await namespace._token_manager.close() + + assert invocations == (["ordinary"] if disconnect else ["ordinary", "load"]) + + +@pytest.mark.parametrize("disconnect", [False, True]) +async def test_disconnect_cancels_queued_hydrate( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, + disconnect: bool, +): + """A disconnect before hydration starts must not launch a page-load chain. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + disconnect: Whether the client disconnects before hydration runs. + """ + loads = [] + + class DisconnectState(State): + @event + def load(self): + """Record that the page-load handler ran.""" + loads.append("loaded") + + wired_app.add_page(lambda: rx.text("load"), route="/", on_load=DisconnectState.load) + wired_app._event_processor = real_base_state_processor + namespace = EventNamespace("/_event", wired_app) + await namespace._token_manager.link_token_to_sid(token, "sid") + async with real_base_state_processor as processor: + future = await processor.enqueue( + token, _client_event(State.hydrate_and_load(), _view("/")) + ) + if disconnect: + cleanup = namespace.on_disconnect("sid") + assert cleanup is not None + await cleanup + assert future.cancelled() + else: + await asyncio.wait_for(future.wait_all(), timeout=5) + await namespace._token_manager.disconnect_all() + assert loads == ([] if disconnect else ["loaded"]) + assert not processor._pending_hydrates + + +@pytest.mark.parametrize("routed", [False, True]) +async def test_disconnect_preserves_independent_events( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, + routed: bool, +): + """Independent foreground and background work survives another chain's disconnect. + + Args: + wired_app: The app sharing the processor's state manager. + real_base_state_processor: The real event processor. + token: The client token. + routed: Whether the foreground event also schedules compatibility page loads. + """ + started = {name: asyncio.Event() for name in ("foreground", "background", "other")} + release = asyncio.Event() + finished = [] + loads = [] + + class IndependentState(State): + @event + def load(self): + """Record whether the disconnected page-load work ran.""" + loads.append("loaded") + + @event + async def foreground(self, name: str): + """Wait for release and record the named foreground invocation. + + Args: + name: The invocation to record. + """ + started[name].set() + await release.wait() + finished.append(name) + + @event(background=True) + async def background(self): + """Wait for release independently of the page-load lifecycle.""" + started["background"].set() + await release.wait() + finished.append("background") + + wired_app.add_page( + lambda: rx.text("load"), route="/", on_load=IndependentState.load + ) + async with real_base_state_processor as processor: + # Start the background event before the foreground handler holds the lock. + background = await processor.enqueue( + token, Event.from_event_type(IndependentState.background())[0] + ) + await asyncio.wait_for(started["background"].wait(), timeout=5) + foreground = await processor.enqueue( + token, + _client_event(IndependentState.foreground("foreground"), _view("/")) + if routed + else Event.from_event_type(IndependentState.foreground("foreground"))[0], + ) + other = await processor.enqueue( + "other-token", + Event.from_event_type(IndependentState.foreground("other"))[0], + ) + await asyncio.wait_for( + asyncio.gather(*(signal.wait() for signal in started.values())), timeout=5 + ) + processor._on_disconnect(token) + assert not any(future.cancelled() for future in (background, foreground, other)) + release.set() + await asyncio.wait_for( + asyncio.gather( + *(future.wait_all() for future in (background, foreground, other)) + ), + timeout=5, + ) + assert sorted(finished) == ["background", "foreground", "other"] + assert not loads diff --git a/tests/units/test_app.py b/tests/units/test_app.py index a29198162f6..ae575628094 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -4922,6 +4922,37 @@ async def event_namespace_with_processor_mock() -> AsyncGenerator[EventNamespace await event_namespace._token_manager.disconnect_all() +@pytest.mark.parametrize("owner", ["current", "replacement", "remote", "missing"]) +async def test_disconnect_only_cancels_current_local_socket( + event_namespace_with_processor_mock: EventNamespace, + owner: str, +): + """A stale disconnect must not cancel a replacement connection's page load. + + Args: + event_namespace_with_processor_mock: Namespace with a mocked processor. + owner: The socket currently owning the client token. + """ + namespace = event_namespace_with_processor_mock + manager = namespace._token_manager + manager.sid_to_token["old-sid"] = "token" + if owner != "missing": + manager.token_to_socket["token"] = SocketRecord( + instance_id="remote" if owner == "remote" else manager.instance_id, + sid="new-sid" if owner == "replacement" else "old-sid", + ) + manager.disconnect_token = AsyncMock() + cleanup = namespace.on_disconnect("old-sid") + assert cleanup is not None + await cleanup + cancel = cast(Mock, namespace.app.event_processor._on_disconnect) + if owner == "current": + cancel.assert_called_once_with("token") + else: + cancel.assert_not_called() + assert namespace.on_disconnect("unknown-sid") is None + + def _connect_environ(token: str) -> dict[str, Any]: return { "QUERY_STRING": f"token={token}",