From 2db865caa9d77cd4dec7f776498a101e9502dba0 Mon Sep 17 00:00:00 2001 From: Farhan Date: Tue, 7 Jul 2026 03:29:35 +0500 Subject: [PATCH] fix: cancel stale on_load chains via handler-declared supersession Navigating away while a page's on_load chain is still running left the stale chain executing, blocking the new page's events behind it and applying its late deltas (#6593). Add a SUPERSEDES_MARKER for event handlers with latest-wins semantics: enqueuing a new chain-root invocation cancels the previous unfinished event chain rooted at the same handler for the same client token. Mark on_load_internal with it so a newer navigation supersedes the previous page's unfinished load. The EventProcessor tracks the active chain root per (event name, token), drops events chained from an already-cancelled parent before they enter the queue, and defers future cleanup while the handler task is still unwinding so late-chained events can find their cancelled parent. --- news/6593.bugfix.md | 1 + .../src/reflex_base/event/__init__.py | 15 ++ .../event/processor/event_processor.py | 75 +++++++- .../src/reflex_base/event/processor/future.py | 4 + reflex/state.py | 15 +- .../event/processor/test_event_processor.py | 180 ++++++++++++++++++ tests/units/test_state.py | 84 ++++++++ 7 files changed, 365 insertions(+), 9 deletions(-) create mode 100644 news/6593.bugfix.md diff --git a/news/6593.bugfix.md b/news/6593.bugfix.md new file mode 100644 index 00000000000..8ff61800b11 --- /dev/null +++ b/news/6593.bugfix.md @@ -0,0 +1 @@ +Cancel the previous unfinished `on_load` event chain when a newer page navigation arrives for the same client, instead of letting stale page-load work block and outlive the navigation. diff --git a/packages/reflex-base/src/reflex_base/event/__init__.py b/packages/reflex-base/src/reflex_base/event/__init__.py index ff94adde754..9fd0e2b452f 100644 --- a/packages/reflex-base/src/reflex_base/event/__init__.py +++ b/packages/reflex-base/src/reflex_base/event/__init__.py @@ -185,6 +185,7 @@ def from_event_type( ) BACKGROUND_TASK_MARKER = "_reflex_background_task" +SUPERSEDES_MARKER = "_reflex_supersedes" EVENT_ACTIONS_MARKER = "_rx_event_actions" UPLOAD_FILES_CLIENT_HANDLER = "uploadFiles" @@ -442,6 +443,19 @@ def is_background(self) -> bool: """ return getattr(self.fn, BACKGROUND_TASK_MARKER, False) + @property + def supersedes(self) -> bool: + """Whether a newer chain-root invocation supersedes an older one. + + When True, enqueuing this handler as a chain root cancels the previous + unfinished event chain rooted at the same handler for the same client + token. + + Returns: + True if the event handler is marked as superseding. + """ + return getattr(self.fn, SUPERSEDES_MARKER, False) + def __call__(self, *args: Any, **kwargs: Any) -> "EventSpec": """Pass arguments to the handler to get an event spec. @@ -2778,6 +2792,7 @@ class EventNamespace: # Constants BACKGROUND_TASK_MARKER = BACKGROUND_TASK_MARKER + SUPERSEDES_MARKER = SUPERSEDES_MARKER EVENT_ACTIONS_MARKER = EVENT_ACTIONS_MARKER _EVENT_FIELDS = _EVENT_FIELDS FORM_DATA = FORM_DATA 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 7d9296fe4dd..a9243775a34 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 @@ -120,6 +120,11 @@ class EventProcessor: _futures: dict[str, EventFuture] = dataclasses.field( default_factory=dict, init=False ) + # Latest-wins tracking for superseding handlers: (event name, token) -> the + # currently active chain root future. + _superseded: dict[tuple[str, str], EventFuture] = dataclasses.field( + default_factory=dict, init=False + ) _token_queues: dict[ str, collections.deque[tuple[EventQueueEntry, RegisteredEventHandler]], @@ -311,6 +316,7 @@ async def stop(self, graceful_shutdown_timeout: float | None = None) -> None: self._queue_task = None # Discard any pending per-token queue entries. self._token_queues.clear() + self._superseded.clear() # Cancel any remaining unresolved futures. for future in self._futures.values(): if not future.done(): @@ -372,6 +378,8 @@ async def enqueue( Returns: An EventFuture that resolves to the result of the associated task. + If the event was chained from an already-cancelled chain, the + returned future is already cancelled and the event is dropped. """ if ev_ctx is None: try: @@ -395,7 +403,14 @@ async def enqueue( tracked.add_done_callback(self._on_future_done) # If this context has a parent, register as a child of the parent's future. if parent_future is not None: + if parent_future.cancelled(): + # The chain this event belongs to was cancelled; the event is + # stillborn and never enters the queue. + tracked.cancel() + return tracked parent_future.add_child(tracked) + if parent_future is None: + self._supersede_previous(token=token, event=event, tracked=tracked) await queue.put(EventQueueEntry(event=event, ctx=ev_ctx)) return tracked @@ -491,14 +506,51 @@ def _try_clean_future(self, future: EventFuture) -> None: # type: ignore[overri """ if not future.done(): return + if future.txid in self._tasks: + # The handler task is still running or unwinding; keep the future + # so late-chained events can find their (possibly cancelled) parent. + return # Not checking future.all_done() to avoid waiting for grandchildren here. if not all(c.done() for c in future.children): return parent = future.parent self._futures.pop(future.txid, None) + if ( + (key := future.supersede_key) is not None + and self._superseded.get(key) is future + and future.all_done() + ): + del self._superseded[key] if parent is not None and parent.txid: self._try_clean_future(parent) + def _supersede_previous( + self, *, token: str, event: Event, tracked: EventFuture + ) -> None: + """Cancel the previous unfinished chain of a superseding event handler. + + Root handlers marked with ``supersedes`` (e.g. ``on_load_internal``) + use latest-wins semantics: enqueuing a new invocation cancels the + previous unfinished event chain for the same handler and client token. + + Args: + token: The client token associated with the event. + event: The event being enqueued. + tracked: The future of the event being enqueued. + """ + try: + registered = RegistrationContext.get().event_handlers.get(event.name) + except LookupError: + return + if registered is None or not registered.handler.supersedes: + return + key = (event.name, token) + previous = self._superseded.get(key) + if previous is not None and not previous.all_done(): + previous.cancel() + self._superseded[key] = tracked + tracked.supersede_key = key + def _on_future_done(self, future: EventFuture) -> None: # type: ignore[override] """Callback invoked when an enqueued future completes. @@ -625,10 +677,13 @@ def _dispatch_next_for_token(self, token: str) -> None: if not token_queue: return entry, registered_handler = token_queue[0] - # Skip cancelled futures. + # Skip cancelled futures. Before a task exists, the only way a future + # can be done (and thus already cleaned up) is cancellation, so a + # missing future also means the entry was cancelled. future = self._futures.get(entry.ctx.txid) - if future is not None and future.cancelled(): - self._try_clean_future(future) + if future is None or future.cancelled(): + if future is not None: + self._try_clean_future(future) token_queue.popleft() if token_queue: self._dispatch_next_for_token(token) @@ -645,10 +700,12 @@ async def _process_queue(self): with contextlib.suppress(QueueShutDown): while True: entry = await queue.get() - if ( - future := self._futures.get(entry.ctx.txid) - ) is not None and future.cancelled(): - self._try_clean_future(future) + # A missing future means the entry was cancelled and already + # cleaned up (see _dispatch_next_for_token). + future = self._futures.get(entry.ctx.txid) + if future is None or future.cancelled(): + if future is not None: + self._try_clean_future(future) queue.task_done() continue try: @@ -757,6 +814,10 @@ def _finish_task(self, task: asyncio.Task): else: if future is not None and not future.done(): future.set_result(result) + if future is not None: + # The task is gone; clean up now in case the future resolved + # earlier (e.g. external cancellation) and cleanup was deferred. + self._try_clean_future(future) __all__ = [ diff --git a/packages/reflex-base/src/reflex_base/event/processor/future.py b/packages/reflex-base/src/reflex_base/event/processor/future.py index 01d27fbdcef..2eca7b585c4 100644 --- a/packages/reflex-base/src/reflex_base/event/processor/future.py +++ b/packages/reflex-base/src/reflex_base/event/processor/future.py @@ -31,6 +31,10 @@ class EventFuture(asyncio.Future): default_factory=asyncio.get_running_loop, repr=False ) + # Key under which this future is registered for latest-wins supersession + # in the EventProcessor, if any. + supersede_key: tuple[str, str] | None = dataclasses.field(default=None, repr=False) + def __post_init__(self) -> None: """Call Future.__init__ for the EventFuture.""" super(EventFuture, self).__init__(loop=self.loop) diff --git a/reflex/state.py b/reflex/state.py index e7c96018654..5c3c9aa6607 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -32,6 +32,7 @@ from reflex_base.event import ( BACKGROUND_TASK_MARKER, EVENT_ACTIONS_MARKER, + SUPERSEDES_MARKER, Event, EventHandler, EventSpec, @@ -737,8 +738,9 @@ def _copy_fn(fn: Callable) -> Callable: closure=fn.__closure__, ) newfn.__annotations__ = fn.__annotations__ - if mark := getattr(fn, BACKGROUND_TASK_MARKER, None): - setattr(newfn, BACKGROUND_TASK_MARKER, mark) + for marker in (BACKGROUND_TASK_MARKER, SUPERSEDES_MARKER): + if mark := getattr(fn, marker, None): + setattr(newfn, marker, mark) # Preserve event_actions from @rx.event decorator if event_actions := getattr(fn, EVENT_ACTIONS_MARKER, None): object.__setattr__(newfn, EVENT_ACTIONS_MARKER, event_actions) @@ -2464,6 +2466,15 @@ def on_load_internal(self) -> list[Event | EventSpec | event.EventCallback] | No ] +# A newer navigation supersedes the previous unfinished on_load chain for the +# same client token, cancelling its stale work (#6593). +setattr( + OnLoadInternalState.event_handlers["on_load_internal"].fn, + SUPERSEDES_MARKER, + True, +) + + class ComponentState(State, mixin=True): """Base class to allow for the creation of a state instance per component. diff --git a/tests/units/reflex_base/event/processor/test_event_processor.py b/tests/units/reflex_base/event/processor/test_event_processor.py index bcc1108be98..2eb19a39c71 100644 --- a/tests/units/reflex_base/event/processor/test_event_processor.py +++ b/tests/units/reflex_base/event/processor/test_event_processor.py @@ -126,6 +126,72 @@ async def _background_slow_logging_handler(value: str = "default"): _background_slow_logging_handler._reflex_background_task = True # type: ignore[attr-defined] +# Gates for coordinating supersede tests; tests create loop-local events here. +_GATES: dict[str, asyncio.Event] = {} + + +async def _gated_logging_handler(value: str = "default"): + """Wait for the gate named ``value`` (if any), then log. + + Args: + value: The value to log; also names the gate to wait for. + """ + gate = _GATES.get(value) + if gate is not None: + await gate.wait() + _CALL_LOG.append({"value": value}) + + +async def _cancellable_load_handler(value: str = "default"): + """Log ``value``; if a gate named ``value`` exists, signal it and block. + + Args: + value: The value to log; also names the gate to signal. + """ + gate = _GATES.get(value) + if gate is not None: + gate.set() + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + _CALL_LOG.append({"value": f"{value}_cancelled"}) + raise + _CALL_LOG.append({"value": value}) + + +async def _resurrecting_load_handler(value: str = "default"): + """Signal the gate named ``value``, block, and chain an event when cancelled. + + Args: + value: The value naming the gate to signal. + """ + gate = _GATES.get(value) + if gate is not None: + gate.set() + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + ctx = EventContext.get() + await ctx.enqueue(Event.from_event_type(logging_event("resurrected"))[0]) + raise + + +async def _superseding_root_handler(value: str = "default", child: str = "load"): + """A superseding handler that chains a load event, like on_load_internal. + + Args: + value: Label forwarded to the chained load event. + child: Which child handler to chain ("load" or "resurrect"). + """ + ctx = EventContext.get() + child_event = ( + resurrecting_load_event if child == "resurrect" else cancellable_load_event + ) + await ctx.enqueue(Event.from_event_type(child_event(value))[0]) + + +_superseding_root_handler._reflex_supersedes = True # type: ignore[attr-defined] + noop_event = EventHandler(fn=_noop_handler) slow_event = EventHandler(fn=_slow_handler) @@ -140,6 +206,10 @@ async def _background_slow_logging_handler(value: str = "default"): background_slow_logging_event = EventHandler(fn=_background_slow_logging_handler) background_then_normal_event = EventHandler(fn=_background_then_normal_handler) error_then_logging_event = EventHandler(fn=_error_then_logging_handler) +gated_logging_event = EventHandler(fn=_gated_logging_handler) +cancellable_load_event = EventHandler(fn=_cancellable_load_handler) +resurrecting_load_event = EventHandler(fn=_resurrecting_load_handler) +superseding_root_event = EventHandler(fn=_superseding_root_handler) @pytest.fixture(autouse=True) @@ -150,6 +220,7 @@ def _register_handlers(forked_registration_context: RegistrationContext): forked_registration_context: Isolated registration context for the test. """ _CALL_LOG.clear() + _GATES.clear() for handler in ( noop_event, slow_event, @@ -164,6 +235,10 @@ def _register_handlers(forked_registration_context: RegistrationContext): background_slow_logging_event, background_then_normal_event, error_then_logging_event, + gated_logging_event, + cancellable_load_event, + resurrecting_load_event, + superseding_root_event, ): RegistrationContext.register_event_handler(handler) @@ -735,3 +810,108 @@ async def _watcher(): # noqa: RUF029 collected = [v async for v in _stream_queue_until_done(queue, _watcher())] assert collected == [99] + + +async def _drain_superseded(ep: EventProcessor) -> None: + """Give done callbacks a few ticks to clean the supersession tracking. + + Args: + ep: The event processor to wait on. + """ + for _ in range(20): + if not ep._superseded: + return + await asyncio.sleep(0) + + +async def test_superseding_event_cancels_previous_chain( + processor: EventProcessor, + token: str, +): + """A newer superseding event cancels the previous running chain (#6593). + + Args: + processor: The event processor fixture. + token: The client token. + """ + _GATES["stale"] = asyncio.Event() + processor.configure() + async with processor as ep: + stale = await ep.enqueue( + token, Event.from_event_type(superseding_root_event("stale"))[0] + ) + await asyncio.wait_for(_GATES["stale"].wait(), timeout=1) + # The root handler returned after chaining, but the chain is live. + assert stale.done() + assert not stale.all_done() + + current = await ep.enqueue( + token, Event.from_event_type(superseding_root_event("fresh"))[0] + ) + await asyncio.wait_for(current.wait_all(), timeout=1) + await _drain_superseded(ep) + assert ep._superseded == {} + + assert _CALL_LOG == [{"value": "stale_cancelled"}, {"value": "fresh"}] + + +async def test_superseding_event_skips_queued_stale_chain( + processor: EventProcessor, + token: str, +): + """A superseded chain that never started is skipped, not executed. + + Args: + processor: The event processor fixture. + token: The client token. + """ + _GATES["blocker"] = asyncio.Event() + processor.configure() + async with processor as ep: + await ep.enqueue( + token, Event.from_event_type(gated_logging_event("blocker"))[0] + ) + stale = await ep.enqueue( + token, Event.from_event_type(superseding_root_event("stale"))[0] + ) + # Let the stale entry reach the per-token queue before superseding it. + await ep.join(timeout=1) + current = await ep.enqueue( + token, Event.from_event_type(superseding_root_event("fresh"))[0] + ) + assert stale.cancelled() + _GATES["blocker"].set() + await asyncio.wait_for(current.wait_all(), timeout=1) + + # The stale root handler never ran at all. + assert _CALL_LOG == [{"value": "blocker"}, {"value": "fresh"}] + + +async def test_superseded_chain_cannot_chain_new_events( + processor: EventProcessor, + token: str, +): + """A cancelled chain cannot resurrect itself by chaining during unwind. + + Args: + processor: The event processor fixture. + token: The client token. + """ + _GATES["stale"] = asyncio.Event() + processor.configure() + async with processor as ep: + await ep.enqueue( + token, + Event.from_event_type(superseding_root_event("stale", "resurrect"))[0], + ) + await asyncio.wait_for(_GATES["stale"].wait(), timeout=1) + + current = await ep.enqueue( + token, Event.from_event_type(superseding_root_event("fresh"))[0] + ) + await asyncio.wait_for(current.wait_all(), timeout=1) + # Drain anything the unwinding stale chain may have enqueued. + await asyncio.wait_for(ep.join(), timeout=1) + + assert {"value": "resurrected"} not in _CALL_LOG + assert {"value": "fresh"} in _CALL_LOG diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 83aa823598a..e3cedbb8ac1 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -4985,3 +4985,87 @@ def child_view(self) -> int: parent_deps = ParentDescState._var_dependencies.get("_shared", set()) assert (ChildDescState.get_full_name(), "child_view") in child_deps assert (ParentDescState.get_full_name(), "parent_view") in parent_deps + + +class OnLoadCancelState(State): + """A test state whose on_load handler blocks until cancelled.""" + + # Signalling gates, populated per-test with loop-local events. + _gates: ClassVar[dict[str, asyncio.Event]] = {} + + @rx.event + async def slow_handler(self): + """Signal start, then block; signal again if cancelled.""" + type(self)._gates["started"].set() + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + type(self)._gates["cancelled"].set() + raise + + +async def test_on_load_internal_supersedes_previous_navigation( + app_module_mock, + token, + mock_root_event_context: EventContext, + mock_base_state_event_processor: BaseStateEventProcessor, +): + """A newer navigation cancels the previous unfinished on_load chain (#6593). + + Args: + app_module_mock: The app module that will be returned by get_app(). + token: A token. + mock_root_event_context: The mock root event context. + mock_base_state_event_processor: The event processor. + """ + OnLoadInternalState._app_ref = None + assert OnLoadInternalState.event_handlers["on_load_internal"].supersedes + assert not State.event_handlers["hydrate"].supersedes + + app = app_module_mock.app = App(_state=State) + app._state_manager = mock_root_event_context.state_manager + + def index(): + return "hello" + + app.add_page(index, on_load=OnLoadCancelState.slow_handler) + app._compile_page("index") + + OnLoadCancelState._gates = { + "started": asyncio.Event(), + "cancelled": asyncio.Event(), + } + on_load_internal_name = format.format_event_handler( + OnLoadInternalState.on_load_internal # pyright: ignore[reportArgumentType] + ) + + async with mock_base_state_event_processor as processor: + stale = await processor.enqueue( + token, + Event( + name=on_load_internal_name, + router_data={ + RouteVar.PATH: "/", + RouteVar.ORIGIN: "/", + RouteVar.QUERY: {}, + }, + ), + ) + await asyncio.wait_for(OnLoadCancelState._gates["started"].wait(), timeout=5) + + # Navigate to a page without on_load events (fast path). + current = await processor.enqueue( + token, + Event( + name=on_load_internal_name, + router_data={ + RouteVar.PATH: "/other", + RouteVar.ORIGIN: "/other", + RouteVar.QUERY: {}, + }, + ), + ) + await asyncio.wait_for(OnLoadCancelState._gates["cancelled"].wait(), timeout=5) + # The fresh navigation completes without waiting behind the stale chain. + await asyncio.wait_for(current.wait_all(), timeout=5) + assert stale.done()