diff --git a/news/6920.bugfix.md b/news/6920.bugfix.md new file mode 100644 index 00000000000..03088174338 --- /dev/null +++ b/news/6920.bugfix.md @@ -0,0 +1 @@ +Fixed a race where a finishing background task could silently discard state updates made by a concurrently running event handler before they reached the frontend, leaving the UI stale until the next write. Background handlers that never enter `async with self` still emit their delta, now computed under the state lock. diff --git a/packages/reflex-base/news/6920.bugfix.md b/packages/reflex-base/news/6920.bugfix.md new file mode 100644 index 00000000000..c446b5d6b35 --- /dev/null +++ b/packages/reflex-base/news/6920.bugfix.md @@ -0,0 +1 @@ +Stop background event handlers from computing a delta and cleaning the root state after the state lock is dropped. On a shared state tree (opportunistic locking, in-memory state manager) that unlocked snapshot-then-clean raced concurrent foreground handlers: a write landing between the background task's dirty-var snapshot and its `_clean()` was silently discarded and never reached any delta. Background state changes are emitted by their `async with self` context exits, which hold the lock. 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 d1b3a53744f..aec3de4a1d3 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 @@ -213,7 +213,12 @@ async def chain_updates( ctx = EventContext.get() if root_state is not None: - # Emit deltas first, so any frontend events are processed with the latest state. + # Emit deltas first, so any frontend events are processed with the + # latest state. The clean deliberately runs after resolution: the + # SharedState fan-out captures its dirty vars at clean time, and + # resolving the delta is what re-marks linked vars through the patch + # machinery, so cleaning earlier would fan out a stale set (see + # tests/integration/test_linked_state.py). try: delta = await root_state._get_resolved_delta() if delta: @@ -228,11 +233,38 @@ async def chain_updates( await _route_events(ctx, fixed_events) +def ensure_locked( + state: BaseState | StateProxy, root_state: BaseState | None +) -> BaseState | None: + """The root to flush deltas from, only while the state lock is held. + + Foreground handlers pass the locked root in. A background handler + yielding from inside ``async with self`` is suspended while its proxy + still holds the lock, so flushing through the proxy's root keeps the + documented ordering: deltas reach the frontend before the yielded event. + Outside the proxy context there is no lock, and flushing there is the + unlocked snapshot/clean that discards concurrent writes. Evaluated per + yield: a generator can move between inside and outside the context. + + Args: + state: The state the handler runs against, possibly a StateProxy. + root_state: The locked root passed by foreground callers, if any. + + Returns: + The root state to flush, or None when the lock is not held. + """ + if root_state is not None: + return root_state + if isinstance(state, StateProxy) and state._is_mutable(): + return state.__wrapped__._get_root_state() + return None + + async def process_event( handler: EventHandler, payload: dict, state: BaseState | StateProxy, - root_state: BaseState, + root_state: BaseState | None, ): """Process event. @@ -240,7 +272,12 @@ async def process_event( handler: EventHandler to process. payload: The event payload. state: State to process the handler. - root_state: The root state of the app, used for emitting deltas. + root_state: The root state of the app, used for emitting deltas. Pass + None when the caller does not hold the state lock (background + tasks): computing and cleaning a delta on an unlocked root races + concurrent events on a shared state tree, and background state + changes are emitted by the ``async with self`` context exits + instead. Raises: ValueError: If a string value is received for an int or float type and cannot be converted. @@ -269,28 +306,44 @@ async def process_event( # Handle async generators. if inspect.isasyncgen(events): async for event in events: - await chain_updates(event, root_state=root_state, handler_name=handler_name) - await chain_updates(None, root_state=root_state, handler_name=handler_name) + await chain_updates( + event, + root_state=ensure_locked(state, root_state), + handler_name=handler_name, + ) + await chain_updates( + None, root_state=ensure_locked(state, root_state), handler_name=handler_name + ) # Handle regular generators. elif inspect.isgenerator(events): try: while True: await chain_updates( - next(events), root_state=root_state, handler_name=handler_name + next(events), + root_state=ensure_locked(state, root_state), + handler_name=handler_name, ) except StopIteration as si: # the "return" value of the generator is not available # in the loop, we must catch StopIteration to access it if si.value is not None: await chain_updates( - si.value, root_state=root_state, handler_name=handler_name + si.value, + root_state=ensure_locked(state, root_state), + handler_name=handler_name, ) - await chain_updates(None, root_state=root_state, handler_name=handler_name) + await chain_updates( + None, root_state=ensure_locked(state, root_state), handler_name=handler_name + ) # Handle regular event chains. else: - await chain_updates(events, root_state=root_state, handler_name=handler_name) + await chain_updates( + events, + root_state=ensure_locked(state, root_state), + handler_name=handler_name, + ) class BaseStateEventProcessor(EventProcessor): @@ -395,13 +448,38 @@ async def _execute_event( root_state=root_state, ) return - # Otherwise drop the state lock and start processing the background task with a proxy state. + # Otherwise drop the state lock and start processing the background task + # with a proxy state. No root_state: the lock is no longer held, and + # under a shared state tree (opportunistic locking, in-memory manager) + # computing a delta here races whatever event holds the lock now -- a + # foreground write landing between this task's dirty-var snapshot and + # its _clean() would be discarded before any delta carries it. A + # background task's own state changes are emitted (and cleaned) by its + # `async with self` context exits, which re-acquire the lock. + proxy = StateProxy(substate) await process_event( handler=registered_handler.handler, - state=StateProxy(substate), + state=proxy, payload=event.payload, - root_state=root_state, + root_state=None, ) + if not proxy._self_entered_context: + # A handler that never entered `async with self` emitted nothing, + # but every background event used to flush a delta (refreshing + # uncached computed vars, and any dirty vars the preamble left, + # like router_data). Preserve that, under the lock this time. + async with ctx.state_manager.modify_state_with_links( + BaseStateToken( + ident=ctx.token, + cls=registered_handler.states[0], + ), + event=event, + ) as flush_state: + await chain_updates( + None, + root_state=flush_state._get_root_state(), + handler_name=registered_handler.handler.fn.__qualname__, + ) async def _handle_backend_exception( self, ex: Exception, ev_ctx: EventContext | None = None diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index ce2eef96834..919827949a2 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -96,6 +96,9 @@ def __init__( self._self_actx_lock = asyncio.Lock() self._self_actx_lock_holder = None self._self_parent_state_proxy = parent_state_proxy + # Whether `async with self` was ever entered; a background handler that + # never did emitted no delta, so the processor flushes once for it. + self._self_entered_context = False def _is_mutable(self) -> bool: """Check if the state is mutable. @@ -134,6 +137,7 @@ async def __aenter__(self) -> Self: State.get_class_substate(self._self_substate_path) ), ) + self._self_entered_context = True return self current_task = asyncio.current_task() if ( @@ -153,6 +157,7 @@ async def __aenter__(self) -> Self: ) mutable_state = await self._self_actx.__aenter__() self._self_mutable = True + self._self_entered_context = True super().__setattr__( "__wrapped__", mutable_state.get_substate(self._self_substate_path) ) 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 ca0a0874aaa..c8d0476edb3 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 @@ -219,6 +219,254 @@ async def preprocess(self, app, state, event) -> StateUpdate: assert "_redirect" in client_event_names +async def test_background_event_does_not_discard_concurrent_foreground_write( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """A foreground write racing a background task's completion reaches a delta. + + Regression: after dropping the state lock, the background branch passed + ``root_state`` into ``process_event``, whose trailing ``chain_updates`` + snapshotted dirty vars, suspended (delta resolution / emit), and then ran + ``_clean()`` -- all unlocked. On a shared state tree (opportunistic + locking, or the in-memory manager used here) a foreground handler's write + landing inside that snapshot->clean window was cleaned before any delta + harvested it: the value never reached the frontend. + + Every step is gate-driven, in both worlds: pre-fix, the background + task's trailing delta resolution (via the uncached async computed var) + parks until the foreground handler has written, and the background + event's future completes strictly after its trailing ``_clean()``; on + fixed code that future completes without any trailing resolution and + the foreground handler proceeds immediately. + + Args: + wired_app: The App wired to the processor's state manager. + real_base_state_processor: The unmocked BaseStateEventProcessor. + emitted_deltas: List to capture emitted deltas. + token: The client token. + """ + bg_started = asyncio.Event() + bg_resolving = asyncio.Event() + fg_wrote = asyncio.Event() + hold_resolution = [False] + bg_future_box: list = [] + + class BgRaceState(State): + victim: str = "" + + @rx.var(cache=False) + async def window(self) -> int: + # Uncached: recomputed in every delta. Once armed, the background + # task's trailing delta resolution parks here until the foreground + # handler has written -- holding the snapshot->clean window open. + if hold_resolution[0]: + bg_resolving.set() + await fg_wrote.wait() + return 0 + + @event(background=True) + async def bg(self): + # Enter the context once so the never-entered compatibility flush + # (which runs under the lock) stays out of this choreography; arm + # the gate only afterwards so the context exit resolves unparked. + async with self: + pass + hold_resolution[0] = True + bg_started.set() + + @event + async def fg(self): + # Proceed once the background trailing resolution has taken its + # snapshot (pre-fix code parks it on fg_wrote), or once the + # background event has fully completed without one (fixed code). + waiter = asyncio.ensure_future(bg_resolving.wait()) + await asyncio.wait( + [waiter, *bg_future_box], return_when=asyncio.FIRST_COMPLETED + ) + waiter.cancel() + self.victim = "written" + fg_wrote.set() + # Pre-fix, the parked resolution now resumes, emits, and cleans; + # the background event's future completes strictly after that + # clean, so awaiting it guarantees the clean landed before this + # handler's own delta snapshot. On fixed code it is already done. + await asyncio.wait(bg_future_box) + + # Seed router_data up front so no event triggers _rehydrate (whose + # full-dict resolution would park on the armed gate before the foreground + # handler could run) and no event needs to carry router_data of its own. + assert real_base_state_processor._root_context is not None + state_manager = real_base_state_processor._root_context.state_manager + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=State) + ) as seed_root: + seed_root.router_data = {"pathname": "/", "query": {}} + try: + async with real_base_state_processor as processor: + bg_future_box.append( + await processor.enqueue( + token, Event.from_event_type(BgRaceState.bg())[0] + ) + ) + started = asyncio.ensure_future(bg_started.wait()) + await asyncio.wait([started], timeout=2) + started.cancel() + assert bg_started.is_set(), "background handler never started" + await processor.enqueue(token, Event.from_event_type(BgRaceState.fg())[0]) + await processor.join(5) + finally: + # The uncached computed var registered BgRaceState as an always-dirty + # substate on the shared State class; later tests' state trees don't + # contain it and would KeyError in get_delta (see reload_state_module). + State._always_dirty_substates.discard(BgRaceState.get_name()) + + state_name = BgRaceState.get_full_name() + victim_key = "victim" + FIELD_MARKER + delivered = [ + d[state_name][victim_key] + for _, d in emitted_deltas + if state_name in d and d[state_name].get(victim_key) == "written" + ] + assert delivered, ( + "The foreground handler's write never reached any delta; it was " + "cleaned by the background task's unlocked trailing update. Deltas: " + f"{emitted_deltas}" + ) + + +async def test_background_yield_inside_context_flushes_delta_before_event( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, +): + """A yield inside ``async with self`` emits the delta before the event. + + A background generator suspended at a yield inside its proxy context + still holds the state lock, so flushing the delta there is safe and + preserves the documented ordering: frontend events are processed with + the latest state. Only yields outside the context (no lock) must skip + the flush. + + Args: + wired_app: The App wired to the processor's state manager. + real_base_state_processor: The unmocked BaseStateEventProcessor. + token: The client token. + """ + timeline: list[tuple[str, Any]] = [] + root_ctx = real_base_state_processor._root_context + assert root_ctx is not None + + async def record_delta(tok: str, delta: Mapping[str, Mapping[str, Any]]) -> None: # noqa: RUF029 + timeline.append(("delta", delta)) + + async def record_event(tok: str, *events: Event) -> None: # noqa: RUF029 + timeline.append(("event", tuple(ev.name for ev in events))) + + object.__setattr__(root_ctx, "emit_delta_impl", record_delta) + object.__setattr__(root_ctx, "emit_event_impl", record_event) + + class BgYieldOrderState(State): + marker: str = "" + + @event(background=True) + async def bg_yield(self): + async with self: + self.marker = "set" + yield rx.call_script("void 0") + + state_manager = root_ctx.state_manager + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=State) + ) as seed_root: + seed_root.router_data = {"pathname": "/", "query": {}} + + async with real_base_state_processor as processor: + await processor.enqueue( + token, Event.from_event_type(BgYieldOrderState.bg_yield())[0] + ) + await processor.join(5) + + state_name = BgYieldOrderState.get_full_name() + marker_key = "marker" + FIELD_MARKER + delta_index = next( + ( + index + for index, (kind, payload) in enumerate(timeline) + if kind == "delta" and payload.get(state_name, {}).get(marker_key) == "set" + ), + None, + ) + event_index = next( + ( + index + for index, (kind, payload) in enumerate(timeline) + if kind == "event" and "_call_script" in payload + ), + None, + ) + assert delta_index is not None, f"marker delta never emitted: {timeline}" + assert event_index is not None, f"yielded event never emitted: {timeline}" + assert delta_index < event_index, ( + "The yielded frontend event was emitted before the delta for the " + f"mutation made in the same proxy context: {timeline}" + ) + + +async def test_background_event_without_context_still_flushes_a_delta( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """A background handler with no ``async with self`` still flushes a delta. + + Backward compatibility: before the unlocked trailing flush was removed, + every background event emitted a delta, which is what refreshed uncached + computed vars for apps driving re-renders off a bare background tick. + That flush now runs under the state lock instead of disappearing. + + Args: + wired_app: The App wired to the processor's state manager. + real_base_state_processor: The unmocked BaseStateEventProcessor. + emitted_deltas: List to capture emitted deltas. + token: The client token. + """ + + class NoContextBgState(State): + @rx.var(cache=False) + def beat(self) -> int: + return 7 + + @event(background=True) + async def bg(self): + pass + + assert real_base_state_processor._root_context is not None + state_manager = real_base_state_processor._root_context.state_manager + async with state_manager.modify_state( + BaseStateToken(ident=token, cls=State) + ) as seed_root: + seed_root.router_data = {"pathname": "/", "query": {}} + + try: + async with real_base_state_processor as processor: + await processor.enqueue( + token, Event.from_event_type(NoContextBgState.bg())[0] + ) + await processor.join(5) + finally: + State._always_dirty_substates.discard(NoContextBgState.get_name()) + + state_name = NoContextBgState.get_full_name() + beat_key = "beat" + FIELD_MARKER + assert any(d.get(state_name, {}).get(beat_key) == 7 for _, d in emitted_deltas), ( + f"no delta refreshed the uncached var: {emitted_deltas}" + ) + + async def test_chained_event_keeps_originating_router_data( wired_app: App, real_base_state_processor: BaseStateEventProcessor, @@ -288,3 +536,82 @@ def client_event(spec, router_data: dict[str, Any]) -> Event: BaseStateToken(ident=token, cls=State) ) assert (await state.get_state(RouterState)).seen == ["/item/abc|abc"] + + +async def test_ensure_locked_returns_a_root_only_while_the_lock_is_held( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, +): + """ensure_locked passes a locked root through and refuses everything else. + + Foreground callers hand in the root they locked; a plain substate or an + un-entered proxy holds no lock, so there is nothing safe to flush. + + Args: + wired_app: The App wired to the processor's state manager. + real_base_state_processor: The unmocked BaseStateEventProcessor. + token: The client token. + """ + from reflex_base.event.processor.base_state_processor import ensure_locked + + from reflex.istate.proxy import StateProxy + + root_ctx = real_base_state_processor._root_context + assert root_ctx is not None + EventContext.set(root_ctx.fork(token=token)) + root = await root_ctx.state_manager.get_state( + BaseStateToken(ident=token, cls=State) + ) + substate = await root.get_state(OnLoadInternalState) + + assert ensure_locked(substate, root) is root + assert ensure_locked(substate, None) is None + assert ensure_locked(StateProxy(substate), None) is None + + +async def test_failed_context_enter_does_not_mark_the_proxy_entered( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + token: str, +): + """A proxy whose enter failed still gets the compatibility flush. + + The entered flag means "a context opened whose exit will flush". If + ``__aenter__`` raises before that and the handler swallows it, the + processor must still run the locked fallback flush, or preamble dirty + vars like router_data would never reach a delta. + + Args: + wired_app: The App wired to the processor's state manager. + real_base_state_processor: The unmocked BaseStateEventProcessor. + token: The client token. + """ + from reflex.istate.proxy import StateProxy + + root_ctx = real_base_state_processor._root_context + assert root_ctx is not None + EventContext.set(root_ctx.fork(token=token)) + root = await root_ctx.state_manager.get_state( + BaseStateToken(ident=token, cls=State) + ) + substate = await root.get_state(OnLoadInternalState) + + proxy = StateProxy(substate) + + def raise_on_modify(*args, **kwargs): + msg = "state manager unavailable" + raise RuntimeError(msg) + + original = root_ctx.state_manager.modify_state_with_links + object.__setattr__( + root_ctx.state_manager, "modify_state_with_links", raise_on_modify + ) + try: + with pytest.raises(RuntimeError, match="state manager unavailable"): + async with proxy: + pass + finally: + object.__setattr__(root_ctx.state_manager, "modify_state_with_links", original) + + assert proxy._self_entered_context is False