fix(events): stop background tasks from cleaning the root state unlocked

Merged
AD
2mo ago
+424-12

The bug

After dropping the state lock, the background branch of _execute_event passes root_state into process_event, whose trailing chain_updates does three things on that root with no lock held:

  1. get_delta() snapshots the dirty vars (sync)
  2. _get_resolved_delta() / emit_delta() suspend (async computed vars, socket write)
  3. finally: root_state._clean() wipes all dirty vars

On a shared state tree (opportunistic locking, or StateManagerMemory) that window races whatever event holds the lock now. A foreground handler's write landing between the snapshot and the clean is discarded before any delta harvests it. The value reaches no delta, ever: the handler completes normally, its delta arrives carrying only always-dirty computed vars, and the frontend keeps the stale value.

Writes made after the window survive, so a single handler can lose its early assignments and deliver its late ones, which makes the failure look impossible from the app side.

How it was found

A merge-queue flake in a downstream app: clicking a card 5-7 ms after a 1 Hz background poll silently did nothing, roughly 1 in 8 such clicks. The click's handler ran, its selected_integration assignment was wiped mid-handler by the poll's trailing clean, and the panel never opened. The same mechanism can wipe any state var, including auth state.

The fix

Pass root_state=None in the background branch. chain_updates already treats None as "route events only, no delta work", so the unlocked snapshot/clean simply stops existing. A background task's own state changes are emitted and cleaned by its async with self context exits, which re-acquire the lock, so nothing legitimate is lost. Yielded and returned events still route exactly as before.

Validation

  • New regression test forces the exact interleaving with gates (no timing dependence): fails in 0.14 s on unpatched code, passes with the fix.
  • A standalone repro (below) driving the real event processor against redis with REFLEX_OPLOCK_ENABLED goes from 13/100 lost foreground deltas to 0/100 with this fix. The oplock-off control loses 0/100 either way.
  • tests/units/test_state.py, tests/units/reflex_base/, tests/units/test_app.py produce a failure set identical to the pre-fix baseline on this checkout.
Standalone repro (redis on localhost:6399) 
import asyncio
import os
import random

os.environ["REFLEX_REDIS_URL"] = "redis://localhost:6399/9"
os.environ["REFLEX_OPLOCK_ENABLED"] = "true"
os.environ["REFLEX_OPLOCK_HOLD_TIME_MS"] = "1000"
os.environ["REFLEX_REDIS_LOCK_EXPIRATION"] = "20000"

import redis.asyncio as aioredis
import reflex as rx
from reflex.config import get_config

get_config()

from reflex.app import BaseStateEventProcessor
from reflex.istate.manager.redis import StateManagerRedis
from reflex.istate.manager.token import BaseStateToken
from reflex_base.event import Event


class Victim(rx.State):
    value: str = ""

    @rx.event
    async def set_value_with_await(self):
        self.value = "clicked"        # early write, the one that gets lost
        await asyncio.sleep(0.005)    # any await: db query, get_state, http


class Poller(rx.State):
    ticks: int = 0

    # Uncached async computed var: keeps every delta non-empty and suspending,
    # like any real app with async computed vars.
    @rx.var(cache=False)
    async def load(self) -> int:
        await asyncio.sleep(0.003)
        return self.ticks

    @rx.event(background=True)
    async def poll(self):
        await asyncio.sleep(random.uniform(0, 0.02))  # unlocked preamble work
        async with self:
            self.ticks += 1


VICTIM = Victim.get_full_name()
CLICK = f"{VICTIM}.set_value_with_await"
POLL = f"{Poller.get_full_name()}.poll"
RD = {"pathname": "/", "asPath": "/", "query": {}}

app = rx.App()


class CapturingNamespace:
    def __init__(self):
        self.updates = []

    async def emit_update(self, update, token):
        await asyncio.sleep(0.002)  # a real socket write suspends
        self.updates.append(update)


async def run_trial(sm):
    ns = CapturingNamespace()
    proc = BaseStateEventProcessor(middleware=app, backend_exception_handler=None)
    app._event_processor = proc
    losses = 0
    async with proc.configure(state_manager=sm, event_namespace=ns):
        with app.set_contexts():
            tok = "tab-1"
            async with sm.modify_state(BaseStateToken(ident=tok, cls=rx.State)) as root:
                root.router_data = dict(RD)
            for _ in range(100):
                ns.updates.clear()
                fut_poll = await proc.enqueue(
                    tok, Event(name=POLL, payload={}, router_data=dict(RD))
                )
                await asyncio.sleep(random.uniform(0, 0.03))
                fut_click = await proc.enqueue(
                    tok, Event(name=CLICK, payload={}, router_data=dict(RD))
                )
                await asyncio.wait([fut_poll, fut_click], timeout=10)
                await asyncio.sleep(0.08)
                delivered = any(
                    VICTIM in (u.delta or {}) and "value_rx_state_" in u.delta[VICTIM]
                    for u in ns.updates
                )
                if not delivered:
                    losses += 1
    return losses


async def main():
    redis = aioredis.Redis.from_url("redis://localhost:6399/9")
    await redis.flushdb()
    lost_on = await run_trial(StateManagerRedis(redis=redis))
    print(f"oplock ON : lost {lost_on}/100 foreground deltas")

    await redis.flushdb()
    sm_off = StateManagerRedis(redis=redis)
    object.__setattr__(sm_off, "_oplock_enabled", False)
    lost_off = await run_trial(sm_off)
    print(f"oplock OFF: lost {lost_off}/100 foreground deltas")
    await redis.aclose()


asyncio.run(main())

Notes for reviewers

  • Background delta flow after this PR: yields from inside async with self still flush ahead of the yielded event (the proxy holds the lock there); proxy context exits emit as documented; and a handler that never enters async with self gets one trailing flush under the state lock, preserving the old per-tick refresh of uncached computed vars. The only removed behavior is the unlocked flush itself.
  • An atomic snapshot-and-clean in chain_updates was attempted and reverted (it broke test_linked_state on every in-memory leg): delta resolution itself re-marks linked vars dirty through the SharedState patch machinery, and SharedState._clean captures its dirty vars for the cross-client fan-out at clean time, so the clean must stay after resolution. That ordering dependency is now documented at the clean site; decoupling the fan-out capture from _clean so the primitive can be hardened is left as follow-up.
  • Writing the regression test surfaced two pre-existing cross-test leaks in the suite, both worked around locally in the new test: an inline state with a cache=False var permanently mutates the shared State._always_dirty_substates (the test discards its entry in a finally, mirroring reload_state_module), and test_router_var_dep leaks a router edge into class-level State._var_dependencies that breaks any later test whose events assign router state (the new test pre-seeds router_data instead of carrying it on events). Happy to file these separately.
  • reflex_enterprise.utils.call_event_from_computed_var has the same defect class (a detached state._clean() outside any lock); that is a separate repo and a separate fix.