diff --git a/news/+rfc-resolution-ledger.misc.md b/news/+rfc-resolution-ledger.misc.md new file mode 100644 index 00000000000..2b437c3476d --- /dev/null +++ b/news/+rfc-resolution-ledger.misc.md @@ -0,0 +1 @@ +RFC: delta flushes clean selectively via a resolution-dirt ledger, so a write landing during delta resolution survives for the next harvest instead of being discarded by a clean that never snapshotted it. diff --git a/packages/reflex-base/news/+rfc-resolution-ledger.misc.md b/packages/reflex-base/news/+rfc-resolution-ledger.misc.md new file mode 100644 index 00000000000..325a007d461 --- /dev/null +++ b/packages/reflex-base/news/+rfc-resolution-ledger.misc.md @@ -0,0 +1 @@ +RFC: chain_updates cleans selectively (dirty snapshot plus resolution ledger), keeping the SharedState fan-out capture complete while concurrent writers' dirt survives. 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 aec3de4a1d3..ffaff8d8ac4 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 @@ -209,22 +209,50 @@ async def chain_updates( root_state: The root state of the app, no delta emitted if omitted. """ from reflex.event import Event + from reflex.state import _recording_resolution_dirt, _resolve_delta ctx = EventContext.get() if root_state is not None: # 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: - await ctx.emit_delta(delta) - finally: - root_state._clean() + # latest state. The clean clears exactly what this flush published: + # the dirty snapshot plus whatever resolution itself dirtied (the + # SharedState patch machinery reached through computed vars records + # into the ledger), so the fan-out capture stays complete while a + # concurrent writer's dirt survives for the next harvest instead of + # being discarded by a clean that never snapshotted it. + delta = root_state.get_delta() + if not delta or not any( + inspect.iscoroutine(value) + for subdelta in delta.values() + for value in subdelta.values() + ): + # A coroutine-free resolution never yields the loop, so nothing + # can interleave with this flush before the emit; the selective + # machinery below would record and subtract nothing. Take the + # plain path the sync pipeline has always had. + try: + if delta: + await ctx.emit_delta(delta) + finally: + root_state._clean() + else: + # The flush suspends while resolving, which is the window the + # selective clean exists for: snapshot what this flush publishes, + # record what resolution itself dirties, and clean exactly that + # union, so a concurrent writer's dirt survives for the next + # harvest instead of being discarded by a clean that never + # snapshotted it. + flushed = root_state._snapshot_dirty_vars() + try: + with _recording_resolution_dirt(flushed_by_resolution := {}): + delta = await _resolve_delta(delta) + for state_name, var_names in flushed_by_resolution.items(): + flushed.setdefault(state_name, set()).update(var_names) + if delta: + await ctx.emit_delta(delta) + finally: + root_state._clean_flushed(flushed) # Convert valid EventHandler and EventSpec into Event if fixed_events := Event.from_event_type( diff --git a/reflex/istate/proxy.py b/reflex/istate/proxy.py index 919827949a2..88097df2534 100644 --- a/reflex/istate/proxy.py +++ b/reflex/istate/proxy.py @@ -614,7 +614,7 @@ def _mark_dirty( Returns: The result of the wrapped function. """ - self._self_state.dirty_vars.add(self._self_field_name) + self._self_state._record_dirty_var(self._self_field_name) self._self_state._mark_dirty() if wrapped is not None: return wrapped(*args, **(kwargs or {})) diff --git a/reflex/istate/shared.py b/reflex/istate/shared.py index e38517fef66..99ff704c606 100644 --- a/reflex/istate/shared.py +++ b/reflex/istate/shared.py @@ -149,6 +149,27 @@ def _clean(self): previous_dirty_vars.update(self.dirty_vars) super()._clean() + @_override_base_method + def _clean_flushed(self, flushed: dict[str, set[str]]) -> None: + """Selective clean that still captures this flush's fan-out seed. + + The capture mirrors ``_clean``: what this event published for this + state, which under a selective clean is exactly its slice of + ``flushed``, including resolution-created dirt from the ledger. A + concurrent writer's dirt is neither captured nor cleared. + + Args: + flushed: Mapping of full state name to the var names the flush + snapshotted or created during resolution. + """ + mine = self.dirty_vars & flushed.get(self.get_full_name(), set()) + if ( + previous_dirty_vars := getattr(self, "_previous_dirty_vars", None) + ) is not None and mine: + previous_dirty_vars.clear() + previous_dirty_vars.update(mine) + super()._clean_flushed(flushed) + @_override_base_method def _mark_dirty(self): """Override BaseState._mark_dirty to avoid marking certain vars as dirty. diff --git a/reflex/state.py b/reflex/state.py index a24b6f376d3..9a89e1b1a54 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -5,6 +5,7 @@ import asyncio import builtins import contextlib +import contextvars import copy import dataclasses import functools @@ -279,6 +280,43 @@ def get_var_for_field(cls: type[BaseState], name: str, f: Field) -> Var: _DROP_FROM_DELTA: Final = object() +_RESOLUTION_DIRT: contextvars.ContextVar[dict[str, set[str]] | None] = ( + contextvars.ContextVar("reflex_resolution_dirt", default=None) +) + +# Count of open recording contexts, so the per-dirty-mark fast path is a plain +# global check instead of a contextvar lookup while no flush is resolving. +_RESOLUTION_RECORDING_DEPTH = 0 + + +@contextlib.contextmanager +def _recording_resolution_dirt(ledger: dict[str, set[str]]) -> Iterator[None]: + """Record vars dirtied by the current task tree into the ledger. + + Delta resolution is not a pure read: computed vars reach ``get_state``, + which drives machinery (e.g. the SharedState patch) that marks vars + dirty. Contextvars are task-local and inherited by tasks created during + resolution, so writes made by the resolving task tree land in the ledger + while concurrent events' writes, made from tasks created elsewhere, do + not. That distinction is what lets a flush clean exactly what it + published and preserve everything else. + + Args: + ledger: Mapping of full state name to var names, filled in place. + + Yields: + None. + """ + global _RESOLUTION_RECORDING_DEPTH + token = _RESOLUTION_DIRT.set(ledger) + _RESOLUTION_RECORDING_DEPTH += 1 + try: + yield + finally: + _RESOLUTION_RECORDING_DEPTH -= 1 + _RESOLUTION_DIRT.reset(token) + + async def _resolve_delta(delta: Delta) -> Delta: """Await all coroutines in the delta, dropping keys that resolve to the drop sentinel. @@ -1519,7 +1557,7 @@ def __setattr__(self, name: str, value: Any): if name in self.backend_vars: self._backend_vars.__setitem__(name, value) - self.dirty_vars.add(name) + self._record_dirty_var(name) self._mark_dirty() return @@ -1552,12 +1590,12 @@ def __setattr__(self, name: str, value: Any): # Add the var to the dirty list. if name in self.base_vars: - self.dirty_vars.add(name) + self._record_dirty_var(name) self._mark_dirty() # For now, handle router_data updates as a special case if name == constants.ROUTER_DATA: - self.dirty_vars.add(name) + self._record_dirty_var(name) self._mark_dirty() def reset(self): @@ -1819,7 +1857,7 @@ def _mark_dirty_computed_vars(self) -> None: defining_state = self._get_root_state().get_substate( tuple(state_name.split(".")) ) - defining_state.dirty_vars.add(cvar) + defining_state._record_dirty_var(cvar) actual_var = defining_state.computed_vars.get(cvar) if actual_var is not None: actual_var.mark_dirty(instance=defining_state) @@ -1888,10 +1926,16 @@ def get_delta(self) -> Delta: if len(subdelta) > 0: delta[self.get_full_name()] = subdelta - # Recursively find the substate deltas. + # Recursively find the substate deltas. A dirty name without an + # attached substate is legitimate under selective cleaning: a prior + # flush preserved another writer's dirt, the tree was then fetched + # partially (redis fetches the handler's slice), and the dirt lives in + # the unfetched substate's own record until an event fetches it. It + # cannot contribute to this delta, so skip it, as _clean always has. substates = self.substates for substate in self.dirty_substates.union(self._always_dirty_substates): - delta.update(substates[substate].get_delta()) + if (substate_instance := substates.get(substate)) is not None: + delta.update(substate_instance.get_delta()) # Return the delta. return delta @@ -1904,6 +1948,79 @@ async def _get_resolved_delta(self) -> Delta: """ return await _resolve_delta(self.get_delta()) + def _record_dirty_var(self, name: str) -> None: + """Mark one var dirty, recording it when delta resolution is active. + + The single choke point for "a var became dirty", so a flush can tell + its own resolution-created dirt apart from concurrent writers. Known + boundary: dirt is name-only, so a concurrent rewrite of a var already + in a flush's snapshot, landing during that flush's resolution, is + still cleaned with it; distinguishing that needs per-var write + versions. Unreachable while every flush holds the token lock, which + all current callers do. + + Args: + name: The var name to mark dirty on this state. + """ + self.dirty_vars.add(name) + if _RESOLUTION_RECORDING_DEPTH: + self._note_resolution_dirt(name) + + def _note_resolution_dirt(self, name: str) -> None: + """Record one var into the active flush's resolution ledger. + + The cold half of ``_record_dirty_var``: reached only while a flush is + resolving, so the hot path pays one set-add and one falsy global + check. + + Args: + name: The var name that became dirty during resolution. + """ + if (ledger := _RESOLUTION_DIRT.get()) is not None: + ledger.setdefault(self.get_full_name(), set()).add(name) + + def _snapshot_dirty_vars(self) -> dict[str, set[str]]: + """The dirty var names per state, for the flush that just snapshotted. + + Walked over the same states ``get_delta`` visits, and including + backend vars, which never reach a delta but are captured for the + SharedState fan-out and must be cleaned once flushed. + + Returns: + Mapping of full state name to a copy of its dirty var names. + """ + snapshot: dict[str, set[str]] = {} + if self.dirty_vars: + snapshot[self.get_full_name()] = set(self.dirty_vars) + for substate_name in self.dirty_substates.union(self._always_dirty_substates): + substate = self.substates.get(substate_name) + if substate is not None: + snapshot.update(substate._snapshot_dirty_vars()) + return snapshot + + def _clean_flushed(self, flushed: dict[str, set[str]]) -> None: + """Clear exactly the dirty vars a flush published, keeping the rest. + + Unlike ``_clean``, a var dirtied by a concurrent writer after the + flush's snapshot survives for the next harvest instead of being + discarded by a clean that never saw it. + + Args: + flushed: Mapping of full state name to the var names the flush + snapshotted or created during resolution. + """ + self._update_was_touched() + for substate_name in tuple( + self.dirty_substates.union(self._always_dirty_substates) + ): + substate = self.substates.get(substate_name) + if substate is None: + continue + substate._clean_flushed(flushed) + if not substate.dirty_vars and not substate.dirty_substates: + self.dirty_substates.discard(substate_name) + self.dirty_vars -= flushed.get(self.get_full_name(), set()) + def _mark_dirty(self): """Mark the substate and all parent states as dirty.""" state_name = self.get_name() 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 c8d0476edb3..bda5fce885a 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 @@ -615,3 +615,72 @@ def raise_on_modify(*args, **kwargs): object.__setattr__(root_ctx.state_manager, "modify_state_with_links", original) assert proxy._self_entered_context is False + + +async def test_chain_updates_keeps_writes_made_during_delta_resolution( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """chain_updates must not clean a write it never snapshotted. + + The snapshot and the clean happen in one step, before the resolved delta + is awaited or emitted, so a concurrent write landing during resolution + stays dirty for the next harvest even if a caller ever runs chain_updates + without holding the state lock. + + 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. + """ + from reflex_base.event.processor.base_state_processor import chain_updates + + resolving = asyncio.Event() + release = asyncio.Event() + hold_resolution = [False] + + class MidResolveState(State): + victim: str = "" + + @rx.var(cache=False) + async def window(self) -> int: + if hold_resolution[0]: + resolving.set() + await release.wait() + return 0 + + root_ctx = real_base_state_processor._root_context + assert root_ctx is not None + EventContext.set(root_ctx.fork(token=token)) + state_manager = root_ctx.state_manager + + try: + root = await state_manager.get_state(BaseStateToken(ident=token, cls=State)) + substate = await root.get_state(MidResolveState) + root._clean() + + hold_resolution[0] = True + flush = asyncio.ensure_future( + chain_updates(None, root_state=root, handler_name="unlocked_flush") + ) + await resolving.wait() + substate.victim = "written" + hold_resolution[0] = False + release.set() + await flush + + assert "victim" in substate.dirty_vars, ( + "the write made during delta resolution was cleaned away" + ) + await chain_updates(None, root_state=root, handler_name="second_flush") + finally: + State._always_dirty_substates.discard(MidResolveState.get_name()) + + state_name = MidResolveState.get_full_name() + victim_key = "victim" + FIELD_MARKER + assert any( + d.get(state_name, {}).get(victim_key) == "written" for _, d in emitted_deltas + ), f"the surviving write never reached a delta: {emitted_deltas}" diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 416e09ade65..401ce34e78d 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -5188,3 +5188,18 @@ async def _coro(value): # noqa: RUF029 - a trivial coroutine value for the delt } resolved = await _resolve_delta(delta) assert resolved == {"s2": {"keep": 1}} + + +def test_get_delta_skips_dirty_names_without_attached_substates() -> None: + """A dirty substate name with no attached instance is skipped, not a KeyError. + + Under selective cleaning a prior flush can preserve another writer's dirt, + and a later redis fetch of the handler's slice restores the parent's + dirty_substates naming a substate that was not fetched. The dirt lives in + the unfetched substate's own record; the partial tree's delta walk must + pass over it the way _clean always has. + """ + root = State(_reflex_internal_init=True) # pyright: ignore [reportCallIssue] + root._clean() + root.dirty_substates.add("not___an___attached____substate") + assert root.get_delta() == {}