From 43b4f00f2b3edfd0181f800d3d7a19a04ef957cd Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 19:03:25 +0000 Subject: [PATCH 01/14] Cut redundant state loads and lock checks on first page load Verify the redis lock and read its TTL once per state-tree save instead of once per substate, flip is_hydrated from the OnLoadInternalState leaf so the final step of the on_load chain no longer fetches and persists every substate, and stop opening the whole state tree on websocket connect just to stamp the session id (the first event after connecting carries it in router_data). Measured on a 20-substate app with the redis state manager: redis commands per first load 291 -> 76, backend CPU 99ms -> 62ms. --- reflex/app.py | 12 +++--------- reflex/istate/manager/redis.py | 28 +++++++++++++++------------- reflex/istate/shared.py | 4 ++-- reflex/state.py | 15 ++++++++++++++- 4 files changed, 34 insertions(+), 25 deletions(-) diff --git a/reflex/app.py b/reflex/app.py index ff18d3bda0c..d94519104aa 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -73,7 +73,6 @@ from reflex.app_mixins import AppMixin, LifespanMixin, MiddlewareMixin from reflex.compiler import compiler from reflex.compiler.compiler import readable_name_from_component -from reflex.istate.data import RouterData from reflex.istate.manager import StateManager, StateModificationContext from reflex.istate.manager.token import BaseStateToken from reflex.route import ( @@ -2277,11 +2276,6 @@ async def link_token_to_sid(self, sid: str, token: str): if new_token: # Duplicate detected, emit new token to client await self.emit("new_token", new_token, to=sid) - - # Update client state to apply new sid/token for running background tasks. - if self.app._state is not None: - async with self.app.state_manager.modify_state( - BaseStateToken(ident=new_token or token, cls=self.app._state) - ) as state: - state.router_data[constants.RouteVar.SESSION_ID] = sid - state.router = RouterData.from_router_data(state.router_data) + # The new sid reaches the state through the router data of the first + # event the client sends after connecting (always the hydrate chain), + # so there is no need to load and persist the whole state tree here. diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 033998d1ef5..7101cb9c4e2 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -423,11 +423,7 @@ async def set_state( await self.redis.set(str(token), pickle_state, ex=self.token_expiration) return - base_state = cast(BaseState, state) - - lock_key = token.lock_key - - if lock_id is not None and lock_key not in self._local_leases: + if lock_id is not None and token.lock_key not in self._local_leases: time_taken = ( self.lock_expiration - (await self.redis.pttl(self._lock_key(token))) ) / 1000 @@ -444,16 +440,22 @@ async def set_state( extra={"dedupe": True}, ) - # Recursively set_state on all known substates. + await self._set_state_tree(token, cast(BaseState, state)) + + async def _set_state_tree(self, token: BaseStateToken, base_state: BaseState): + """Persist a state and, concurrently, every substate attached to it. + + The lock check and the hold-time warning happen once in ``set_state``; + this recursion only writes the keys that were touched. + + Args: + token: The token (any state class) identifying the client. + base_state: The state instance whose tree to persist. + """ tasks = [ asyncio.create_task( - self.set_state( - token, - substate, - lock_id=lock_id, - **context, - ), - name=f"reflex_set_state|{lock_key}|{substate.get_full_name()}", + self._set_state_tree(token, substate), + name=f"reflex_set_state|{token.lock_key}|{substate.get_full_name()}", ) for substate in base_state.substates.values() ] diff --git a/reflex/istate/shared.py b/reflex/istate/shared.py index 432cdadbb52..a271d118b84 100644 --- a/reflex/istate/shared.py +++ b/reflex/istate/shared.py @@ -13,7 +13,7 @@ from typing_extensions import Self from reflex.istate.manager.token import BaseStateToken -from reflex.state import BaseState, State, _override_base_method +from reflex.state import BaseState, OnLoadInternalState, State, _override_base_method logger = logging.getLogger(__name__) @@ -179,7 +179,7 @@ def _rehydrate(self): Event( name=get_hydrate_event(self._get_root_state()), ), - State.set_is_hydrated(True), + OnLoadInternalState.set_is_hydrated(True), ] async def _resolve_linked_state( diff --git a/reflex/state.py b/reflex/state.py index 1ea080cc863..ba118b867d6 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -2630,9 +2630,22 @@ def on_load_internal(self) -> list[Event | EventSpec | event.EventCallback] | No load_events, router_data=self.router_data, ), - State.set_is_hydrated(True), + OnLoadInternalState.set_is_hydrated(True), ] + @event + def set_is_hydrated(self, value: bool) -> None: + """Set the hydrated flag from this leaf substate. + + Targeting the leaf instead of ``State.set_is_hydrated`` keeps the + redis state manager from fetching and persisting every substate in the + app just to flip one root var. + + Args: + value: The hydrated state. + """ + self.is_hydrated = value + class ComponentState(State, mixin=True): """Base class to allow for the creation of a state instance per component. From c3e69307788c4c2f40d501c5f82871054e23010b Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 19:25:42 +0000 Subject: [PATCH 02/14] Hydrate in one lock cycle, from the connect packet, as a diff of the compiled defaults Replace the three events the frontend sent on every websocket (re)connect (hydrate, update_vars_internal, on_load_internal) with a single State.hydrate_and_load event that resets and applies client storage, sends the state snapshot and queues the page's on_load handlers under one state lock. The event rides in the socket.io CONNECT packet (auth), so the backend starts loading state as soon as the namespace connects instead of after the connect acknowledgement round trip. On the first hydrate of a page the frontend still holds the compiled initialState, so it sends per-state hashes of it; for every state whose compiled defaults match the backend's, only vars that differ are sent. Reconnects and re-hydrates after storage resets still get the full state. Measured on a 20-substate app (redis state manager, 40ms RTT, 4Mbps, no websocket compression as with granian): time to hydrated 498ms -> 631ms is 716ms without the diff; hydrate payload 36KB -> 2KB on the wire; websocket frames after connect 5 -> 3; redis commands per first load 76 -> 55. --- .../reflex_base/.templates/web/utils/state.js | 43 +++-- .../src/reflex_base/compiler/templates.py | 44 ++++- .../src/reflex_base/constants/compiler.py | 2 + .../event/processor/base_state_processor.py | 9 +- reflex/app.py | 8 +- reflex/compiler/compiler.py | 6 +- reflex/state.py | 180 ++++++++++++++++-- tests/units/istate/manager/test_redis.py | 46 +++++ tests/units/test_app.py | 29 +++ tests/units/test_state.py | 172 ++++++++++++++++- 10 files changed, 476 insertions(+), 63 deletions(-) diff --git a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js index 8ba6d00509c..d97b82675f6 100644 --- a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js +++ b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js @@ -395,7 +395,19 @@ export const applyEvent = async (event, socket, navigate, params) => { return; } - // Update token and router data (if missing). + // Send the event to the server. + if (socket) { + socket.emit("event", withRouterData(event, params)); + } +}; + +/** + * Fill in the event's router data from the current location, if missing. + * @param event The event to send. + * @param params The params object from useParams + * @returns The same event, with router_data populated. + */ +const withRouterData = (event, params) => { if ( event.router_data === undefined || Object.keys(event.router_data).length === 0 @@ -423,11 +435,7 @@ export const applyEvent = async (event, socket, navigate, params) => { event.router_data.query = query; } } - - // Send the event to the server. - if (socket) { - socket.emit("event", event); - } + return event; }; /** @@ -589,6 +597,13 @@ export const connect = async ( const endpoint = getBackendURL(EVENTURL); const on_hydrated_queue = []; + // The hydrate event rides in the socket.io CONNECT packet, so the backend + // starts loading state as soon as the namespace connects instead of after + // an extra round trip for the connect acknowledgement. + const bootAuth = (first) => ({ + event: withRouterData(initialEvents(first)[0], params), + }); + // Create the socket. socket.current = io(endpoint.href, { path: endpoint["pathname"], @@ -596,6 +611,7 @@ export const connect = async ( protocols: [reflexEnvironment.version], autoUnref: false, query: { token: getToken() }, + auth: bootAuth(true), reconnection: false, // Reconnection will be handled manually. }); socket.current.wait_connect = !socket.current.connected; @@ -623,8 +639,9 @@ export const connect = async ( !socket.current.wait_connect ) { socket.current.wait_connect = true; - socket.current.rehydrate = true; socket.current.io.opts.query = { token: getToken() }; // Update token for reconnect. + // A reconnect rehydrates in full: the reducers no longer hold the defaults. + socket.current.auth = bootAuth(false); socket.current.connect(); } }; @@ -675,10 +692,6 @@ export const connect = async ( setConnectErrors([]); window.addEventListener("pagehide", pagehideHandler); window.addEventListener("beforeunload", disconnectTrigger); - if (socket.current.rehydrate) { - socket.current.rehydrate = false; - queueEvents(initialEvents(), socket, true, navigate, params); - } // Drain any initial events from the queue. while (event_queue.length > 0) { await processEvent(socket.current, navigate, params); @@ -1062,14 +1075,6 @@ export const useEventLoop = ( ); }, []); - const sentHydrate = useRef(false); // Avoid double-hydrate due to React strict-mode - useEffect(() => { - if (!sentHydrate.current) { - queueEvents(initial_events(), socket, true, navigate, params); - sentHydrate.current = true; - } - }, []); - // Handle frontend errors and send them to the backend via websocket. useEffect(() => { if (typeof window === "undefined") { diff --git a/packages/reflex-base/src/reflex_base/compiler/templates.py b/packages/reflex-base/src/reflex_base/compiler/templates.py index be18fb87e9d..568e98f48e6 100644 --- a/packages/reflex-base/src/reflex_base/compiler/templates.py +++ b/packages/reflex-base/src/reflex_base/compiler/templates.py @@ -275,6 +275,7 @@ def context_template( is_dev_mode: bool, default_color_mode: str, initial_state: dict[str, Any] | None = None, + initial_state_hashes: list[str] | None = None, state_name: str | None = None, client_storage: dict[str, dict[str, dict[str, Any]]] | None = None, disable_react_owner_stacks: bool = False, @@ -283,6 +284,9 @@ def context_template( Args: initial_state: The initial state for the context. + initial_state_hashes: Per-state hashes of ``initial_state`` in sorted + state name order, sent with the first hydrate so the backend can + skip vars still at their default. state_name: The name of the state. client_storage: The client storage for the context. is_dev_mode: Whether the app is in development mode. @@ -315,14 +319,22 @@ def context_template( export const exception_state_name = "{constants.CompileVars.FRONTEND_EXCEPTION_STATE_FULL}" -// These events are triggered on initial load and each page navigation. +// Tracked cookie and local storage vars set in the browser, or undefined if none. +const clientStorageVars = () => {{ + const client_storage_vars = hydrateClientStorage(clientStorage); + if (client_storage_vars && Object.keys(client_storage_vars).length !== 0) {{ + return client_storage_vars; + }} + return undefined; +}} + +// These events are triggered on each client-side page navigation. export const onLoadInternalEvent = () => {{ const internal_events = []; - // Get tracked cookie and local storage vars to send to the backend. - const client_storage_vars = hydrateClientStorage(clientStorage); - // But only send the vars if any are actually set in the browser. - if (client_storage_vars && Object.keys(client_storage_vars).length !== 0) {{ + // Only send the client storage vars if any are actually set in the browser. + const client_storage_vars = clientStorageVars(); + if (client_storage_vars !== undefined) {{ internal_events.push( ReflexEvent( '{state_name}.{constants.CompileVars.UPDATE_VARS_INTERNAL}', @@ -338,11 +350,22 @@ def context_template( return internal_events; }} -// The following events are sent when the websocket connects or reconnects. -export const initialEvents = () => [ - ReflexEvent('{state_name}.{constants.CompileVars.HYDRATE}'), - ...onLoadInternalEvent() -] +// The single event sent when the websocket connects or reconnects: it resets and +// applies client storage, sends the state, and queues the page's on_load events. +// On the first connect the frontend still holds the compiled defaults, so it +// sends their hashes and the backend only returns the vars that differ; any +// later (re)hydrate gets the full state. +export const initialEvents = (first = false) => {{ + const client_storage_vars = clientStorageVars(); + const payload = {{}}; + if (client_storage_vars !== undefined) {{ + payload.vars = client_storage_vars; + }} + if (first) {{ + payload.hashes = initialStateHashes; + }} + return [ReflexEvent('{state_name}.{constants.CompileVars.HYDRATE_AND_LOAD}', payload)]; +}} """ if state_name else """ @@ -404,6 +427,7 @@ def context_template( import {{ jsx }} from "@emotion/react"; {disable_owner_stacks_str} export const initialState = {"{}" if not initial_state else json_dumps(initial_state)} +export const initialStateHashes = {"[]" if not initial_state_hashes else json_dumps(initial_state_hashes)} export const defaultColorMode = {default_color_mode} export const ColorModeContext = createContext({{ diff --git a/packages/reflex-base/src/reflex_base/constants/compiler.py b/packages/reflex-base/src/reflex_base/constants/compiler.py index 036dfe1799a..63106adb7ff 100644 --- a/packages/reflex-base/src/reflex_base/constants/compiler.py +++ b/packages/reflex-base/src/reflex_base/constants/compiler.py @@ -57,6 +57,8 @@ class CompileVars(SimpleNamespace): EVENTS = "events" # The name of the initial hydrate event. HYDRATE = "hydrate" + # The name of the event sent on (re)connect: hydrate plus on_load in one step. + HYDRATE_AND_LOAD = "hydrate_and_load" # The name of the is_hydrated variable. IS_HYDRATED = "is_hydrated" # The name of the function to add events to the queue. 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 de04bac1e93..46c5cd068cd 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 @@ -36,10 +36,13 @@ @functools.lru_cache(maxsize=1) -def _hydrate_event_name(): +def _hydrate_event_names() -> frozenset[str]: from reflex.state import State - return format_event_handler(State.event_handlers["hydrate"]) + return frozenset( + format_event_handler(State.event_handlers[name]) + for name in ("hydrate", "hydrate_and_load") + ) def _check_valid_yield(events: Any, handler_name: str = "unknown") -> Any: @@ -410,7 +413,7 @@ async def _execute_event( ) as state: # Compatibility hack rehydrate the state before processing this event. needs_to_rehydrate = bool( - not state.router_data and event.name != _hydrate_event_name() + not state.router_data and event.name not in _hydrate_event_names() ) # re-assign only when the value is set and different diff --git a/reflex/app.py b/reflex/app.py index d94519104aa..bf9a708b267 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -2002,12 +2002,15 @@ def sid_to_token(self) -> dict[str, str]: # For backward compatibility, expose the underlying dict return self._token_manager.sid_to_token - async def on_connect(self, sid: str, environ: dict): + async def on_connect(self, sid: str, environ: dict, auth: Any = None): """Event for when the websocket is connected. Args: sid: The Socket.IO session id. environ: The request information, including HTTP headers. + auth: The payload of the socket.io CONNECT packet. The frontend + puts its hydrate event here so it is processed without waiting + for the connect acknowledgement round trip. """ if isinstance(self._token_manager, RedisTokenManager): # Make sure this instance is watching for updates from other instances. @@ -2025,6 +2028,9 @@ async def on_connect(self, sid: str, environ: dict): f"Frontend version {subprotocol} for session {sid} does not match the backend version {constants.Reflex.VERSION}." ) + if isinstance(auth, dict) and (boot_event := auth.get("event")) is not None: + await self.on_event(sid, boot_event) + def on_disconnect(self, sid: str) -> asyncio.Task | None: """Event for when the websocket disconnects. diff --git a/reflex/compiler/compiler.py b/reflex/compiler/compiler.py index 3ba8746316e..1d266ca8e46 100644 --- a/reflex/compiler/compiler.py +++ b/reflex/compiler/compiler.py @@ -48,7 +48,7 @@ from reflex.compiler.plugins import default_page_plugins from reflex.compiler.plugins.builtin import collect_var_app_wraps_in_subtree from reflex.compiler.plugins.memoize import MemoizeStatefulPlugin -from reflex.state import BaseState, code_uses_state_contexts +from reflex.state import BaseState, code_uses_state_contexts, state_snapshot_hashes from reflex.utils import console, frontend_skeleton, path_ops, prerequisites from reflex.utils.exec import get_compile_context, is_prod_mode from reflex.utils.prerequisites import get_web_dir @@ -218,9 +218,11 @@ def _compile_contexts(state: type[BaseState] | None, theme: Component | None) -> not is_prod_mode() and not environment.REFLEX_REACT_OWNER_STACKS.get() ) + initial_state = utils.compile_state(state) if state else None return ( templates.context_template( - initial_state=utils.compile_state(state), + initial_state=initial_state, + initial_state_hashes=state_snapshot_hashes(initial_state), state_name=state.get_name(), client_storage=utils.compile_client_storage(state), is_dev_mode=not is_prod_mode(), diff --git a/reflex/state.py b/reflex/state.py index ba118b867d6..f33b2381bd8 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -8,6 +8,7 @@ import copy import dataclasses import functools +import hashlib import inspect import logging import pickle @@ -2474,10 +2475,166 @@ def set_is_hydrated(self, value: bool) -> None: """ self.is_hydrated = value + @event(supersedes=True) + async def hydrate_and_load( + self, + vars: dict[str, Any] | None = None, + hashes: list[str] | None = None, + ) -> list[Event | EventSpec | event.EventCallback] | None: + """Hydrate the frontend and queue the current page's on_load handlers. + + Sent by the frontend once per websocket (re)connect. Doing the client + storage reset, the browser-provided client storage values, the state + snapshot and the on_load enumeration under one state lock avoids three + separate load/persist cycles of the state tree. + + Args: + vars: Client storage vars set in the browser, keyed by fully + qualified var name. + hashes: On the first hydrate of a page, the per-state hashes of the + compiled ``initialState`` the frontend still holds, in sorted + state name order. States whose hash matches the backend's + default snapshot only get the vars that differ from it; + everything else is sent in full. + + Returns: + The on_load events for the current page, if any. + """ + from reflex_base.event.context import EventContext + + self._reset_client_storage() + if vars: + await _apply_client_storage_vars(self, vars) + # The snapshot must carry is_hydrated=False: the frontend skips + # writing client storage for a delta that is not yet hydrated, and + # the reset defaults above must not be written back to the browser. + self.is_hydrated = False + ctx = EventContext.get() + if ctx.emit_delta_impl is not None: + delta = await _resolve_delta(self.dict()) + if hashes: + delta = await _diff_against_initial_state(type(self), delta, hashes) + await ctx.emit_delta(delta=delta) + self._clean() + return _load_events_for_page(self) + T = TypeVar("T", bound=BaseState) +def state_snapshot_hashes(snapshot: Delta) -> list[str]: + """Hash each state's entry of a full-tree snapshot as the frontend receives it. + + Used at compile time for the ``initialState`` baked into the frontend and + at runtime for the backend's own default snapshot, so equal hashes mean the + frontend already holds exactly the backend's defaults for that state. + + Args: + snapshot: A resolved full-tree dict, as returned by ``BaseState.dict``. + + Returns: + A short hex digest of each state's serialized vars, in sorted state + name order. + """ + return [ + hashlib.sha1( + format.json_dumps(snapshot[state_name], sort_keys=True).encode() + ).hexdigest()[:16] + for state_name in sorted(snapshot) + ] + + +# Per root state class: the resolved default snapshot and its per-state hashes. +_initial_snapshot_cache: dict[type[BaseState], tuple[Delta, dict[str, str]]] = {} + + +async def _diff_against_initial_state( + root_cls: type[BaseState], delta: Delta, hashes: list[str] +) -> Delta: + """Drop vars the frontend already holds at their default value. + + Args: + root_cls: The root state class; its default snapshot is computed once. + delta: The resolved full snapshot about to be sent. + hashes: Per-state hashes of the frontend's compiled ``initialState``, + in sorted state name order. + + Returns: + The delta with unchanged vars removed for every state whose compiled + defaults match the backend's, and left untouched for the others. + """ + cached = _initial_snapshot_cache.get(root_cls) + if cached is None: + snapshot = await _resolve_delta( + root_cls(_reflex_internal_init=True).dict(initial=True) + ) + cached = _initial_snapshot_cache[root_cls] = ( + snapshot, + dict(zip(sorted(snapshot), state_snapshot_hashes(snapshot), strict=True)), + ) + defaults, default_hashes = cached + if len(hashes) != len(default_hashes): + # The frontend was compiled against a different set of states. + return delta + frontend_hashes = dict(zip(sorted(default_hashes), hashes, strict=True)) + diff: Delta = {} + for state_name, state_vars in delta.items(): + if frontend_hashes.get(state_name) != default_hashes.get(state_name): + diff[state_name] = state_vars + continue + default_vars = defaults[state_name] + changed = { + name: value + for name, value in state_vars.items() + if name not in default_vars or value != default_vars[name] + } + if changed: + diff[state_name] = changed + return diff + + +async def _apply_client_storage_vars(state: BaseState, vars: dict[str, Any]) -> None: + """Apply browser-provided client storage values to the states that own them. + + Args: + state: Any state in the tree; used to reach the owning substates. + vars: Fully qualified var names mapped to their browser values. + """ + for var, value in vars.items(): + state_name, _, var_name = var.rpartition(".") + var_name = var_name.removesuffix(FIELD_MARKER) + var_state_cls = State.get_class_substate(state_name) + if var_state_cls._is_client_storage(var_name): + var_state = await state.get_state(var_state_cls) + setattr(var_state, var_name, value) + + +def _load_events_for_page( + state: BaseState, +) -> list[Event | EventSpec | event.EventCallback] | None: + """Queue the on_load handlers for the page the client is on. + + Sets ``is_hydrated`` directly when the page has no on_load handlers, so no + extra event round trip is needed for the common case. + + Args: + state: Any state in the tree; ``is_hydrated`` is set through it. + + Returns: + The on_load events followed by the hydrated flip, or None. + """ + load_events = RegistrationContext.get().app.get_load_events(state.router.url.path) + if not load_events: + state.is_hydrated = True + return None + if state.is_hydrated: + state.is_hydrated = False + return [ + *Event.from_event_type(load_events, router_data=state.router_data), + OnLoadInternalState.set_is_hydrated(True), + ] + + def dynamic(func: Callable[[T], Component]): """Create a dynamically generated components from a state class. @@ -2594,13 +2751,7 @@ async def update_vars_internal(self, vars: dict[str, Any]) -> None: Args: vars: The fully qualified vars and values to update. """ - for var, value in vars.items(): - state_name, _, var_name = var.rpartition(".") - var_name = var_name.removesuffix(FIELD_MARKER) - var_state_cls = State.get_class_substate(state_name) - if var_state_cls._is_client_storage(var_name): - var_state = await self.get_state(var_state_cls) - setattr(var_state, var_name, value) + await _apply_client_storage_vars(self, vars) class OnLoadInternalState(State): @@ -2618,20 +2769,7 @@ def on_load_internal(self) -> list[Event | EventSpec | event.EventCallback] | No Returns: The list of events to queue for on load handling. """ - load_events = RegistrationContext.get().app.get_load_events( - self.router.url.path - ) - if not load_events: - self.is_hydrated = True - return None # Fast path for navigation with no on_load events defined. - self.is_hydrated = False - return [ - *Event.from_event_type( - load_events, - router_data=self.router_data, - ), - OnLoadInternalState.set_is_hydrated(True), - ] + return _load_events_for_page(self) @event def set_is_hydrated(self, value: bool) -> None: diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 5ec8de77be6..8178f5b6900 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -736,3 +736,49 @@ async def modify(): ) assert isinstance(final_state, root_state) assert final_state.count == 2 + + +async def test_set_state_checks_lock_once_per_tree( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """Saving a state tree verifies the lock and reads its TTL once, not per substate. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + """ + state_manager_redis._oplock_enabled = False + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + redis = state_manager_redis.redis + lock_key = state_manager_redis._lock_key(token) + real_get, real_pttl = redis.get, redis.pttl + lock_gets: list[Any] = [] + pttls: list[Any] = [] + + async def counting_get(key): + if key == lock_key: + lock_gets.append(key) + return await real_get(key) + + async def counting_pttl(key): + pttls.append(key) + return await real_pttl(key) + + async with state_manager_redis.modify_state(token) as state: + assert len(state.substates) == 2 + state.count = 1 + lock_id = await real_get(lock_key) + redis.get = counting_get + redis.pttl = counting_pttl + try: + await state_manager_redis.set_state(token, state, lock_id=lock_id) + finally: + redis.get = real_get + redis.pttl = real_pttl + + # One lock check and one TTL read for a tree of three states. + assert len(lock_gets) == 1 + assert len(pttls) == 1 + saved = await state_manager_redis.get_state(token) + assert saved.count == 1 diff --git a/tests/units/test_app.py b/tests/units/test_app.py index ac7fca79696..44d1e43f907 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -4344,3 +4344,32 @@ def fake_compile_app(*_args: Any, **_kwargs: Any) -> bool: app._compile() assert not _hash_str_encodings + + +@pytest.mark.asyncio +async def test_on_connect_processes_boot_event_from_auth( + event_namespace: EventNamespace, +): + """The hydrate event carried in the socket.io CONNECT packet is processed on connect. + + Args: + event_namespace: The event namespace. + """ + event_namespace._token_manager = Mock() + event_namespace._token_manager.link_token_to_sid = AsyncMock(return_value=None) + event_namespace.on_event = AsyncMock() + boot_event = {"name": "state.hydrate_and_load", "payload": {}, "router_data": {}} + + await event_namespace.on_connect( + "new_sid", {"QUERY_STRING": "token=abc"}, {"event": boot_event} + ) + event_namespace._token_manager.link_token_to_sid.assert_awaited_once_with( + "abc", "new_sid" + ) + event_namespace.on_event.assert_awaited_once_with("new_sid", boot_event) + + # Without a boot event (or without auth at all) nothing is processed. + event_namespace.on_event.reset_mock() + await event_namespace.on_connect("new_sid", {"QUERY_STRING": "token=abc"}, None) + await event_namespace.on_connect("new_sid", {"QUERY_STRING": "token=abc"}) + event_namespace.on_event.assert_not_awaited() diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 37c0ed2fc4c..f5453ce5bcf 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -50,6 +50,7 @@ from reflex.istate.manager.redis import StateManagerRedis from reflex.istate.manager.token import BaseStateToken from reflex.istate.proxy import MutableProxy, StateProxy +from reflex.istate.storage import Cookie from reflex.state import BaseState, ImmutableStateError, OnLoadInternalState, State from reflex.testing import chdir from reflex.utils import prerequisites @@ -2171,7 +2172,8 @@ async def _coro_blocker(): # When Oplock is enabled, we don't warn when lock is held too long. assert not lock_warnings else: - assert len(lock_warnings) == 7 + # One warning per state-tree save, not one per substate. + assert len(lock_warnings) == 1 class CopyingAsyncMock(AsyncMock): @@ -3388,14 +3390,15 @@ def index(): ) await on_load_future.wait_all() - # The processor chains all events: on_load_internal sets is_hydrated=False, - # then the on_load handler runs, then set_is_hydrated(True) runs. - # First delta: router + is_hydrated=False + # The processor chains all events: hydrate leaves is_hydrated=False, then + # the on_load handler runs, then set_is_hydrated(True) runs. + # First delta: router only. on_load_internal does not re-send the + # is_hydrated=False the hydrate snapshot already carried. assert len(emitted_deltas) == 1 + len(expected) first_token, first_delta = emitted_deltas[0] assert first_token == token assert first_delta[State.get_full_name()].pop("router" + FIELD_MARKER) is not None - assert first_delta == exp_is_hydrated(State, False) + assert first_delta == {State.get_full_name(): {}} # Find the deltas containing the test handler's state change for (delta_token, actual_delta), expected_delta in zip( @@ -3451,11 +3454,11 @@ def index(): ) await processor.join() - # First delta: router + is_hydrated=False + # First delta: router only (is_hydrated=False was already in the hydrate snapshot) assert len(emitted_deltas) >= 2 first_delta = emitted_deltas[0][1] assert first_delta[State.get_full_name()].pop("router" + FIELD_MARKER) is not None - assert first_delta == exp_is_hydrated(State, False) + assert first_delta == {State.get_full_name(): {}} # Find deltas containing the test handler's state change (num incremented twice) handler_deltas = [ @@ -5380,3 +5383,158 @@ def test_setattr_alias_annotated_var(mocker: MockerFixture): state.key = 1 # pyright: ignore[reportAttributeAccessIssue] assert state.key == 1 error_mock.assert_called_once() + + +class BootCookieState(State): + """A state with a cookie var and an on_load handler for hydrate_and_load tests.""" + + flavor: str = Cookie("plain") + loads: int = 0 + + def on_load_handler(self): + """Count page loads.""" + self.loads += 1 + + +async def test_hydrate_and_load_single_lock_cycle( + app_module_mock, + token, + mock_root_event_context: EventContext, + mock_base_state_event_processor: BaseStateEventProcessor, + emitted_deltas: list, +): + """One hydrate_and_load event resets/applies client storage, snapshots, and queues on_load. + + 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. + emitted_deltas: List to capture emitted deltas. + """ + assert State.event_handlers["hydrate_and_load"].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=BootCookieState.on_load_handler) + app._compile_page("index") + + boot_name = format.format_event_handler( + State.hydrate_and_load # pyright: ignore[reportArgumentType] + ) + cookie_key = f"{BootCookieState.get_full_name()}.flavor{FIELD_MARKER}" + router_data = {RouteVar.PATH: "/", RouteVar.ORIGIN: "/", RouteVar.QUERY: {}} + + async with mock_base_state_event_processor as processor: + future = await processor.enqueue( + token, + Event( + name=boot_name, + payload={"vars": {cookie_key: "chocolate"}}, + router_data=router_data, + ), + ) + await future.wait_all() + + state_name = State.get_full_name() + hydrated_key = CompileVars.IS_HYDRATED + FIELD_MARKER + # Snapshot (not hydrated, browser cookie applied), on_load delta, hydrated. + snapshot = emitted_deltas[0][1] + assert snapshot[state_name][hydrated_key] is False + assert ( + snapshot[BootCookieState.get_full_name()]["flavor" + FIELD_MARKER] + == "chocolate" + ) + assert snapshot[BootCookieState.get_full_name()]["loads" + FIELD_MARKER] == 0 + assert [d for _, d in emitted_deltas[1:]] == [ + {BootCookieState.get_full_name(): {"loads" + FIELD_MARKER: 1}}, + exp_is_hydrated(State, True), + ] + + # The next hydrate resets the cookie var to its default when the browser + # no longer sends it, and the on_load chain runs again. + emitted_deltas.clear() + async with mock_base_state_event_processor as processor: + future = await processor.enqueue( + token, Event(name=boot_name, payload={}, router_data=router_data) + ) + await future.wait_all() + snapshot = emitted_deltas[0][1] + assert snapshot[BootCookieState.get_full_name()]["flavor" + FIELD_MARKER] == "plain" + assert emitted_deltas[1][1] == { + BootCookieState.get_full_name(): {"loads" + FIELD_MARKER: 2} + } + + +async def test_hydrate_and_load_diffs_against_compiled_defaults( + app_module_mock, + token, + mock_root_event_context: EventContext, + mock_base_state_event_processor: BaseStateEventProcessor, + emitted_deltas: list, +): + """With matching initialState hashes only vars that differ from the defaults are sent. + + 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. + emitted_deltas: List to capture emitted deltas. + """ + from reflex.compiler.utils import compile_state + from reflex.state import state_snapshot_hashes + + 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) + app._compile_page("index") + + boot_name = format.format_event_handler( + State.hydrate_and_load # pyright: ignore[reportArgumentType] + ) + cookie_key = f"{BootCookieState.get_full_name()}.flavor{FIELD_MARKER}" + router_data = {RouteVar.PATH: "/", RouteVar.ORIGIN: "/", RouteVar.QUERY: {}} + hashes = state_snapshot_hashes(compile_state(State)) + + async with mock_base_state_event_processor as processor: + future = await processor.enqueue( + token, + Event( + name=boot_name, + payload={"vars": {cookie_key: "chocolate"}, "hashes": hashes}, + router_data=router_data, + ), + ) + await future.wait_all() + + snapshot = emitted_deltas[0][1] + # Only the root router and the changed cookie var differ from the compiled defaults. + assert set(snapshot) == {State.get_full_name(), BootCookieState.get_full_name()} + assert set(snapshot[State.get_full_name()]) == {"router" + FIELD_MARKER} + assert snapshot[BootCookieState.get_full_name()] == { + "flavor" + FIELD_MARKER: "chocolate" + } + + # Hashes compiled against a different set of states fall back to the full snapshot. + emitted_deltas.clear() + async with mock_base_state_event_processor as processor: + future = await processor.enqueue( + token, + Event( + name=boot_name, + payload={"hashes": hashes[:-1]}, + router_data=router_data, + ), + ) + await future.wait_all() + snapshot = emitted_deltas[0][1] + assert "loads" + FIELD_MARKER in snapshot[BootCookieState.get_full_name()] From fa02db87dd04f78988b3329034118eda0ab0a086 Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 19:54:11 +0000 Subject: [PATCH 03/14] Rebuild the default-state snapshot when state classes change; add news and tests The cached default snapshot used to diff the first hydrate against the compiled initialState is keyed by the number of registered state classes, so states defined after the first hydrate (dynamic component states, test suites) never compare against a stale snapshot; states missing from it are sent in full. The hydrate tests move to the isolated-registry processor tests so states other unit tests leave behind cannot break a full snapshot. --- news/+first-load-hydrate.performance.md | 1 + .../news/+first-load-hydrate.performance.md | 1 + reflex/state.py | 20 ++- tests/units/istate/manager/test_redis.py | 5 +- .../processor/test_base_state_processor.py | 128 ++++++++++++++ tests/units/test_state.py | 156 ------------------ 6 files changed, 146 insertions(+), 165 deletions(-) create mode 100644 news/+first-load-hydrate.performance.md create mode 100644 packages/reflex-base/news/+first-load-hydrate.performance.md diff --git a/news/+first-load-hydrate.performance.md b/news/+first-load-hydrate.performance.md new file mode 100644 index 00000000000..0199db48f6e --- /dev/null +++ b/news/+first-load-hydrate.performance.md @@ -0,0 +1 @@ +Faster first page load: the frontend now sends a single `hydrate_and_load` event inside the websocket connect packet instead of three events after the connect acknowledgement, the backend handles it under one state lock, and on the first hydrate of a page only the vars that differ from the compiled defaults are sent (a 20-substate app went from 36 KB to 2 KB on the wire and from 5 to 3 frames). Redis-backed apps also do far fewer commands per page load: the lock is verified once per state-tree save instead of once per substate, and neither the websocket connect nor the final `is_hydrated` flip load every substate any more. diff --git a/packages/reflex-base/news/+first-load-hydrate.performance.md b/packages/reflex-base/news/+first-load-hydrate.performance.md new file mode 100644 index 00000000000..1b49d818e03 --- /dev/null +++ b/packages/reflex-base/news/+first-load-hydrate.performance.md @@ -0,0 +1 @@ +The compiled frontend hydrates with one `hydrate_and_load` event carried in the socket.io connect packet, sending per-state hashes of its compiled `initialState` so the backend can skip vars still at their default; client-side navigation keeps using `on_load_internal` as before. diff --git a/reflex/state.py b/reflex/state.py index f33b2381bd8..a325f9381c1 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -2544,8 +2544,9 @@ def state_snapshot_hashes(snapshot: Delta) -> list[str]: ] -# Per root state class: the resolved default snapshot and its per-state hashes. -_initial_snapshot_cache: dict[type[BaseState], tuple[Delta, dict[str, str]]] = {} +# Per root state class: the number of state classes the snapshot was built +# for, the resolved default snapshot and its per-state hashes. +_initial_snapshot_cache: dict[type[BaseState], tuple[int, Delta, dict[str, str]]] = {} async def _diff_against_initial_state( @@ -2564,25 +2565,30 @@ async def _diff_against_initial_state( defaults match the backend's, and left untouched for the others. """ cached = _initial_snapshot_cache.get(root_cls) - if cached is None: + n_state_classes = len(all_base_state_classes) + if cached is None or cached[0] != n_state_classes: + # Rebuilt when state classes were defined after the last snapshot. snapshot = await _resolve_delta( root_cls(_reflex_internal_init=True).dict(initial=True) ) cached = _initial_snapshot_cache[root_cls] = ( + n_state_classes, snapshot, dict(zip(sorted(snapshot), state_snapshot_hashes(snapshot), strict=True)), ) - defaults, default_hashes = cached + _, defaults, default_hashes = cached if len(hashes) != len(default_hashes): # The frontend was compiled against a different set of states. return delta frontend_hashes = dict(zip(sorted(default_hashes), hashes, strict=True)) diff: Delta = {} for state_name, state_vars in delta.items(): - if frontend_hashes.get(state_name) != default_hashes.get(state_name): + default_vars = defaults.get(state_name) + if default_vars is None or frontend_hashes.get( + state_name + ) != default_hashes.get(state_name): diff[state_name] = state_vars continue - default_vars = defaults[state_name] changed = { name: value for name, value in state_vars.items() @@ -2610,7 +2616,7 @@ async def _apply_client_storage_vars(state: BaseState, vars: dict[str, Any]) -> def _load_events_for_page( - state: BaseState, + state: State, ) -> list[Event | EventSpec | event.EventCallback] | None: """Queue the on_load handlers for the page the client is on. diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 8178f5b6900..d7498b36382 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -769,8 +769,8 @@ async def counting_pttl(key): assert len(state.substates) == 2 state.count = 1 lock_id = await real_get(lock_key) - redis.get = counting_get - redis.pttl = counting_pttl + redis.get = counting_get # pyright: ignore[reportAttributeAccessIssue] + redis.pttl = counting_pttl # pyright: ignore[reportAttributeAccessIssue] try: await state_manager_redis.set_state(token, state, lock_id=lock_id) finally: @@ -781,4 +781,5 @@ async def counting_pttl(key): assert len(lock_gets) == 1 assert len(pttls) == 1 saved = await state_manager_redis.get_state(token) + assert isinstance(saved, root_state) assert saved.count == 1 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 a70130f70c1..85681d7fa4f 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 @@ -745,3 +745,131 @@ def raise_on_modify(*args, **kwargs): object.__setattr__(root_ctx.state_manager, "modify_state_with_links", original) assert proxy._self_entered_context is False + + +def _boot_event(name: str, payload: dict[str, Any]) -> Event: + return Event( + name=name, + payload=payload, + router_data={"pathname": "/", "asPath": "/", "query": {}}, + ) + + +async def test_hydrate_and_load_single_lock_cycle( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """One hydrate_and_load event resets/applies client storage, snapshots, and queues on_load. + + 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. + """ + assert State.event_handlers["hydrate_and_load"].supersedes + + class CookieState(State): + flavor: str = rx.Cookie("plain") + loads: int = 0 + + @event + def on_load_handler(self): + self.loads += 1 + + wired_app.add_page( + lambda: rx.text(CookieState.flavor), + route="/", + on_load=CookieState.on_load_handler, + ) + wired_app._compile_page("index") + boot_name = Event.from_event_type(State.hydrate_and_load())[0].name # pyright: ignore[reportCallIssue] + cookie_key = f"{CookieState.get_full_name()}.flavor{FIELD_MARKER}" + state_name = State.get_full_name() + hydrated_key = CompileVars.IS_HYDRATED + FIELD_MARKER + + async with real_base_state_processor as processor: + future = await processor.enqueue( + token, _boot_event(boot_name, {"vars": {cookie_key: "chocolate"}}) + ) + await future.wait_all() + + # Snapshot (not hydrated, browser cookie applied), the on_load chain, hydrated. + snapshot = emitted_deltas[0][1] + assert snapshot[state_name][hydrated_key] is False + assert snapshot[CookieState.get_full_name()]["flavor" + FIELD_MARKER] == "chocolate" + assert snapshot[CookieState.get_full_name()]["loads" + FIELD_MARKER] == 0 + assert [d for _, d in emitted_deltas[1:]] == [ + {CookieState.get_full_name(): {"loads" + FIELD_MARKER: 1}}, + {state_name: {hydrated_key: True}}, + ] + + # The next hydrate resets the cookie var when the browser no longer sends + # it, and the on_load chain runs again. + emitted_deltas.clear() + async with real_base_state_processor as processor: + future = await processor.enqueue(token, _boot_event(boot_name, {})) + await future.wait_all() + snapshot = emitted_deltas[0][1] + assert snapshot[CookieState.get_full_name()]["flavor" + FIELD_MARKER] == "plain" + assert emitted_deltas[1][1] == { + CookieState.get_full_name(): {"loads" + FIELD_MARKER: 2} + } + + +async def test_hydrate_and_load_diffs_against_compiled_defaults( + wired_app: App, + real_base_state_processor: BaseStateEventProcessor, + emitted_deltas: list[tuple[str, Mapping[str, Mapping[str, Any]]]], + token: str, +): + """With matching initialState hashes only vars that differ from the defaults are sent. + + 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.compiler.utils import compile_state + from reflex.state import state_snapshot_hashes + + class CookieState(State): + flavor: str = rx.Cookie("plain") + loads: int = 0 + + wired_app.add_page(lambda: rx.text(CookieState.flavor), route="/") + wired_app._compile_page("index") + boot_name = Event.from_event_type(State.hydrate_and_load())[0].name # pyright: ignore[reportCallIssue] + cookie_key = f"{CookieState.get_full_name()}.flavor{FIELD_MARKER}" + state_name = State.get_full_name() + hashes = state_snapshot_hashes(compile_state(State)) + + async with real_base_state_processor as processor: + future = await processor.enqueue( + token, + _boot_event( + boot_name, {"vars": {cookie_key: "chocolate"}, "hashes": hashes} + ), + ) + await future.wait_all() + + snapshot = emitted_deltas[0][1] + # Only the root router and the changed cookie var differ from the compiled defaults. + assert set(snapshot) == {state_name, CookieState.get_full_name()} + assert set(snapshot[state_name]) == {"router" + FIELD_MARKER} + assert snapshot[CookieState.get_full_name()] == { + "flavor" + FIELD_MARKER: "chocolate" + } + + # Hashes compiled against a different set of states fall back to the full snapshot. + emitted_deltas.clear() + async with real_base_state_processor as processor: + future = await processor.enqueue( + token, _boot_event(boot_name, {"hashes": hashes[:-1]}) + ) + await future.wait_all() + snapshot = emitted_deltas[0][1] + assert "loads" + FIELD_MARKER in snapshot[CookieState.get_full_name()] diff --git a/tests/units/test_state.py b/tests/units/test_state.py index f5453ce5bcf..129d5ce4f6e 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -50,7 +50,6 @@ from reflex.istate.manager.redis import StateManagerRedis from reflex.istate.manager.token import BaseStateToken from reflex.istate.proxy import MutableProxy, StateProxy -from reflex.istate.storage import Cookie from reflex.state import BaseState, ImmutableStateError, OnLoadInternalState, State from reflex.testing import chdir from reflex.utils import prerequisites @@ -5383,158 +5382,3 @@ def test_setattr_alias_annotated_var(mocker: MockerFixture): state.key = 1 # pyright: ignore[reportAttributeAccessIssue] assert state.key == 1 error_mock.assert_called_once() - - -class BootCookieState(State): - """A state with a cookie var and an on_load handler for hydrate_and_load tests.""" - - flavor: str = Cookie("plain") - loads: int = 0 - - def on_load_handler(self): - """Count page loads.""" - self.loads += 1 - - -async def test_hydrate_and_load_single_lock_cycle( - app_module_mock, - token, - mock_root_event_context: EventContext, - mock_base_state_event_processor: BaseStateEventProcessor, - emitted_deltas: list, -): - """One hydrate_and_load event resets/applies client storage, snapshots, and queues on_load. - - 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. - emitted_deltas: List to capture emitted deltas. - """ - assert State.event_handlers["hydrate_and_load"].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=BootCookieState.on_load_handler) - app._compile_page("index") - - boot_name = format.format_event_handler( - State.hydrate_and_load # pyright: ignore[reportArgumentType] - ) - cookie_key = f"{BootCookieState.get_full_name()}.flavor{FIELD_MARKER}" - router_data = {RouteVar.PATH: "/", RouteVar.ORIGIN: "/", RouteVar.QUERY: {}} - - async with mock_base_state_event_processor as processor: - future = await processor.enqueue( - token, - Event( - name=boot_name, - payload={"vars": {cookie_key: "chocolate"}}, - router_data=router_data, - ), - ) - await future.wait_all() - - state_name = State.get_full_name() - hydrated_key = CompileVars.IS_HYDRATED + FIELD_MARKER - # Snapshot (not hydrated, browser cookie applied), on_load delta, hydrated. - snapshot = emitted_deltas[0][1] - assert snapshot[state_name][hydrated_key] is False - assert ( - snapshot[BootCookieState.get_full_name()]["flavor" + FIELD_MARKER] - == "chocolate" - ) - assert snapshot[BootCookieState.get_full_name()]["loads" + FIELD_MARKER] == 0 - assert [d for _, d in emitted_deltas[1:]] == [ - {BootCookieState.get_full_name(): {"loads" + FIELD_MARKER: 1}}, - exp_is_hydrated(State, True), - ] - - # The next hydrate resets the cookie var to its default when the browser - # no longer sends it, and the on_load chain runs again. - emitted_deltas.clear() - async with mock_base_state_event_processor as processor: - future = await processor.enqueue( - token, Event(name=boot_name, payload={}, router_data=router_data) - ) - await future.wait_all() - snapshot = emitted_deltas[0][1] - assert snapshot[BootCookieState.get_full_name()]["flavor" + FIELD_MARKER] == "plain" - assert emitted_deltas[1][1] == { - BootCookieState.get_full_name(): {"loads" + FIELD_MARKER: 2} - } - - -async def test_hydrate_and_load_diffs_against_compiled_defaults( - app_module_mock, - token, - mock_root_event_context: EventContext, - mock_base_state_event_processor: BaseStateEventProcessor, - emitted_deltas: list, -): - """With matching initialState hashes only vars that differ from the defaults are sent. - - 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. - emitted_deltas: List to capture emitted deltas. - """ - from reflex.compiler.utils import compile_state - from reflex.state import state_snapshot_hashes - - 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) - app._compile_page("index") - - boot_name = format.format_event_handler( - State.hydrate_and_load # pyright: ignore[reportArgumentType] - ) - cookie_key = f"{BootCookieState.get_full_name()}.flavor{FIELD_MARKER}" - router_data = {RouteVar.PATH: "/", RouteVar.ORIGIN: "/", RouteVar.QUERY: {}} - hashes = state_snapshot_hashes(compile_state(State)) - - async with mock_base_state_event_processor as processor: - future = await processor.enqueue( - token, - Event( - name=boot_name, - payload={"vars": {cookie_key: "chocolate"}, "hashes": hashes}, - router_data=router_data, - ), - ) - await future.wait_all() - - snapshot = emitted_deltas[0][1] - # Only the root router and the changed cookie var differ from the compiled defaults. - assert set(snapshot) == {State.get_full_name(), BootCookieState.get_full_name()} - assert set(snapshot[State.get_full_name()]) == {"router" + FIELD_MARKER} - assert snapshot[BootCookieState.get_full_name()] == { - "flavor" + FIELD_MARKER: "chocolate" - } - - # Hashes compiled against a different set of states fall back to the full snapshot. - emitted_deltas.clear() - async with mock_base_state_event_processor as processor: - future = await processor.enqueue( - token, - Event( - name=boot_name, - payload={"hashes": hashes[:-1]}, - router_data=router_data, - ), - ) - await future.wait_all() - snapshot = emitted_deltas[0][1] - assert "loads" + FIELD_MARKER in snapshot[BootCookieState.get_full_name()] From 2d9abce959bcd93156b1105b577ceb06cfe297a4 Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 19:58:58 +0000 Subject: [PATCH 04/14] Keep the is_hydrated=False delta on on_load_internal The frontend sets is_hydrated=false locally on navigation, but the on_load chain has always re-sent it in its first delta and tests and downstream code depend on that ordering; keep it rather than saving one 89-byte frame. --- reflex/state.py | 3 +-- .../event/processor/test_base_state_processor.py | 3 ++- tests/units/test_state.py | 13 ++++++------- 3 files changed, 9 insertions(+), 10 deletions(-) diff --git a/reflex/state.py b/reflex/state.py index a325f9381c1..ef86a25e508 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -2633,8 +2633,7 @@ def _load_events_for_page( if not load_events: state.is_hydrated = True return None - if state.is_hydrated: - state.is_hydrated = False + state.is_hydrated = False return [ *Event.from_event_type(load_events, router_data=state.router_data), OnLoadInternalState.set_is_hydrated(True), 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 85681d7fa4f..336ae12765a 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 @@ -802,6 +802,7 @@ def on_load_handler(self): assert snapshot[CookieState.get_full_name()]["flavor" + FIELD_MARKER] == "chocolate" assert snapshot[CookieState.get_full_name()]["loads" + FIELD_MARKER] == 0 assert [d for _, d in emitted_deltas[1:]] == [ + {state_name: {hydrated_key: False}}, {CookieState.get_full_name(): {"loads" + FIELD_MARKER: 1}}, {state_name: {hydrated_key: True}}, ] @@ -814,7 +815,7 @@ def on_load_handler(self): await future.wait_all() snapshot = emitted_deltas[0][1] assert snapshot[CookieState.get_full_name()]["flavor" + FIELD_MARKER] == "plain" - assert emitted_deltas[1][1] == { + assert emitted_deltas[2][1] == { CookieState.get_full_name(): {"loads" + FIELD_MARKER: 2} } diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 129d5ce4f6e..824848da4ac 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -3389,15 +3389,14 @@ def index(): ) await on_load_future.wait_all() - # The processor chains all events: hydrate leaves is_hydrated=False, then - # the on_load handler runs, then set_is_hydrated(True) runs. - # First delta: router only. on_load_internal does not re-send the - # is_hydrated=False the hydrate snapshot already carried. + # The processor chains all events: on_load_internal sets is_hydrated=False, + # then the on_load handler runs, then set_is_hydrated(True) runs. + # First delta: router + is_hydrated=False assert len(emitted_deltas) == 1 + len(expected) first_token, first_delta = emitted_deltas[0] assert first_token == token assert first_delta[State.get_full_name()].pop("router" + FIELD_MARKER) is not None - assert first_delta == {State.get_full_name(): {}} + assert first_delta == exp_is_hydrated(State, False) # Find the deltas containing the test handler's state change for (delta_token, actual_delta), expected_delta in zip( @@ -3453,11 +3452,11 @@ def index(): ) await processor.join() - # First delta: router only (is_hydrated=False was already in the hydrate snapshot) + # First delta: router + is_hydrated=False assert len(emitted_deltas) >= 2 first_delta = emitted_deltas[0][1] assert first_delta[State.get_full_name()].pop("router" + FIELD_MARKER) is not None - assert first_delta == {State.get_full_name(): {}} + assert first_delta == exp_is_hydrated(State, False) # Find deltas containing the test handler's state change (num incremented twice) handler_deltas = [ From 2850fa20e000f7e74c59e3ba0885d8e3d70e0c2c Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 20:49:26 +0000 Subject: [PATCH 05/14] Type the compiled initial state as non-optional before hashing it --- reflex/compiler/compiler.py | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/reflex/compiler/compiler.py b/reflex/compiler/compiler.py index 1d266ca8e46..867453d8eb0 100644 --- a/reflex/compiler/compiler.py +++ b/reflex/compiler/compiler.py @@ -218,23 +218,21 @@ def _compile_contexts(state: type[BaseState] | None, theme: Component | None) -> not is_prod_mode() and not environment.REFLEX_REACT_OWNER_STACKS.get() ) - initial_state = utils.compile_state(state) if state else None - return ( - templates.context_template( - initial_state=initial_state, - initial_state_hashes=state_snapshot_hashes(initial_state), - state_name=state.get_name(), - client_storage=utils.compile_client_storage(state), - is_dev_mode=not is_prod_mode(), - default_color_mode=default_color_mode, - disable_react_owner_stacks=disable_react_owner_stacks, - ) - if state - else templates.context_template( + if state is None: + return templates.context_template( is_dev_mode=not is_prod_mode(), default_color_mode=default_color_mode, disable_react_owner_stacks=disable_react_owner_stacks, ) + initial_state = utils.compile_state(state) + return templates.context_template( + initial_state=initial_state, + initial_state_hashes=state_snapshot_hashes(initial_state), + state_name=state.get_name(), + client_storage=utils.compile_client_storage(state), + is_dev_mode=not is_prod_mode(), + default_color_mode=default_color_mode, + disable_react_owner_stacks=disable_react_owner_stacks, ) From 16d25bbd1facc97811361f5704e0cc7faaa6717d Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 20:53:47 +0000 Subject: [PATCH 06/14] Address review findings on the first-load hydrate - Serialize touched states first and verify the lock immediately before one pipelined write, instead of a check followed by concurrent per-substate writes; this also collapses N SET round trips into one. - Drop the sid/token link when the boot event carried in the CONNECT packet fails, since a refused connect never reaches on_disconnect. - Compare hydrate values with their compiled defaults in serialized form, so Python-equal but JSON-distinct values (1 vs 1.0, 0 vs False) are still sent. - Seed the default snapshot from compile_state, so a backend in the compiling process diffs against the exact compiled values and hot reloads refresh it. - Use CompileVars for the hydrate handler names in the processor. --- .../event/processor/base_state_processor.py | 3 +- reflex/app.py | 9 +- reflex/compiler/utils.py | 12 +-- reflex/istate/manager/redis.py | 71 ++++++++------- reflex/state.py | 88 ++++++++++++++----- .../processor/test_base_state_processor.py | 24 ++++- tests/units/test_app.py | 24 +++++ 7 files changed, 168 insertions(+), 63 deletions(-) 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 46c5cd068cd..6d5b815dc24 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 @@ -16,6 +16,7 @@ from reflex.istate.manager.token import BaseStateToken from reflex.istate.proxy import StateProxy from reflex.utils import types +from reflex_base.constants import CompileVars from reflex_base.event.context import EventContext from reflex_base.event.processor.event_processor import EventProcessor, EventQueueEntry from reflex_base.registry import RegisteredEventHandler @@ -41,7 +42,7 @@ def _hydrate_event_names() -> frozenset[str]: return frozenset( format_event_handler(State.event_handlers[name]) - for name in ("hydrate", "hydrate_and_load") + for name in (CompileVars.HYDRATE, CompileVars.HYDRATE_AND_LOAD) ) diff --git a/reflex/app.py b/reflex/app.py index bf9a708b267..5ecbf21f010 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -2029,7 +2029,14 @@ async def on_connect(self, sid: str, environ: dict, auth: Any = None): ) if isinstance(auth, dict) and (boot_event := auth.get("event")) is not None: - await self.on_event(sid, boot_event) + try: + await self.on_event(sid, boot_event) + except Exception: + # A refused connect never reaches on_disconnect, so drop the + # token link made above before the error refuses the connect. + if (linked_token := self.sid_to_token.get(sid)) is not None: + await self._token_manager.disconnect_token(linked_token, sid) + raise def on_disconnect(self, sid: str) -> asyncio.Task | None: """Event for when the websocket disconnects. diff --git a/reflex/compiler/utils.py b/reflex/compiler/utils.py index a3e2c4a2c49..23c97c26bc8 100644 --- a/reflex/compiler/utils.py +++ b/reflex/compiler/utils.py @@ -38,7 +38,7 @@ from reflex_components_core.el.elements.sectioning import Body from reflex.istate.storage import Cookie, LocalStorage, SessionStorage -from reflex.state import BaseState, _resolve_delta +from reflex.state import BaseState, _resolve_delta, cache_initial_snapshot from reflex.utils import path_ops from reflex.utils.prerequisites import get_web_dir @@ -222,16 +222,16 @@ def compile_state(state: type[BaseState]) -> dict: try: _ = asyncio.get_running_loop() except RuntimeError: - pass + # Normally the compile runs before any event loop starts, we asyncio.run is available for calling. + resolved_initial_state = asyncio.run(_resolve_delta(initial_state)) else: with concurrent.futures.ThreadPoolExecutor() as pool: resolved_initial_state = pool.submit( asyncio.run, _resolve_delta(initial_state) ).result() - return _sorted_keys(resolved_initial_state) - - # Normally the compile runs before any event loop starts, we asyncio.run is available for calling. - return _sorted_keys(asyncio.run(_resolve_delta(initial_state))) + # A backend in this process diffs first hydrates against these exact values. + cache_initial_snapshot(state, resolved_initial_state) + return _sorted_keys(resolved_initial_state) def _compile_client_storage_field( diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 7101cb9c4e2..dc9d599ce9e 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -397,6 +397,15 @@ async def set_state( RuntimeError: If the state instance doesn't match the state name in the token. """ token = self._coerce_token(token) + if isinstance(token, BaseStateToken): + # Serialize the touched states before checking the lock, so the + # check sits right before the single pipelined write. + writes = self._collect_state_writes(token, cast(BaseState, state)) + else: + # Non-BaseState token: simple single-key write. + pickle_state = token.serialize(state) + writes = [(str(token), pickle_state)] if pickle_state else [] + # Check that we're holding the lock. if ( lock_id is not None @@ -416,14 +425,11 @@ async def set_state( ) raise LockExpiredError(msg) - if not isinstance(token, BaseStateToken): - # Non-BaseState token: simple single-key write. - pickle_state = token.serialize(state) - if pickle_state: - await self.redis.set(str(token), pickle_state, ex=self.token_expiration) - return - - if lock_id is not None and token.lock_key not in self._local_leases: + if ( + isinstance(token, BaseStateToken) + and lock_id is not None + and token.lock_key not in self._local_leases + ): time_taken = ( self.lock_expiration - (await self.redis.pttl(self._lock_key(token))) ) / 1000 @@ -440,38 +446,35 @@ async def set_state( extra={"dedupe": True}, ) - await self._set_state_tree(token, cast(BaseState, state)) + if len(writes) == 1: + await self.redis.set(writes[0][0], writes[0][1], ex=self.token_expiration) + elif writes: + pipeline = self.redis.pipeline() + for key, pickle_state in writes: + pipeline.set(key, pickle_state, ex=self.token_expiration) + await pipeline.execute() - async def _set_state_tree(self, token: BaseStateToken, base_state: BaseState): - """Persist a state and, concurrently, every substate attached to it. - - The lock check and the hold-time warning happen once in ``set_state``; - this recursion only writes the keys that were touched. + def _collect_state_writes( + self, token: BaseStateToken, base_state: BaseState + ) -> list[tuple[str, bytes]]: + """Serialize a state and every substate attached to it that was touched. Args: token: The token (any state class) identifying the client. base_state: The state instance whose tree to persist. - """ - tasks = [ - asyncio.create_task( - self._set_state_tree(token, substate), - name=f"reflex_set_state|{token.lock_key}|{substate.get_full_name()}", - ) - for substate in base_state.substates.values() - ] - # Persist only the given state (parents or substates are excluded by BaseState.__getstate__). - if base_state._get_was_touched(): - pickle_state = base_state._serialize() - if pickle_state: - await self.redis.set( - str(token.with_cls(type(base_state))), - pickle_state, - ex=self.token_expiration, - ) - # Wait for substates to be persisted. - for t in tasks: - await t + Returns: + The redis keys and pickled payloads to write. + """ + writes: list[tuple[str, bytes]] = [] + stack = [base_state] + while stack: + state = stack.pop() + # Persist only the given state (parents or substates are excluded by BaseState.__getstate__). + if state._get_was_touched() and (pickle_state := state._serialize()): + writes.append((str(token.with_cls(type(state))), pickle_state)) + stack.extend(state.substates.values()) + return writes @contextlib.asynccontextmanager async def _try_modify_state( diff --git a/reflex/state.py b/reflex/state.py index ef86a25e508..ecc6b66f7d6 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -2544,9 +2544,55 @@ def state_snapshot_hashes(snapshot: Delta) -> list[str]: ] -# Per root state class: the number of state classes the snapshot was built -# for, the resolved default snapshot and its per-state hashes. -_initial_snapshot_cache: dict[type[BaseState], tuple[int, Delta, dict[str, str]]] = {} +@dataclasses.dataclass(frozen=True) +class _InitialSnapshot: + """The defaults the compiled frontend holds, as the backend last computed them.""" + + # The number of state classes registered when the snapshot was taken. + n_state_classes: int + # Per state full name, each var's serialized default value. + serialized: dict[str, dict[str, str]] + # Per state full name, the hash of its serialized defaults. + hashes: dict[str, str] + + +# Per root state class: the resolved default snapshot the frontend was compiled with. +_initial_snapshot_cache: dict[type[BaseState], _InitialSnapshot] = {} + + +def _serialize_var(value: Any) -> str: + """Serialize a var value the way it reaches the frontend. + + Args: + value: The resolved var value. + + Returns: + The JSON text of the value. + """ + return format.json_dumps(value, sort_keys=True) + + +def cache_initial_snapshot(root_cls: type[BaseState], snapshot: Delta) -> None: + """Remember the default snapshot the frontend was compiled with. + + Called by the compiler so a backend running in the compiling process + diffs hydrates against exactly the values baked into the bundle, and + picks up new defaults on every hot reload. + + Args: + root_cls: The root state class the snapshot was taken from. + snapshot: The resolved full-tree default snapshot. + """ + _initial_snapshot_cache[root_cls] = _InitialSnapshot( + n_state_classes=len(all_base_state_classes), + serialized={ + state_name: {name: _serialize_var(value) for name, value in vars.items()} + for state_name, vars in snapshot.items() + }, + hashes=dict( + zip(sorted(snapshot), state_snapshot_hashes(snapshot), strict=True) + ), + ) async def _diff_against_initial_state( @@ -2554,6 +2600,10 @@ async def _diff_against_initial_state( ) -> Delta: """Drop vars the frontend already holds at their default value. + Values are compared in their serialized form, so a value that is + Python-equal but JSON-distinct from its default (``1`` vs ``1.0``, + ``0`` vs ``False``) is still sent. + Args: root_cls: The root state class; its default snapshot is computed once. delta: The resolved full snapshot about to be sent. @@ -2565,34 +2615,32 @@ async def _diff_against_initial_state( defaults match the backend's, and left untouched for the others. """ cached = _initial_snapshot_cache.get(root_cls) - n_state_classes = len(all_base_state_classes) - if cached is None or cached[0] != n_state_classes: - # Rebuilt when state classes were defined after the last snapshot. - snapshot = await _resolve_delta( - root_cls(_reflex_internal_init=True).dict(initial=True) + if cached is None or cached.n_state_classes != len(all_base_state_classes): + # No compile happened in this process, or state classes were defined + # after the last snapshot: compute the defaults the same way. + cache_initial_snapshot( + root_cls, + await _resolve_delta( + root_cls(_reflex_internal_init=True).dict(initial=True) + ), ) - cached = _initial_snapshot_cache[root_cls] = ( - n_state_classes, - snapshot, - dict(zip(sorted(snapshot), state_snapshot_hashes(snapshot), strict=True)), - ) - _, defaults, default_hashes = cached - if len(hashes) != len(default_hashes): + cached = _initial_snapshot_cache[root_cls] + if len(hashes) != len(cached.hashes): # The frontend was compiled against a different set of states. return delta - frontend_hashes = dict(zip(sorted(default_hashes), hashes, strict=True)) + frontend_hashes = dict(zip(sorted(cached.hashes), hashes, strict=True)) diff: Delta = {} for state_name, state_vars in delta.items(): - default_vars = defaults.get(state_name) - if default_vars is None or frontend_hashes.get( + default_vars = cached.serialized.get(state_name) + if default_vars is None or frontend_hashes.get(state_name) != cached.hashes.get( state_name - ) != default_hashes.get(state_name): + ): diff[state_name] = state_vars continue changed = { name: value for name, value in state_vars.items() - if name not in default_vars or value != default_vars[name] + if name not in default_vars or _serialize_var(value) != default_vars[name] } if changed: diff[state_name] = changed 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 336ae12765a..d459359a5fc 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 @@ -840,13 +840,19 @@ async def test_hydrate_and_load_diffs_against_compiled_defaults( class CookieState(State): flavor: str = rx.Cookie("plain") loads: int = 0 + ratio: float = 1.0 + + @event + def set_ratio_int(self): + self.ratio = 1 # Python-equal to the default, JSON-distinct. wired_app.add_page(lambda: rx.text(CookieState.flavor), route="/") wired_app._compile_page("index") boot_name = Event.from_event_type(State.hydrate_and_load())[0].name # pyright: ignore[reportCallIssue] cookie_key = f"{CookieState.get_full_name()}.flavor{FIELD_MARKER}" state_name = State.get_full_name() - hashes = state_snapshot_hashes(compile_state(State)) + compiled = compile_state(State) + hashes = state_snapshot_hashes(compiled) async with real_base_state_processor as processor: future = await processor.enqueue( @@ -865,6 +871,22 @@ class CookieState(State): "flavor" + FIELD_MARKER: "chocolate" } + # A value that is Python-equal but serializes differently is still sent. + emitted_deltas.clear() + async with real_base_state_processor as processor: + await ( + await processor.enqueue( + token, Event.from_event_type(CookieState.set_ratio_int())[0] + ) + ).wait_all() + emitted_deltas.clear() + future = await processor.enqueue( + token, _boot_event(boot_name, {"hashes": hashes}) + ) + await future.wait_all() + snapshot = emitted_deltas[0][1] + assert snapshot[CookieState.get_full_name()]["ratio" + FIELD_MARKER] == 1 + # Hashes compiled against a different set of states fall back to the full snapshot. emitted_deltas.clear() async with real_base_state_processor as processor: diff --git a/tests/units/test_app.py b/tests/units/test_app.py index 44d1e43f907..3770de9bf04 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -4373,3 +4373,27 @@ async def test_on_connect_processes_boot_event_from_auth( await event_namespace.on_connect("new_sid", {"QUERY_STRING": "token=abc"}, None) await event_namespace.on_connect("new_sid", {"QUERY_STRING": "token=abc"}) event_namespace.on_event.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_on_connect_unlinks_token_when_boot_event_fails( + event_namespace: EventNamespace, +): + """A boot event that fails to process drops the sid/token link before refusing the connect. + + Args: + event_namespace: The event namespace. + """ + event_namespace._token_manager = Mock() + event_namespace._token_manager.link_token_to_sid = AsyncMock(return_value=None) + event_namespace._token_manager.disconnect_token = AsyncMock() + event_namespace._token_manager.sid_to_token = {"new_sid": "abc"} + event_namespace.on_event = AsyncMock(side_effect=ValueError("bad boot event")) + + with pytest.raises(ValueError, match="bad boot event"): + await event_namespace.on_connect( + "new_sid", {"QUERY_STRING": "token=abc"}, {"event": {"name": "x"}} + ) + event_namespace._token_manager.disconnect_token.assert_awaited_once_with( + "abc", "new_sid" + ) From 2ff4e515494e4740236e4f7922ff9a25ce7594b3 Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 20:59:30 +0000 Subject: [PATCH 07/14] Bind hydrate hashes to state names, share the on_load supersede group, shorten news - The hash list the frontend sends now starts with a digest of the compiled state names, so hashes are never matched positionally against a different set of states. - @event(supersedes=...) accepts a group name; hydrate_and_load and on_load_internal share the on_load group so a reconnect cancels an unfinished navigation chain the way a navigation always did. - News fragments describe the user-visible change only. --- news/+first-load-hydrate.performance.md | 2 +- .../news/+first-load-hydrate.performance.md | 2 +- .../src/reflex_base/constants/compiler.py | 3 ++ .../src/reflex_base/event/__init__.py | 33 +++++++++--- .../event/processor/event_processor.py | 7 +-- reflex/state.py | 54 ++++++++++++------- .../processor/test_base_state_processor.py | 5 +- tests/units/test_state.py | 6 +++ 8 files changed, 79 insertions(+), 33 deletions(-) diff --git a/news/+first-load-hydrate.performance.md b/news/+first-load-hydrate.performance.md index 0199db48f6e..8b7c9f82f29 100644 --- a/news/+first-load-hydrate.performance.md +++ b/news/+first-load-hydrate.performance.md @@ -1 +1 @@ -Faster first page load: the frontend now sends a single `hydrate_and_load` event inside the websocket connect packet instead of three events after the connect acknowledgement, the backend handles it under one state lock, and on the first hydrate of a page only the vars that differ from the compiled defaults are sent (a 20-substate app went from 36 KB to 2 KB on the wire and from 5 to 3 frames). Redis-backed apps also do far fewer commands per page load: the lock is verified once per state-tree save instead of once per substate, and neither the websocket connect nor the final `is_hydrated` flip load every substate any more. +Faster first page load: the frontend hydrates in a single event sent with the websocket connect, and only the values that differ from the compiled defaults are sent (a 20-substate app went from 36 KB to 2 KB and from 5 to 3 frames). Redis-backed apps also do about 80% fewer redis commands per page load. diff --git a/packages/reflex-base/news/+first-load-hydrate.performance.md b/packages/reflex-base/news/+first-load-hydrate.performance.md index 1b49d818e03..ac601cae98f 100644 --- a/packages/reflex-base/news/+first-load-hydrate.performance.md +++ b/packages/reflex-base/news/+first-load-hydrate.performance.md @@ -1 +1 @@ -The compiled frontend hydrates with one `hydrate_and_load` event carried in the socket.io connect packet, sending per-state hashes of its compiled `initialState` so the backend can skip vars still at their default; client-side navigation keeps using `on_load_internal` as before. +The compiled frontend hydrates with a single `hydrate_and_load` event sent along with the websocket connect, so the first page load needs one fewer round trip. diff --git a/packages/reflex-base/src/reflex_base/constants/compiler.py b/packages/reflex-base/src/reflex_base/constants/compiler.py index 63106adb7ff..e3c8da50a9e 100644 --- a/packages/reflex-base/src/reflex_base/constants/compiler.py +++ b/packages/reflex-base/src/reflex_base/constants/compiler.py @@ -59,6 +59,9 @@ class CompileVars(SimpleNamespace): HYDRATE = "hydrate" # The name of the event sent on (re)connect: hydrate plus on_load in one step. HYDRATE_AND_LOAD = "hydrate_and_load" + # The supersede group shared by hydrate_and_load and on_load_internal, so a + # reconnect or navigation cancels the previous unfinished on_load chain. + ON_LOAD_SUPERSEDE_GROUP = "on_load" # The name of the is_hydrated variable. IS_HYDRATED = "is_hydrated" # The name of the function to add events to the queue. diff --git a/packages/reflex-base/src/reflex_base/event/__init__.py b/packages/reflex-base/src/reflex_base/event/__init__.py index b191bd93349..7b6250ba17a 100644 --- a/packages/reflex-base/src/reflex_base/event/__init__.py +++ b/packages/reflex-base/src/reflex_base/event/__init__.py @@ -595,7 +595,22 @@ def supersedes(self) -> bool: Returns: True if the event handler is marked as superseding. """ - return getattr(self.fn, SUPERSEDES_MARKER, False) + return bool(getattr(self.fn, SUPERSEDES_MARKER, False)) + + @property + def supersede_group(self) -> str | None: + """The name under which this handler's chains supersede each other. + + Handlers sharing a group cancel each other's unfinished chains, so a + reconnect's hydrate and a navigation's on_load never run side by side. + + Returns: + The group name, or None for handlers that do not supersede. + """ + marker = getattr(self.fn, SUPERSEDES_MARKER, False) + if isinstance(marker, str): + return marker + return format.format_event_handler(self) if marker else None def __call__(self, *args: Any, **kwargs: Any) -> "EventSpec": """Pass arguments to the handler to get an event spec. @@ -2956,7 +2971,7 @@ def __new__( func: None = None, *, background: bool | None = None, - supersedes: bool | None = None, + supersedes: bool | str | None = None, stop_propagation: bool | None = None, prevent_default: bool | None = None, throttle: int | None = None, @@ -2972,7 +2987,7 @@ def __new__( func: "Callable[[BASE_STATE, Unpack[P]], Any]", *, background: bool | None = None, - supersedes: bool | None = None, + supersedes: bool | str | None = None, stop_propagation: bool | None = None, prevent_default: bool | None = None, throttle: int | None = None, @@ -2985,7 +3000,7 @@ def __new__( func: "Callable[[BASE_STATE, Unpack[P]], Any] | None" = None, *, background: bool | None = None, - supersedes: bool | None = None, + supersedes: bool | str | None = None, stop_propagation: bool | None = None, prevent_default: bool | None = None, throttle: int | None = None, @@ -2999,8 +3014,10 @@ def __new__( background: Whether the event should be run in the background. Defaults to False. supersedes: Whether enqueuing the event cancels the previous unfinished chain of the same event for the same client token (latest-wins). - Cancellation is cooperative, so a handler that never yields to the - event loop is not interrupted. Defaults to False. + A string names a group instead: handlers sharing the group + supersede each other's chains. Cancellation is cooperative, so + a handler that never yields to the event loop is not + interrupted. Defaults to False. stop_propagation: Whether to stop the event from bubbling up the DOM tree. prevent_default: Whether to prevent the default behavior of the event. throttle: Throttle the event handler to limit calls (in milliseconds). @@ -3052,8 +3069,8 @@ def wrapper( msg = "Background task must be async function or generator." raise TypeError(msg) setattr(func, BACKGROUND_TASK_MARKER, True) - if supersedes is True: - setattr(func, SUPERSEDES_MARKER, True) + if supersedes: + setattr(func, SUPERSEDES_MARKER, supersedes) if getattr(func, "__name__", "").startswith("_"): msg = "Event handlers cannot be private." raise ValueError(msg) 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 96408c6bf84..bcd7f23ada9 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 @@ -550,7 +550,8 @@ def _supersede_previous( 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. + previous unfinished event chain of the same supersede group (the + handler itself, or the group it names) for the same client token. Args: token: The client token associated with the event. @@ -561,9 +562,9 @@ def _supersede_previous( registered = RegistrationContext.get().event_handlers.get(event.name) except LookupError: return - if registered is None or not registered.handler.supersedes: + if registered is None or (group := registered.handler.supersede_group) is None: return - key = (event.name, token) + key = (group, token) previous = self._superseded.get(key) if previous is not None and not previous.all_done(): logger.debug( diff --git a/reflex/state.py b/reflex/state.py index ecc6b66f7d6..eeffe277cb3 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -2475,7 +2475,7 @@ def set_is_hydrated(self, value: bool) -> None: """ self.is_hydrated = value - @event(supersedes=True) + @event(supersedes=constants.CompileVars.ON_LOAD_SUPERSEDE_GROUP) async def hydrate_and_load( self, vars: dict[str, Any] | None = None, @@ -2522,8 +2522,20 @@ async def hydrate_and_load( T = TypeVar("T", bound=BaseState) +def _short_digest(text: str) -> str: + """Digest text into a short hex string. + + Args: + text: The text to digest. + + Returns: + The first 16 hex digits of its SHA-1. + """ + return hashlib.sha1(text.encode()).hexdigest()[:16] + + def state_snapshot_hashes(snapshot: Delta) -> list[str]: - """Hash each state's entry of a full-tree snapshot as the frontend receives it. + """Hash a full-tree snapshot as the frontend receives it. Used at compile time for the ``initialState`` baked into the frontend and at runtime for the backend's own default snapshot, so equal hashes mean the @@ -2533,14 +2545,17 @@ def state_snapshot_hashes(snapshot: Delta) -> list[str]: snapshot: A resolved full-tree dict, as returned by ``BaseState.dict``. Returns: - A short hex digest of each state's serialized vars, in sorted state - name order. + A digest of the sorted state names, followed by a digest of each + state's serialized vars in that order; the first entry binds the + rest to the state names they were computed for. """ + names = sorted(snapshot) return [ - hashlib.sha1( - format.json_dumps(snapshot[state_name], sort_keys=True).encode() - ).hexdigest()[:16] - for state_name in sorted(snapshot) + _short_digest("\n".join(names)), + *( + _short_digest(format.json_dumps(snapshot[state_name], sort_keys=True)) + for state_name in names + ), ] @@ -2552,6 +2567,8 @@ class _InitialSnapshot: n_state_classes: int # Per state full name, each var's serialized default value. serialized: dict[str, dict[str, str]] + # The digest of the sorted state names. + names_digest: str # Per state full name, the hash of its serialized defaults. hashes: dict[str, str] @@ -2583,15 +2600,15 @@ def cache_initial_snapshot(root_cls: type[BaseState], snapshot: Delta) -> None: root_cls: The root state class the snapshot was taken from. snapshot: The resolved full-tree default snapshot. """ + names_digest, *state_hashes = state_snapshot_hashes(snapshot) _initial_snapshot_cache[root_cls] = _InitialSnapshot( n_state_classes=len(all_base_state_classes), serialized={ state_name: {name: _serialize_var(value) for name, value in vars.items()} for state_name, vars in snapshot.items() }, - hashes=dict( - zip(sorted(snapshot), state_snapshot_hashes(snapshot), strict=True) - ), + names_digest=names_digest, + hashes=dict(zip(sorted(snapshot), state_hashes, strict=True)), ) @@ -2607,8 +2624,9 @@ async def _diff_against_initial_state( Args: root_cls: The root state class; its default snapshot is computed once. delta: The resolved full snapshot about to be sent. - hashes: Per-state hashes of the frontend's compiled ``initialState``, - in sorted state name order. + hashes: The digest of the frontend's compiled state names followed + by its per-state hashes of the compiled ``initialState``, in + sorted state name order. Returns: The delta with unchanged vars removed for every state whose compiled @@ -2625,10 +2643,10 @@ async def _diff_against_initial_state( ), ) cached = _initial_snapshot_cache[root_cls] - if len(hashes) != len(cached.hashes): + if not hashes or hashes[0] != cached.names_digest: # The frontend was compiled against a different set of states. return delta - frontend_hashes = dict(zip(sorted(cached.hashes), hashes, strict=True)) + frontend_hashes = dict(zip(sorted(cached.hashes), hashes[1:], strict=True)) diff: Delta = {} for state_name, state_vars in delta.items(): default_vars = cached.serialized.get(state_name) @@ -2813,9 +2831,9 @@ class OnLoadInternalState(State): This is a separate substate to avoid deserializing the entire state tree for every page navigation. """ - # A newer navigation supersedes the previous unfinished on_load chain for - # the same client token, cancelling its stale work (#6593). - @event(supersedes=True) + # A newer navigation or reconnect supersedes the previous unfinished + # on_load chain for the same client token, cancelling its stale work (#6593). + @event(supersedes=constants.CompileVars.ON_LOAD_SUPERSEDE_GROUP) def on_load_internal(self) -> list[Event | EventSpec | event.EventCallback] | None: """Queue on_load handlers for the current page. 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 d459359a5fc..52ee00fe864 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 @@ -887,11 +887,12 @@ def set_ratio_int(self): snapshot = emitted_deltas[0][1] assert snapshot[CookieState.get_full_name()]["ratio" + FIELD_MARKER] == 1 - # Hashes compiled against a different set of states fall back to the full snapshot. + # Hashes compiled against a different set of states (a different names + # digest) fall back to the full snapshot. emitted_deltas.clear() async with real_base_state_processor as processor: future = await processor.enqueue( - token, _boot_event(boot_name, {"hashes": hashes[:-1]}) + token, _boot_event(boot_name, {"hashes": ["0" * 16, *hashes[1:]]}) ) await future.wait_all() snapshot = emitted_deltas[0][1] diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 824848da4ac..08cda78c1e4 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -5235,6 +5235,12 @@ async def test_on_load_internal_supersedes_previous_navigation( """ assert OnLoadInternalState.event_handlers["on_load_internal"].supersedes assert not State.event_handlers["hydrate"].supersedes + # A reconnect's hydrate and a navigation's on_load cancel each other's chains. + assert ( + OnLoadInternalState.event_handlers["on_load_internal"].supersede_group + == State.event_handlers["hydrate_and_load"].supersede_group + == "on_load" + ) app = app_module_mock.app = App(_state=State) app._state_manager = mock_root_event_context.state_manager From 48f9c5cdeef7459a8e620850d22ff3660fefa087 Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 21:03:51 +0000 Subject: [PATCH 08/14] Fence state saves on the redis lock with WATCH/MULTI/EXEC The pipelined write of a state tree now runs as a transaction that watches the lock key: if the lock expired or changed hands between the ownership check and the write, EXEC aborts and the save raises LockExpiredError instead of overwriting a newer writer's state. The mocked redis models WATCH/MULTI/EXEC so the unit tests cover the aborted-save path. --- reflex/istate/manager/redis.py | 98 +++++++++++++----------- tests/units/istate/manager/test_redis.py | 51 ++++++++---- tests/units/mock_redis.py | 81 +++++++++++++++----- 3 files changed, 150 insertions(+), 80 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index dc9d599ce9e..f042bc1494c 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -14,6 +14,7 @@ from redis import ResponseError from redis.asyncio import Redis +from redis.exceptions import WatchError from reflex_base.config import get_config from reflex_base.environment import environment from reflex_base.utils.exceptions import ( @@ -406,53 +407,64 @@ async def set_state( pickle_state = token.serialize(state) writes = [(str(token), pickle_state)] if pickle_state else [] - # Check that we're holding the lock. - if ( - lock_id is not None - and (existing_lock_id := await self.redis.get(self._lock_key(token))) - != lock_id - ): - msg = ( - f"Lock expired for token {token} while processing. Consider increasing " - f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " - "or use `@rx.event(background=True)` decorator for long-running tasks. " - f"Current lock id: {existing_lock_id!r}, expected lock id: {lock_id!r}." - + ( - f" Happened in event: {event.name}" - if (event := context.get("event")) is not None - else "" - ) - ) - raise LockExpiredError(msg) + if lock_id is None: + pipeline = self.redis.pipeline(transaction=False) + for key, pickle_state in writes: + pipeline.set(key, pickle_state, ex=self.token_expiration) + if writes: + await pipeline.execute() + return - if ( - isinstance(token, BaseStateToken) - and lock_id is not None - and token.lock_key not in self._local_leases - ): - time_taken = ( - self.lock_expiration - (await self.redis.pttl(self._lock_key(token))) - ) / 1000 - if time_taken > self.lock_warning_threshold / 1000: - event_suffix = ( - f" Happened in event: {event.name}" - if (event := context.get("event")) is not None - else "" - ) - logger.warning( - f"Lock for token {token} was held too long {time_taken=}s, " - "use `@rx.event(background=True)` decorator for long-running " - f"tasks.{event_suffix}", - extra={"dedupe": True}, + event_suffix = ( + f" Happened in event: {event.name}" + if (event := context.get("event")) is not None + else "" + ) + lock_key = self._lock_key(token) + # Fence the write on the lock: WATCH makes EXEC fail if the lock key + # changed hands or expired between the ownership check and the write. + async with self.redis.pipeline(transaction=True) as pipeline: + await pipeline.watch(lock_key) + existing_lock_id = await pipeline.get(lock_key) + if existing_lock_id != lock_id: + msg = ( + f"Lock expired for token {token} while processing. Consider increasing " + f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " + "or use `@rx.event(background=True)` decorator for long-running tasks. " + f"Current lock id: {existing_lock_id!r}, expected lock id: {lock_id!r}." + + event_suffix ) - - if len(writes) == 1: - await self.redis.set(writes[0][0], writes[0][1], ex=self.token_expiration) - elif writes: - pipeline = self.redis.pipeline() + raise LockExpiredError(msg) + if ( + isinstance(token, BaseStateToken) + and token.lock_key not in self._local_leases + ): + time_taken = ( + self.lock_expiration - (await self.redis.pttl(lock_key)) + ) / 1000 + if time_taken > self.lock_warning_threshold / 1000: + logger.warning( + f"Lock for token {token} was held too long {time_taken=}s, " + "use `@rx.event(background=True)` decorator for long-running " + f"tasks.{event_suffix}", + extra={"dedupe": True}, + ) + if not writes: + return + pipeline.multi() for key, pickle_state in writes: pipeline.set(key, pickle_state, ex=self.token_expiration) - await pipeline.execute() + try: + await pipeline.execute() + except WatchError: + msg = ( + f"Lock for token {token} changed while its state was being " + "saved, so the save was discarded. Consider increasing " + f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " + "or use `@rx.event(background=True)` decorator for long-running tasks." + + event_suffix + ) + raise LockExpiredError(msg) from None def _collect_state_writes( self, token: BaseStateToken, base_state: BaseState diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index d7498b36382..d9040fe68d6 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -742,7 +742,7 @@ async def test_set_state_checks_lock_once_per_tree( state_manager_redis: StateManagerRedis, root_state: type[RedisTestState], ): - """Saving a state tree verifies the lock and reads its TTL once, not per substate. + """Saving a state tree reads the lock TTL once, not once per substate. Args: state_manager_redis: The StateManagerRedis to test. @@ -751,16 +751,9 @@ async def test_set_state_checks_lock_once_per_tree( state_manager_redis._oplock_enabled = False token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) redis = state_manager_redis.redis - lock_key = state_manager_redis._lock_key(token) - real_get, real_pttl = redis.get, redis.pttl - lock_gets: list[Any] = [] + real_pttl = redis.pttl pttls: list[Any] = [] - async def counting_get(key): - if key == lock_key: - lock_gets.append(key) - return await real_get(key) - async def counting_pttl(key): pttls.append(key) return await real_pttl(key) @@ -768,18 +761,44 @@ async def counting_pttl(key): async with state_manager_redis.modify_state(token) as state: assert len(state.substates) == 2 state.count = 1 - lock_id = await real_get(lock_key) - redis.get = counting_get # pyright: ignore[reportAttributeAccessIssue] redis.pttl = counting_pttl # pyright: ignore[reportAttributeAccessIssue] try: - await state_manager_redis.set_state(token, state, lock_id=lock_id) + await state_manager_redis.set_state( + token, + state, + lock_id=await redis.get(state_manager_redis._lock_key(token)), + ) finally: - redis.get = real_get - redis.pttl = real_pttl + redis.pttl = real_pttl # pyright: ignore[reportAttributeAccessIssue] - # One lock check and one TTL read for a tree of three states. - assert len(lock_gets) == 1 + # One TTL read for a tree of three states. assert len(pttls) == 1 saved = await state_manager_redis.get_state(token) assert isinstance(saved, root_state) assert saved.count == 1 + + +async def test_set_state_discards_writes_when_lock_changes_hands( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """A save whose lock expired or was re-acquired between the check and the write is discarded. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + """ + from reflex_base.utils.exceptions import LockExpiredError + + state_manager_redis._oplock_enabled = False + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + + with pytest.raises(LockExpiredError): + async with state_manager_redis.modify_state(token) as state: + state.count = 5 + # Another worker takes over the lock before this save lands. + await state_manager_redis.redis.set(lock_key, b"someone-else") + saved = await state_manager_redis.get_state(token) + assert isinstance(saved, root_state) + assert saved.count == 0 diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 4eb0255ac41..31b6ec3d1a5 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, Mock from redis.asyncio import Redis +from redis.exceptions import WatchError from redis.typing import EncodableT, KeyT from reflex.utils import prerequisites @@ -142,39 +143,77 @@ async def mock_pexpire(key: KeyT, px: int, xx: bool = False) -> bool: # noqa: R return True return False - def pipeline(): - pipeline_mock = Mock() - results = [] + class _Pipeline: + """A pipeline that also models WATCH/MULTI/EXEC on the mocked keys.""" - def get_pipeline(key: KeyT): - results.append(redis_mock.get(key=key)) + def __init__(self): + self.results = [] + self.watched: bytes | None = None + self.watched_value: Any = None + self.in_multi = False - def set_pipeline( + async def __aenter__(self): + return self + + async def __aexit__(self, *exc_info): + await self.reset() + + async def reset(self): + self.results = [] + self.watched = None + self.watched_value = None + self.in_multi = False + + async def watch(self, key: KeyT): + _expire_keys() + self.watched = _key_bytes(key) + self.watched_value = keys.get(self.watched) + + def multi(self): + self.in_multi = True + + def get(self, key: KeyT): + if self.watched is not None and not self.in_multi: + # Immediate mode while watching, before MULTI. + return redis_mock.get(key=key) + self.results.append(redis_mock.get(key=key)) + return None + + def set( + self, key: KeyT, value: EncodableT, ex: int | None = None, px: int | None = None, nx: bool = False, ): - results.append(redis_mock.set(key=key, value=value, ex=ex, px=px, nx=nx)) + self.results.append( + redis_mock.set(key=key, value=value, ex=ex, px=px, nx=nx) + ) - def sadd_pipeline(key: KeyT, value: EncodableT): - results.append(redis_mock.sadd(key=key, value=value)) + def sadd(self, key: KeyT, value: EncodableT): + self.results.append(redis_mock.sadd(key=key, value=value)) - def pexpire_pipeline(key: KeyT, px: int, xx: bool = False): - results.append(redis_mock.pexpire(key=key, px=px, xx=xx)) + def pexpire(self, key: KeyT, px: int, xx: bool = False): + self.results.append(redis_mock.pexpire(key=key, px=px, xx=xx)) - async def execute(): + async def execute(self): _expire_keys() - return await asyncio.gather(*results) - - pipeline_mock.get = get_pipeline - pipeline_mock.set = set_pipeline - pipeline_mock.sadd = sadd_pipeline - pipeline_mock.pexpire = pexpire_pipeline - pipeline_mock.execute = execute - - return pipeline_mock + if ( + self.watched is not None + and keys.get(self.watched) != self.watched_value + ): + for pending in self.results: + pending.close() + self.results = [] + msg = "Watched variable changed." + raise WatchError(msg) + results = await asyncio.gather(*self.results) + self.results = [] + return results + + def pipeline(transaction: bool = True): + return _Pipeline() async def pttl(key: KeyT) -> int: # noqa: RUF029 _expire_keys() From e923225dd31bbbfd6e237037d9d1a8f536d33667 Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 21:11:35 +0000 Subject: [PATCH 09/14] Read the lock TTL on the watching pipeline and name the boot protocol keys The TTL warning in the redis save now reads PTTL through the pipeline that holds WATCH, so the ownership check, the warning and the write use one connection instead of borrowing a second pooled connection while the first is held. The connect-packet and hydrate payload keys are named constants on CompileVars and used by the backend and the generated frontend template. --- .../reflex_base/.templates/web/utils/state.js | 3 ++- .../src/reflex_base/compiler/templates.py | 6 +++--- .../src/reflex_base/constants/compiler.py | 6 ++++++ reflex/app.py | 6 +++++- reflex/istate/manager/redis.py | 4 +++- tests/units/istate/manager/test_redis.py | 19 +++++++++++++------ tests/units/mock_redis.py | 6 ++++++ 7 files changed, 38 insertions(+), 12 deletions(-) diff --git a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js index d97b82675f6..c6e88096376 100644 --- a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js +++ b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js @@ -599,7 +599,8 @@ export const connect = async ( // The hydrate event rides in the socket.io CONNECT packet, so the backend // starts loading state as soon as the namespace connects instead of after - // an extra round trip for the connect acknowledgement. + // an extra round trip for the connect acknowledgement. The key is read by + // the backend as CompileVars.CONNECT_AUTH_EVENT. const bootAuth = (first) => ({ event: withRouterData(initialEvents(first)[0], params), }); diff --git a/packages/reflex-base/src/reflex_base/compiler/templates.py b/packages/reflex-base/src/reflex_base/compiler/templates.py index 568e98f48e6..c16b82223f6 100644 --- a/packages/reflex-base/src/reflex_base/compiler/templates.py +++ b/packages/reflex-base/src/reflex_base/compiler/templates.py @@ -338,7 +338,7 @@ def context_template( internal_events.push( ReflexEvent( '{state_name}.{constants.CompileVars.UPDATE_VARS_INTERNAL}', - {{vars: client_storage_vars}}, + {{{constants.CompileVars.PAYLOAD_VARS}: client_storage_vars}}, ), ); }} @@ -359,10 +359,10 @@ def context_template( const client_storage_vars = clientStorageVars(); const payload = {{}}; if (client_storage_vars !== undefined) {{ - payload.vars = client_storage_vars; + payload["{constants.CompileVars.PAYLOAD_VARS}"] = client_storage_vars; }} if (first) {{ - payload.hashes = initialStateHashes; + payload["{constants.CompileVars.PAYLOAD_HASHES}"] = initialStateHashes; }} return [ReflexEvent('{state_name}.{constants.CompileVars.HYDRATE_AND_LOAD}', payload)]; }} diff --git a/packages/reflex-base/src/reflex_base/constants/compiler.py b/packages/reflex-base/src/reflex_base/constants/compiler.py index e3c8da50a9e..977451b7456 100644 --- a/packages/reflex-base/src/reflex_base/constants/compiler.py +++ b/packages/reflex-base/src/reflex_base/constants/compiler.py @@ -62,6 +62,12 @@ class CompileVars(SimpleNamespace): # The supersede group shared by hydrate_and_load and on_load_internal, so a # reconnect or navigation cancels the previous unfinished on_load chain. ON_LOAD_SUPERSEDE_GROUP = "on_load" + # The key of the socket.io CONNECT auth packet that carries the boot event. + CONNECT_AUTH_EVENT = "event" + # Payload keys of hydrate_and_load / update_vars_internal; they are passed + # through as handler kwargs, so they must match those parameter names. + PAYLOAD_VARS = "vars" + PAYLOAD_HASHES = "hashes" # The name of the is_hydrated variable. IS_HYDRATED = "is_hydrated" # The name of the function to add events to the queue. diff --git a/reflex/app.py b/reflex/app.py index 5ecbf21f010..6694651ddd2 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -2028,7 +2028,11 @@ async def on_connect(self, sid: str, environ: dict, auth: Any = None): f"Frontend version {subprotocol} for session {sid} does not match the backend version {constants.Reflex.VERSION}." ) - if isinstance(auth, dict) and (boot_event := auth.get("event")) is not None: + if ( + isinstance(auth, dict) + and (boot_event := auth.get(constants.CompileVars.CONNECT_AUTH_EVENT)) + is not None + ): try: await self.on_event(sid, boot_event) except Exception: diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index f042bc1494c..d2475b80f8c 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -439,8 +439,10 @@ async def set_state( isinstance(token, BaseStateToken) and token.lock_key not in self._local_leases ): + # Immediate mode on the watching connection: the parent client + # would need a second pooled connection while this one is held. time_taken = ( - self.lock_expiration - (await self.redis.pttl(lock_key)) + self.lock_expiration - (await pipeline.pttl(lock_key)) ) / 1000 if time_taken > self.lock_warning_threshold / 1000: logger.warning( diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index d9040fe68d6..93644500a33 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -751,17 +751,24 @@ async def test_set_state_checks_lock_once_per_tree( state_manager_redis._oplock_enabled = False token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) redis = state_manager_redis.redis - real_pttl = redis.pttl + real_pipeline = redis.pipeline pttls: list[Any] = [] - async def counting_pttl(key): - pttls.append(key) - return await real_pttl(key) + def counting_pipeline(*args, **kwargs): + pipe = real_pipeline(*args, **kwargs) + real_pttl = pipe.pttl + + def counting_pttl(key): + pttls.append(key) + return real_pttl(key) + + pipe.pttl = counting_pttl # pyright: ignore[reportAttributeAccessIssue] + return pipe async with state_manager_redis.modify_state(token) as state: assert len(state.substates) == 2 state.count = 1 - redis.pttl = counting_pttl # pyright: ignore[reportAttributeAccessIssue] + redis.pipeline = counting_pipeline # pyright: ignore[reportAttributeAccessIssue] try: await state_manager_redis.set_state( token, @@ -769,7 +776,7 @@ async def counting_pttl(key): lock_id=await redis.get(state_manager_redis._lock_key(token)), ) finally: - redis.pttl = real_pttl # pyright: ignore[reportAttributeAccessIssue] + redis.pipeline = real_pipeline # pyright: ignore[reportAttributeAccessIssue] # One TTL read for a tree of three states. assert len(pttls) == 1 diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 31b6ec3d1a5..9ad1ba0ea99 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -179,6 +179,12 @@ def get(self, key: KeyT): self.results.append(redis_mock.get(key=key)) return None + def pttl(self, key: KeyT): + if self.watched is not None and not self.in_multi: + return redis_mock.pttl(key=key) + self.results.append(redis_mock.pttl(key=key)) + return None + def set( self, key: KeyT, From b7295cc4638f2c245615cb4716703ba177cecb1d Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 21:25:21 +0000 Subject: [PATCH 10/14] Retry the fenced redis save once when the connection drops while watching redis-py reports a connection error during WATCH as a WatchError, since the watch state is lost with the connection, and a lease refresh touching the lock key trips the watch as well. The fenced save now retries once on a fresh connection, which re-checks the lock before writing, and only a second WatchError discards the save as a lost lock. This surfaced in the lifespan integration test: the harness closes the redis pool on shutdown while a lifespan task's save is in flight, and the old unfenced path reconnected silently where the fenced one raised. --- reflex/istate/manager/redis.py | 55 ++++++++++++++++++------ tests/units/istate/manager/test_redis.py | 51 ++++++++++++++++++++++ 2 files changed, 93 insertions(+), 13 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index d2475b80f8c..c81f3444b62 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -420,9 +420,48 @@ async def set_state( if (event := context.get("event")) is not None else "" ) + # A dropped connection also surfaces as WatchError, since it loses the + # WATCH state, and so does a lease refresh touching the lock key: retry + # once on a fresh connection, which re-checks the lock before writing. + try: + await self._fenced_save(token, lock_id, writes, event_suffix) + except WatchError: + try: + await self._fenced_save(token, lock_id, writes, event_suffix) + except WatchError: + msg = ( + f"Lock for token {token} changed while its state was being " + "saved, so the save was discarded. Consider increasing " + f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " + "or use `@rx.event(background=True)` decorator for long-running tasks." + + event_suffix + ) + raise LockExpiredError(msg) from None + + async def _fenced_save( + self, + token: StateToken[Any], + lock_id: bytes, + writes: list[tuple[str, bytes]], + event_suffix: str, + ) -> None: + """Write the serialized states in one transaction fenced on the lock. + + WATCH makes EXEC fail if the lock key changed hands or expired between + the ownership check and the write. + + Args: + token: The token whose lock fences the write. + lock_id: The lock id the caller holds. + writes: The redis keys and pickled payloads to write. + event_suffix: Event context appended to warning and error messages. + + Raises: + LockExpiredError: If the lock is not held by lock_id. + WatchError: If the lock key changed or the connection dropped + while watching it. + """ lock_key = self._lock_key(token) - # Fence the write on the lock: WATCH makes EXEC fail if the lock key - # changed hands or expired between the ownership check and the write. async with self.redis.pipeline(transaction=True) as pipeline: await pipeline.watch(lock_key) existing_lock_id = await pipeline.get(lock_key) @@ -456,17 +495,7 @@ async def set_state( pipeline.multi() for key, pickle_state in writes: pipeline.set(key, pickle_state, ex=self.token_expiration) - try: - await pipeline.execute() - except WatchError: - msg = ( - f"Lock for token {token} changed while its state was being " - "saved, so the save was discarded. Consider increasing " - f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " - "or use `@rx.event(background=True)` decorator for long-running tasks." - + event_suffix - ) - raise LockExpiredError(msg) from None + await pipeline.execute() def _collect_state_writes( self, token: BaseStateToken, base_state: BaseState diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 93644500a33..67781940dc9 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -9,6 +9,7 @@ import pytest import pytest_asyncio +from redis.exceptions import WatchError from reflex.istate.manager.redis import StateManagerRedis from reflex.istate.manager.token import BaseStateToken @@ -785,6 +786,56 @@ def counting_pttl(key): assert saved.count == 1 +async def test_set_state_retries_once_when_connection_drops_while_watching( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """A dropped connection during the fenced save is retried on a fresh one. + + redis-py reports a connection error while watching as WatchError, so the + save must re-check the lock on a new connection instead of failing. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + """ + state_manager_redis._oplock_enabled = False + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + redis = state_manager_redis.redis + real_pipeline = redis.pipeline + pipelines: list[Any] = [] + + def dropping_pipeline(*args, **kwargs): + pipe = real_pipeline(*args, **kwargs) + pipelines.append(pipe) + if len(pipelines) == 1: + + async def dropped_execute(): + await pipe.reset() + msg = "A ConnectionError occurred while watching one or more keys" + raise WatchError(msg) + + pipe.execute = dropped_execute # pyright: ignore[reportAttributeAccessIssue] + return pipe + + async with state_manager_redis.modify_state(token) as state: + state.count = 7 + redis.pipeline = dropping_pipeline # pyright: ignore[reportAttributeAccessIssue] + try: + await state_manager_redis.set_state( + token, + state, + lock_id=await redis.get(state_manager_redis._lock_key(token)), + ) + finally: + redis.pipeline = real_pipeline # pyright: ignore[reportAttributeAccessIssue] + + assert len(pipelines) == 2 + saved = await state_manager_redis.get_state(token) + assert isinstance(saved, root_state) + assert saved.count == 7 + + async def test_set_state_discards_writes_when_lock_changes_hands( state_manager_redis: StateManagerRedis, root_state: type[RedisTestState], From e3975ac661d37c567b6ab9634a3dc1d57ebaf387 Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 21:32:32 +0000 Subject: [PATCH 11/14] Save state trees with one atomic lock-checked script instead of WATCH The redis client is created with retry_on_error=[RedisError], so when EXEC returned nil for a changed lock, redis-py reconnected and re-ran the transaction without the WATCH and the stale write landed anyway; the mock did not model that, only a real redis showed it. The save is now a single EVAL: the script compares the lock id, writes every touched state with its expiration, and returns the lock's PTTL, or nil without writing when the lock changed hands. That makes the check and write atomic on the server, safe under command retries, and one round trip instead of four (WATCH, GET, PTTL, EXEC). The mock emulates that script, and the WATCH-specific retry and tests go away with it. Also fall back to the full snapshot when a hydrate payload's hash list has a different length than the compiled states, instead of failing the event. --- reflex/istate/manager/redis.py | 131 +++++++----------- reflex/state.py | 5 +- tests/units/istate/manager/test_redis.py | 83 ++--------- tests/units/mock_redis.py | 76 ++++------ .../processor/test_base_state_processor.py | 11 ++ 5 files changed, 107 insertions(+), 199 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index c81f3444b62..cf45e961379 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -9,12 +9,11 @@ import sys import time import uuid -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable from typing import Any, TypedDict, cast from redis import ResponseError from redis.asyncio import Redis -from redis.exceptions import WatchError from reflex_base.config import get_config from reflex_base.environment import environment from reflex_base.utils.exceptions import ( @@ -115,6 +114,21 @@ class OplockFound(Exception): # noqa: N818 """Indicates that an opportunistic lock was found.""" +# KEYS: the lock key, then the state keys. ARGV: the lock id, the state +# expiration in seconds, then the serialized states in KEYS order. Writes only +# while the lock is held by ARGV[1], and returns the lock's remaining PTTL, or +# nil when the lock changed hands or expired and nothing was written. +_FENCED_SAVE_SCRIPT = """ +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return nil +end +for i = 2, #KEYS do + redis.call('SET', KEYS[i], ARGV[i + 1], 'EX', ARGV[2]) +end +return redis.call('PTTL', KEYS[1]) +""" + + @dataclasses.dataclass class StateManagerRedis(StateManager): """A state manager that stores states in redis.""" @@ -399,8 +413,7 @@ async def set_state( """ token = self._coerce_token(token) if isinstance(token, BaseStateToken): - # Serialize the touched states before checking the lock, so the - # check sits right before the single pipelined write. + # Serialize the touched states up front so they go out in one write. writes = self._collect_state_writes(token, cast(BaseState, state)) else: # Non-BaseState token: simple single-key write. @@ -420,82 +433,44 @@ async def set_state( if (event := context.get("event")) is not None else "" ) - # A dropped connection also surfaces as WatchError, since it loses the - # WATCH state, and so does a lease refresh touching the lock key: retry - # once on a fresh connection, which re-checks the lock before writing. - try: - await self._fenced_save(token, lock_id, writes, event_suffix) - except WatchError: - try: - await self._fenced_save(token, lock_id, writes, event_suffix) - except WatchError: - msg = ( - f"Lock for token {token} changed while its state was being " - "saved, so the save was discarded. Consider increasing " - f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " - "or use `@rx.event(background=True)` decorator for long-running tasks." - + event_suffix - ) - raise LockExpiredError(msg) from None - - async def _fenced_save( - self, - token: StateToken[Any], - lock_id: bytes, - writes: list[tuple[str, bytes]], - event_suffix: str, - ) -> None: - """Write the serialized states in one transaction fenced on the lock. - - WATCH makes EXEC fail if the lock key changed hands or expired between - the ownership check and the write. - - Args: - token: The token whose lock fences the write. - lock_id: The lock id the caller holds. - writes: The redis keys and pickled payloads to write. - event_suffix: Event context appended to warning and error messages. - - Raises: - LockExpiredError: If the lock is not held by lock_id. - WatchError: If the lock key changed or the connection dropped - while watching it. - """ lock_key = self._lock_key(token) - async with self.redis.pipeline(transaction=True) as pipeline: - await pipeline.watch(lock_key) - existing_lock_id = await pipeline.get(lock_key) - if existing_lock_id != lock_id: - msg = ( - f"Lock expired for token {token} while processing. Consider increasing " - f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " - "or use `@rx.event(background=True)` decorator for long-running tasks. " - f"Current lock id: {existing_lock_id!r}, expected lock id: {lock_id!r}." - + event_suffix + # One round trip: the script checks the lock and writes atomically on + # the server, so an expired or re-acquired lock discards every write, + # and a retried command re-checks the lock instead of bypassing it. + pttl = await cast( + "Awaitable[int | None]", + self.redis.eval( + _FENCED_SAVE_SCRIPT, + 1 + len(writes), + lock_key, + *(key for key, _ in writes), + lock_id, + self.token_expiration, + *(pickle_state for _, pickle_state in writes), + ), + ) + if pttl is None: + existing_lock_id = await self.redis.get(lock_key) + msg = ( + f"Lock expired for token {token} while processing. Consider increasing " + f"`app.state_manager.lock_expiration` (currently {self.lock_expiration}) " + "or use `@rx.event(background=True)` decorator for long-running tasks. " + f"Current lock id: {existing_lock_id!r}, expected lock id: {lock_id!r}." + + event_suffix + ) + raise LockExpiredError(msg) + if ( + isinstance(token, BaseStateToken) + and token.lock_key not in self._local_leases + ): + time_taken = (self.lock_expiration - pttl) / 1000 + if time_taken > self.lock_warning_threshold / 1000: + logger.warning( + f"Lock for token {token} was held too long {time_taken=}s, " + "use `@rx.event(background=True)` decorator for long-running " + f"tasks.{event_suffix}", + extra={"dedupe": True}, ) - raise LockExpiredError(msg) - if ( - isinstance(token, BaseStateToken) - and token.lock_key not in self._local_leases - ): - # Immediate mode on the watching connection: the parent client - # would need a second pooled connection while this one is held. - time_taken = ( - self.lock_expiration - (await pipeline.pttl(lock_key)) - ) / 1000 - if time_taken > self.lock_warning_threshold / 1000: - logger.warning( - f"Lock for token {token} was held too long {time_taken=}s, " - "use `@rx.event(background=True)` decorator for long-running " - f"tasks.{event_suffix}", - extra={"dedupe": True}, - ) - if not writes: - return - pipeline.multi() - for key, pickle_state in writes: - pipeline.set(key, pickle_state, ex=self.token_expiration) - await pipeline.execute() def _collect_state_writes( self, token: BaseStateToken, base_state: BaseState diff --git a/reflex/state.py b/reflex/state.py index eeffe277cb3..945e6a14da0 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -2643,8 +2643,9 @@ async def _diff_against_initial_state( ), ) cached = _initial_snapshot_cache[root_cls] - if not hashes or hashes[0] != cached.names_digest: - # The frontend was compiled against a different set of states. + if len(hashes) != len(cached.hashes) + 1 or hashes[0] != cached.names_digest: + # The frontend was compiled against a different set of states, or the + # payload is malformed. return delta frontend_hashes = dict(zip(sorted(cached.hashes), hashes[1:], strict=True)) diff: Delta = {} diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 67781940dc9..082ebaeb34e 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -9,7 +9,6 @@ import pytest import pytest_asyncio -from redis.exceptions import WatchError from reflex.istate.manager.redis import StateManagerRedis from reflex.istate.manager.token import BaseStateToken @@ -739,11 +738,11 @@ async def modify(): assert final_state.count == 2 -async def test_set_state_checks_lock_once_per_tree( +async def test_set_state_saves_tree_in_one_round_trip( state_manager_redis: StateManagerRedis, root_state: type[RedisTestState], ): - """Saving a state tree reads the lock TTL once, not once per substate. + """Saving a state tree checks the lock and writes every touched state in one command. Args: state_manager_redis: The StateManagerRedis to test. @@ -752,24 +751,17 @@ async def test_set_state_checks_lock_once_per_tree( state_manager_redis._oplock_enabled = False token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) redis = state_manager_redis.redis - real_pipeline = redis.pipeline - pttls: list[Any] = [] + real_eval = redis.eval + evals: list[tuple[Any, ...]] = [] - def counting_pipeline(*args, **kwargs): - pipe = real_pipeline(*args, **kwargs) - real_pttl = pipe.pttl - - def counting_pttl(key): - pttls.append(key) - return real_pttl(key) - - pipe.pttl = counting_pttl # pyright: ignore[reportAttributeAccessIssue] - return pipe + def counting_eval(script, numkeys, *keys_and_args): + evals.append(keys_and_args) + return real_eval(script, numkeys, *keys_and_args) async with state_manager_redis.modify_state(token) as state: assert len(state.substates) == 2 state.count = 1 - redis.pipeline = counting_pipeline # pyright: ignore[reportAttributeAccessIssue] + redis.eval = counting_eval # pyright: ignore[reportAttributeAccessIssue] try: await state_manager_redis.set_state( token, @@ -777,70 +769,21 @@ def counting_pttl(key): lock_id=await redis.get(state_manager_redis._lock_key(token)), ) finally: - redis.pipeline = real_pipeline # pyright: ignore[reportAttributeAccessIssue] + redis.eval = real_eval # pyright: ignore[reportAttributeAccessIssue] - # One TTL read for a tree of three states. - assert len(pttls) == 1 + # One command for a tree of three states, all of them touched by the load. + assert len(evals) == 1 + assert str(token) in evals[0] saved = await state_manager_redis.get_state(token) assert isinstance(saved, root_state) assert saved.count == 1 -async def test_set_state_retries_once_when_connection_drops_while_watching( - state_manager_redis: StateManagerRedis, - root_state: type[RedisTestState], -): - """A dropped connection during the fenced save is retried on a fresh one. - - redis-py reports a connection error while watching as WatchError, so the - save must re-check the lock on a new connection instead of failing. - - Args: - state_manager_redis: The StateManagerRedis to test. - root_state: The root state class. - """ - state_manager_redis._oplock_enabled = False - token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) - redis = state_manager_redis.redis - real_pipeline = redis.pipeline - pipelines: list[Any] = [] - - def dropping_pipeline(*args, **kwargs): - pipe = real_pipeline(*args, **kwargs) - pipelines.append(pipe) - if len(pipelines) == 1: - - async def dropped_execute(): - await pipe.reset() - msg = "A ConnectionError occurred while watching one or more keys" - raise WatchError(msg) - - pipe.execute = dropped_execute # pyright: ignore[reportAttributeAccessIssue] - return pipe - - async with state_manager_redis.modify_state(token) as state: - state.count = 7 - redis.pipeline = dropping_pipeline # pyright: ignore[reportAttributeAccessIssue] - try: - await state_manager_redis.set_state( - token, - state, - lock_id=await redis.get(state_manager_redis._lock_key(token)), - ) - finally: - redis.pipeline = real_pipeline # pyright: ignore[reportAttributeAccessIssue] - - assert len(pipelines) == 2 - saved = await state_manager_redis.get_state(token) - assert isinstance(saved, root_state) - assert saved.count == 7 - - async def test_set_state_discards_writes_when_lock_changes_hands( state_manager_redis: StateManagerRedis, root_state: type[RedisTestState], ): - """A save whose lock expired or was re-acquired between the check and the write is discarded. + """A save whose lock expired or was re-acquired before the write is discarded. Args: state_manager_redis: The StateManagerRedis to test. diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 9ad1ba0ea99..4f1bb6f81db 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -9,7 +9,6 @@ from unittest.mock import AsyncMock, Mock from redis.asyncio import Redis -from redis.exceptions import WatchError from redis.typing import EncodableT, KeyT from reflex.utils import prerequisites @@ -144,46 +143,8 @@ async def mock_pexpire(key: KeyT, px: int, xx: bool = False) -> bool: # noqa: R return False class _Pipeline: - """A pipeline that also models WATCH/MULTI/EXEC on the mocked keys.""" - def __init__(self): self.results = [] - self.watched: bytes | None = None - self.watched_value: Any = None - self.in_multi = False - - async def __aenter__(self): - return self - - async def __aexit__(self, *exc_info): - await self.reset() - - async def reset(self): - self.results = [] - self.watched = None - self.watched_value = None - self.in_multi = False - - async def watch(self, key: KeyT): - _expire_keys() - self.watched = _key_bytes(key) - self.watched_value = keys.get(self.watched) - - def multi(self): - self.in_multi = True - - def get(self, key: KeyT): - if self.watched is not None and not self.in_multi: - # Immediate mode while watching, before MULTI. - return redis_mock.get(key=key) - self.results.append(redis_mock.get(key=key)) - return None - - def pttl(self, key: KeyT): - if self.watched is not None and not self.in_multi: - return redis_mock.pttl(key=key) - self.results.append(redis_mock.pttl(key=key)) - return None def set( self, @@ -197,6 +158,9 @@ def set( redis_mock.set(key=key, value=value, ex=ex, px=px, nx=nx) ) + def get(self, key: KeyT): + self.results.append(redis_mock.get(key=key)) + def sadd(self, key: KeyT, value: EncodableT): self.results.append(redis_mock.sadd(key=key, value=value)) @@ -204,16 +168,6 @@ def pexpire(self, key: KeyT, px: int, xx: bool = False): self.results.append(redis_mock.pexpire(key=key, px=px, xx=xx)) async def execute(self): - _expire_keys() - if ( - self.watched is not None - and keys.get(self.watched) != self.watched_value - ): - for pending in self.results: - pending.close() - self.results = [] - msg = "Watched variable changed." - raise WatchError(msg) results = await asyncio.gather(*self.results) self.results = [] return results @@ -221,6 +175,29 @@ async def execute(self): def pipeline(transaction: bool = True): return _Pipeline() + async def mock_eval(script: str, numkeys: int, *keys_and_args: Any) -> Any: + """Emulate the one script the redis state manager runs. + + It is the fenced save: check the lock held in KEYS[1] against ARGV[1], + write the remaining keys with the expiration in ARGV[2], and return + the lock's PTTL, or None without writing when the lock is not held. + + Args: + script: The Lua source, unused. + numkeys: How many leading entries of keys_and_args are keys. + keys_and_args: The keys followed by the arguments. + + Returns: + The lock's PTTL after writing, or None when nothing was written. + """ + lock_key, *state_keys = keys_and_args[:numkeys] + lock_id, expiration, *payloads = keys_and_args[numkeys:] + if await redis_mock.get(lock_key) != lock_id: + return None + for key, payload in zip(state_keys, payloads, strict=True): + await redis_mock.set(key, payload, ex=int(expiration)) + return await redis_mock.pttl(lock_key) + async def pttl(key: KeyT) -> int: # noqa: RUF029 _expire_keys() return ( @@ -289,6 +266,7 @@ async def listen() -> AsyncGenerator[dict[str, Any] | None, None]: redis_mock.scard = mock_scard redis_mock.pexpire = mock_pexpire redis_mock.pipeline = pipeline + redis_mock.eval = mock_eval redis_mock.pttl = pttl redis_mock.pubsub = pubsub redis_mock.config_set = AsyncMock() 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 52ee00fe864..5bb78c2a755 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 @@ -897,3 +897,14 @@ def set_ratio_int(self): await future.wait_all() snapshot = emitted_deltas[0][1] assert "loads" + FIELD_MARKER in snapshot[CookieState.get_full_name()] + + # A matching names digest with a truncated hash list is malformed and + # also falls back to the full snapshot instead of failing the event. + emitted_deltas.clear() + async with real_base_state_processor as processor: + future = await processor.enqueue( + token, _boot_event(boot_name, {"hashes": hashes[:-1]}) + ) + await future.wait_all() + snapshot = emitted_deltas[0][1] + assert "loads" + FIELD_MARKER in snapshot[CookieState.get_full_name()] From 62280c0ea6053bc778a2546ed3b304c65b1aec4d Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 21:33:47 +0000 Subject: [PATCH 12/14] Document the millisecond to second conversion in the lock warning --- reflex/istate/manager/redis.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index cf45e961379..e904f1fc8c2 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -463,6 +463,8 @@ async def set_state( isinstance(token, BaseStateToken) and token.lock_key not in self._local_leases ): + # lock_expiration, the PTTL and the threshold are milliseconds; the + # warning reports the time held in seconds. time_taken = (self.lock_expiration - pttl) / 1000 if time_taken > self.lock_warning_threshold / 1000: logger.warning( From 8a5abccbeaaf038227936ca56bac39142188bfe1 Mon Sep 17 00:00:00 2001 From: Farhan Ali Raza Date: Mon, 7 Sep 2026 21:35:14 +0000 Subject: [PATCH 13/14] Type the fenced save EVAL call for the minimum redis-py version Older redis-py stubs declare the EVAL arguments and reply as str, while the save passes bytes keys and payloads and reads back an int or nil. --- reflex/istate/manager/redis.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index e904f1fc8c2..364e78c475d 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -9,7 +9,7 @@ import sys import time import uuid -from collections.abc import AsyncIterator, Awaitable +from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any, TypedDict, cast from redis import ResponseError @@ -437,17 +437,17 @@ async def set_state( # One round trip: the script checks the lock and writes atomically on # the server, so an expired or re-acquired lock discards every write, # and a retried command re-checks the lock instead of bypassing it. - pttl = await cast( - "Awaitable[int | None]", - self.redis.eval( - _FENCED_SAVE_SCRIPT, - 1 + len(writes), - lock_key, - *(key for key, _ in writes), - lock_id, - self.token_expiration, - *(pickle_state for _, pickle_state in writes), - ), + # redis-py types EVAL arguments as str and its reply as str; both keys + # and payloads are bytes here and the script replies with an int or nil. + fenced_save = cast("Callable[..., Awaitable[int | None]]", self.redis.eval) + pttl = await fenced_save( + _FENCED_SAVE_SCRIPT, + 1 + len(writes), + lock_key, + *(key for key, _ in writes), + lock_id, + self.token_expiration, + *(pickle_state for _, pickle_state in writes), ) if pttl is None: existing_lock_id = await self.redis.get(lock_key) From 306c2386405ca7050742864db84dee24909bad7a Mon Sep 17 00:00:00 2001 From: Farhan Date: Tue, 8 Sep 2026 23:08:45 +0500 Subject: [PATCH 14/14] Warm the websocket transport early and cache state metadata per class Open the socket.io transport in a microtask before React mounts, then hand the warm socket to connect() so the first hydrate_and_load round trip does not wait for the mount. Unclaimed warm sockets are discarded on a timeout, on pagehide, on a transport mismatch, or on HMR dispose. Read the Redis state tree with one MGET instead of a pipelined GET per state, and skip the read when the tree is already populated. Store the immutable per-class metadata (parent, root, name, full name) on the owning class so the 128-entry LRU cannot evict it in apps with many states, and add a hydration benchmark that covers both sides of the LRU capacity. Claude-Session: https://claude.ai/code/session_017ahSHCfgq16R8hSLBWH4p2 --- news/+first-load-hydrate.performance.md | 2 +- .../news/+first-load-hydrate.performance.md | 2 +- .../reflex_base/.templates/web/utils/state.js | 92 +++++- reflex/istate/manager/redis.py | 16 +- reflex/state.py | 53 +++- tests/benchmarks/test_state_hydration.py | 68 +++++ tests/units/compiler/state_js.test.mjs | 273 ++++++++++++++++++ .../units/compiler/test_state_js_template.py | 23 ++ tests/units/istate/manager/test_redis.py | 47 +++ tests/units/mock_redis.py | 15 + tests/units/test_state.py | 19 ++ 11 files changed, 590 insertions(+), 20 deletions(-) create mode 100644 tests/benchmarks/test_state_hydration.py create mode 100644 tests/units/compiler/state_js.test.mjs diff --git a/news/+first-load-hydrate.performance.md b/news/+first-load-hydrate.performance.md index 8b7c9f82f29..5a034a31790 100644 --- a/news/+first-load-hydrate.performance.md +++ b/news/+first-load-hydrate.performance.md @@ -1 +1 @@ -Faster first page load: the frontend hydrates in a single event sent with the websocket connect, and only the values that differ from the compiled defaults are sent (a 20-substate app went from 36 KB to 2 KB and from 5 to 3 frames). Redis-backed apps also do about 80% fewer redis commands per page load. +Speed up first page loads by combining hydration with the websocket connect and sending only values that differ from compiled defaults. Reduce Redis state-tree read/write overhead and avoid repeated class metadata computation in apps with many states. diff --git a/packages/reflex-base/news/+first-load-hydrate.performance.md b/packages/reflex-base/news/+first-load-hydrate.performance.md index ac601cae98f..cb3e8a81216 100644 --- a/packages/reflex-base/news/+first-load-hydrate.performance.md +++ b/packages/reflex-base/news/+first-load-hydrate.performance.md @@ -1 +1 @@ -The compiled frontend hydrates with a single `hydrate_and_load` event sent along with the websocket connect, so the first page load needs one fewer round trip. +Begin opening the websocket transport before React mounts, then hydrate with a single `hydrate_and_load` event sent along with the websocket connect to save a round trip. diff --git a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js index c6e88096376..ad9e59b6ee7 100644 --- a/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js +++ b/packages/reflex-base/src/reflex_base/.templates/web/utils/state.js @@ -144,6 +144,74 @@ export const isBackendDisabled = () => { return cookie !== undefined && cookie.split("=")[1] == "false"; }; +/** + * Create a socket without starting its namespace or hydration events. + * @param endpoint The backend URL. + * @param transports The configured transports. + * @returns The disconnected socket. + */ +const createSocket = (endpoint, transports) => + io(endpoint.href, { + path: endpoint.pathname, + transports, + protocols: [reflexEnvironment.version], + autoUnref: false, + autoConnect: false, + query: { token: getToken() }, + reconnection: false, + }); + +let warmSocket = null; +let cancelWarmup = () => {}; +let socketStarted = false; + +/** Close an unclaimed transport and remove its cleanup handlers. */ +const discardWarmSocket = () => { + const socket = warmSocket; + warmSocket = null; + cancelWarmup(); + socket?.disconnect(); +}; + +// Start only the transport while React is still preparing to mount. The +// namespace stays disconnected until connect() installs all its handlers. +// Defer past module evaluation because context.js imports this module too. +if (typeof window !== "undefined") { + queueMicrotask(() => { + if ( + socketStarted || + Object.keys(initialState).length <= 1 || + isBackendDisabled() || + document.visibilityState === "hidden" + ) { + return; + } + try { + warmSocket = createSocket(getBackendURL(EVENTURL), [env.TRANSPORT]); + } catch { + // Speculative setup may fail (for example, blocked session storage). + // The normal connection path will report failures when the app mounts. + return; + } + const timeout = setTimeout(discardWarmSocket, 10000); + window.addEventListener("pagehide", discardWarmSocket); + cancelWarmup = () => { + clearTimeout(timeout); + window.removeEventListener("pagehide", discardWarmSocket); + }; + warmSocket.io.open((error) => { + if (error) discardWarmSocket(); + }); + }); +} + +if (import.meta.hot) { + import.meta.hot.dispose(() => { + socketStarted = true; + discardWarmSocket(); + }); +} + /** * Determine if any event in the event queue is stateful. * @@ -606,15 +674,20 @@ export const connect = async ( }); // Create the socket. - socket.current = io(endpoint.href, { - path: endpoint["pathname"], - transports: transports, - protocols: [reflexEnvironment.version], - autoUnref: false, - query: { token: getToken() }, - auth: bootAuth(true), - reconnection: false, // Reconnection will be handled manually. - }); + socketStarted = true; + if ( + warmSocket && + (warmSocket.io.opts.transports.length !== transports.length || + transports.some( + (transport, i) => transport !== warmSocket.io.opts.transports[i], + )) + ) { + discardWarmSocket(); + } + socket.current = warmSocket ?? createSocket(endpoint, transports); + warmSocket = null; + cancelWarmup(); + socket.current.auth = bootAuth(true); socket.current.wait_connect = !socket.current.connected; // Ensure undefined fields in events are sent as null instead of removed socket.current.io.encoder.replacer = (k, v) => (v === undefined ? null : v); @@ -804,6 +877,7 @@ export const connect = async ( }); document.addEventListener("visibilitychange", checkVisibility); + socket.current.connect(); }; /** diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 364e78c475d..09e4dd25542 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -344,14 +344,20 @@ async def get_state( key=lambda x: x.get_full_name(), ) - redis_pipeline = self.redis.pipeline() - for state_cls in required_state_classes: - redis_pipeline.get(str(token.with_cls(state_cls))) + # Read the tree atomically with one command instead of a transaction + # containing a GET for each state. An already populated tree needs no IO. + redis_states = ( + await self.redis.mget([ + str(token.with_cls(state_cls)) for state_cls in required_state_classes + ]) + if required_state_classes + else [] + ) for state_cls, redis_state in zip( required_state_classes, - await redis_pipeline.execute(), - strict=False, + redis_states, + strict=True, ): state = None diff --git a/reflex/state.py b/reflex/state.py index 945e6a14da0..dd851054313 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -327,6 +327,45 @@ def _override_base_method(fn: Callable[PARAMS, RETURN]) -> Callable[PARAMS, RETU return fn +def _cache_per_class( + fn: Callable[[type[BaseState]], RETURN], +) -> Callable[[type[BaseState]], RETURN]: + """Cache immutable metadata on the class that owns it. + + A small LRU keeps hot lookups fast; evicted values remain on their owning + classes so large apps never recompute them. Read the class's own dict so + subclasses never inherit their parent's cached result. + + Args: + fn: The class method to cache. + + Returns: + A method that computes its value once per class. + """ + cache_key = fn.__name__ + + @functools.lru_cache + @functools.wraps(fn) + def wrapped(cls: type[BaseState]) -> RETURN: + """Return the metadata owned by this class. + + Args: + cls: The state class. + + Returns: + The cached metadata. + """ + cache = cls.__dict__["_reflex_internal_class_cache"] + try: + return cache[cache_key] + except KeyError: + value = fn(cls) + cache[cache_key] = value + return value + + return wrapped + + def _has_data_descriptor(cls: type, name: str) -> bool: """Whether the class provides a descriptor that handles assignment for `name`. @@ -412,6 +451,7 @@ def _is_user_descriptor(value: Any) -> bool: }) CLASS_VAR_NAMES = frozenset({ + "_reflex_internal_class_cache", "_fast_attr_names", "vars", "base_vars", @@ -430,6 +470,9 @@ def _is_user_descriptor(value: Any) -> bool: class BaseState(EvenMoreBasicBaseState): """The state of the app.""" + # Immutable metadata belongs to each class, including when an LRU evicts it. + _reflex_internal_class_cache: ClassVar[builtins.dict[str, Any]] = {} + # A map from the var name to the var. vars: ClassVar[builtins.dict[str, Var]] = {} @@ -607,6 +650,8 @@ def __init_subclass__(cls, mixin: bool = False, **kwargs): super().__init_subclass__(**kwargs) + cls._reflex_internal_class_cache = {} + if cls._mixin: return @@ -1131,7 +1176,7 @@ def get_skip_vars(cls) -> set[str]: ) @classmethod - @functools.lru_cache + @_cache_per_class def get_parent_state(cls) -> type[BaseState] | None: """Get the parent state. @@ -1159,7 +1204,7 @@ def get_parent_state(cls) -> type[BaseState] | None: return None # No known parent @classmethod - @functools.lru_cache + @_cache_per_class def get_root_state(cls) -> type[BaseState]: """Get the root state. @@ -1179,7 +1224,7 @@ def get_substates(cls) -> set[type[BaseState]]: return RegistrationContext.get().get_substates(cls) @classmethod - @functools.lru_cache + @_cache_per_class def get_name(cls) -> str: """Get the name of the state. @@ -1190,7 +1235,7 @@ def get_name(cls) -> str: return format.to_snake_case(f"{module}___{cls.__name__}") @classmethod - @functools.lru_cache + @_cache_per_class def get_full_name(cls) -> str: """Get the full name of the state. diff --git a/tests/benchmarks/test_state_hydration.py b/tests/benchmarks/test_state_hydration.py new file mode 100644 index 00000000000..dcc077b83a0 --- /dev/null +++ b/tests/benchmarks/test_state_hydration.py @@ -0,0 +1,68 @@ +"""Track state-tree hydration costs below and above the metadata LRU capacity.""" + +import pytest +from pytest_codspeed import BenchmarkFixture + +from reflex.state import BaseState + + +@pytest.fixture(scope="module", params=[20, 200]) +def hydration_state(request: pytest.FixtureRequest) -> BaseState: + """Create a state tree with the requested number of substates. + + Args: + request: The parametrized fixture request. + + Returns: + A populated state tree. + """ + count = request.param + root = type(f"HydrationRoot{count}", (BaseState,), {"__module__": __name__}) + for index in range(count): + type( + f"HydrationChild{count}_{index}", + (root,), + { + "__module__": __name__, + "__annotations__": {"count": int, "label": str}, + "count": index, + "label": f"sub-{index}", + }, + ) + return root() + + +def test_hydration_snapshot( + hydration_state: BaseState, benchmark: BenchmarkFixture +) -> None: + """Measure complete snapshots of small and large state trees. + + Args: + hydration_state: The state tree to snapshot. + benchmark: The benchmark fixture. + """ + benchmark(hydration_state.dict) + + +def test_hydration_metadata( + hydration_state: BaseState, benchmark: BenchmarkFixture +) -> None: + """Measure repeated metadata walks without state serialization costs. + + Args: + hydration_state: The populated tree whose classes to visit. + benchmark: The benchmark fixture. + """ + classes = [ + type(hydration_state), + *(type(s) for s in hydration_state.substates.values()), + ] + + @benchmark + def walk() -> None: + """Visit the metadata used to fetch and reconnect a state tree.""" + for cls in classes: + cls.get_name() + cls.get_full_name() + cls.get_parent_state() + cls.get_root_state() diff --git a/tests/units/compiler/state_js.test.mjs b/tests/units/compiler/state_js.test.mjs new file mode 100644 index 00000000000..9bb0ed24ec0 --- /dev/null +++ b/tests/units/compiler/state_js.test.mjs @@ -0,0 +1,273 @@ +import assert from "node:assert/strict"; +import { readFile } from "node:fs/promises"; +import test from "node:test"; +import vm from "node:vm"; + +const source = await readFile( + new URL( + "../../../packages/reflex-base/src/reflex_base/.templates/web/utils/state.js", + import.meta.url, + ), + "utf8", +); + +/** Evaluate the actual frontend module with controlled transport and browser APIs. */ +async function setup({ + browser = true, + stateful = true, + disabled = false, + hidden = false, +} = {}) { + const sockets = []; + const microtasks = []; + const timers = new Map(); + const firstHydrates = []; + const updates = []; + const window = new EventTarget(); + window.location = new URL("http://localhost:3000/page?query=value"); + const storage = new Map([["token", "original-token"]]); + window.sessionStorage = { + getItem: (key) => storage.get(key), + setItem: (key, value) => storage.set(key, value), + }; + const document = new EventTarget(); + document.cookie = disabled ? "backend-enabled=false" : ""; + document.visibilityState = hidden ? "hidden" : "visible"; + let dispose; + const context = vm.createContext({ + console, + URL, + URLSearchParams, + performance, + ...(browser ? { window, document } : {}), + queueMicrotask: (fn) => microtasks.push(fn), + setTimeout: (fn, ms) => { + timers.set(fn, ms); + return fn; + }, + clearTimeout: (fn) => timers.delete(fn), + }); + window.setTimeout = context.setTimeout; + + function io(url, opts) { + const handlers = new Map(); + const socket = { + connected: false, + namespaceConnects: 0, + disconnects: 0, + io: { + opts, + encoder: {}, + decoder: {}, + opens: 0, + open(callback) { + this.opens++; + this.openCallback = callback; + }, + }, + on(name, handler) { + handlers.set(name, handler); + }, + connect() { + // A ready transport can deliver these synchronously on namespace connect. + for (const name of [ + "connect", + "connect_error", + "event", + "new_token", + "disconnect", + ]) { + assert.ok(handlers.has(name), `missing handler: ${name}`); + } + this.namespaceConnects++; + this.connected = true; + handlers.get("connect")(); + handlers.get("new_token")("assigned-token"); + handlers.get("event")({ delta: { child: { count: 7 } } }); + }, + disconnect() { + this.disconnects++; + this.connected = false; + }, + }; + sockets.push(socket); + return socket; + } + const imports = { + "socket.io-client": { default: io }, + "$/env.json": { + default: { EVENT: "ws://localhost:8000/_event", TRANSPORT: "websocket" }, + }, + "$/reflex.json": { default: { version: "test" } }, + "universal-cookie": { default: class {} }, + react: { useCallback() {}, useEffect() {}, useRef() {}, useState() {} }, + "react-router": { + useLocation() {}, + useNavigate() {}, + useSearchParams() {}, + useParams() {}, + }, + "$/utils/context": { + initialEvents(first) { + firstHydrates.push(first); + return [ + { + name: "root.hydrate_and_load", + payload: first ? { hashes: ["defaults"] } : {}, + }, + ]; + }, + initialState: stateful ? { root: {}, child: {} } : {}, + onLoadInternalEvent() {}, + state_name: "root", + exception_state_name: "exception", + }, + "$/utils/helpers/debounce": { default() {} }, + "$/utils/helpers/throttle": { default() {} }, + "$/utils/helpers/upload": { uploadFiles() {} }, + }; + const module = new vm.SourceTextModule(source, { + context, + initializeImportMeta(meta) { + meta.hot = { + dispose(fn) { + dispose = fn; + }, + }; + }, + }); + await module.link((name) => { + assert.ok(imports[name], `unexpected import: ${name}`); + return new vm.SyntheticModule( + Object.keys(imports[name]), + function () { + for (const [key, value] of Object.entries(imports[name])) + this.setExport(key, value); + }, + { context }, + ); + }); + await module.evaluate(); + const socket = { current: null }; + return { + sockets, + timers, + firstHydrates, + window, + socket, + updates, + dispose, + flush: () => microtasks.splice(0).forEach((fn) => fn()), + connect: (transports = ["websocket"]) => + module.namespace.connect( + socket, + { child: (delta) => updates.push(delta.count) }, + transports, + () => {}, + {}, + () => {}, + { current: {} }, + ), + }; +} + +test("warm the transport without hydrating, then reuse it with every handler attached", async () => { + const app = await setup(); + app.flush(); + assert.equal(app.sockets.length, 1); + assert.equal(app.sockets[0].io.opens, 1); + assert.equal(app.sockets[0].namespaceConnects, 0); + assert.equal(app.sockets[0].io.opts.autoConnect, false); + assert.deepEqual(app.firstHydrates, []); + await app.connect(); + assert.equal(app.sockets.length, 1); + assert.equal(app.socket.current.namespaceConnects, 1); + assert.deepEqual(app.updates, [7]); + assert.deepEqual(app.firstHydrates, [true]); + assert.equal(app.timers.size, 0); + assert.equal(app.window.sessionStorage.getItem("token"), "assigned-token"); + assert.equal(app.socket.current.auth.event.router_data.pathname, "/page"); +}); + +test("reconnect uses the assigned token and requests a full hydrate", async () => { + const app = await setup(); + app.flush(); + await app.connect(); + app.socket.current.connected = false; + app.socket.current.reconnect(); + assert.equal(app.sockets.length, 1); + assert.equal(app.socket.current.io.opts.query.token, "assigned-token"); + assert.deepEqual(app.firstHydrates, [true, false]); + assert.equal(app.socket.current.namespaceConnects, 2); +}); + +for (const reason of ["error", "timeout", "pagehide", "hot reload"]) { + test(`discard an unclaimed warm transport on ${reason}`, async () => { + const app = await setup(); + app.flush(); + const warm = app.sockets[0]; + if (reason === "error") warm.io.openCallback(new Error("offline")); + if (reason === "timeout") [...app.timers.keys()][0](); + if (reason === "pagehide") app.window.dispatchEvent(new Event("pagehide")); + if (reason === "hot reload") app.dispose(); + assert.equal(warm.disconnects, 1); + assert.equal(app.timers.size, 0); + await app.connect(); + assert.equal(app.sockets.length, 2); + assert.equal(app.socket.current.namespaceConnects, 1); + }); +} + +test("a late warmup error cannot close the socket after React claims it", async () => { + const app = await setup(); + app.flush(); + await app.connect(); + app.socket.current.io.openCallback(new Error("late failure")); + assert.equal(app.socket.current.disconnects, 0); +}); + +test("a synchronous mount does not leave a second speculative connection", async () => { + const app = await setup(); + await app.connect(); + app.flush(); + assert.equal(app.sockets.length, 1); +}); + +test("hot reload before the microtask runs prevents speculative setup", async () => { + const app = await setup(); + app.dispose(); + app.flush(); + assert.equal(app.sockets.length, 0); +}); + +test("blocked session storage does not throw from speculative setup", async () => { + const app = await setup(); + app.window.sessionStorage.getItem = () => { + throw new Error("storage blocked"); + }; + app.flush(); + assert.equal(app.sockets.length, 0); + await assert.rejects(app.connect(), /storage blocked/); +}); + +test("a caller selecting another transport gets its requested configuration", async () => { + const app = await setup(); + app.flush(); + await app.connect(["polling"]); + assert.equal(app.sockets[0].disconnects, 1); + assert.equal(app.socket.current.io.opts.transports[0], "polling"); +}); + +for (const config of [ + { browser: false }, + { stateful: false }, + { disabled: true }, + { hidden: true }, +]) { + test(`skip speculative connections for ${JSON.stringify(config)}`, async () => { + const app = await setup(config); + app.flush(); + assert.equal(app.sockets.length, 0); + assert.equal(app.timers.size, 0); + }); +} diff --git a/tests/units/compiler/test_state_js_template.py b/tests/units/compiler/test_state_js_template.py index 3cc81b70858..04eb4ec8725 100644 --- a/tests/units/compiler/test_state_js_template.py +++ b/tests/units/compiler/test_state_js_template.py @@ -1,13 +1,36 @@ """Regression tests for the state.js frontend template.""" +import shutil +import subprocess from pathlib import Path +import pytest + STATE_JS_TEMPLATE = ( Path(__file__).parents[3] / "packages/reflex-base/src/reflex_base/.templates/web/utils/state.js" ) +def test_socket_startup_lifecycle() -> None: + """Execute the frontend startup and reconnect tests against the real template.""" + node = shutil.which("node") + if node is None: + pytest.skip("Node.js is required for the frontend runtime tests") + result = subprocess.run( + [ + node, + "--experimental-vm-modules", + str(Path(__file__).with_name("state_js.test.mjs")), + ], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_state_js_does_not_register_deprecated_unload_listener() -> None: """The template must not register the deprecated `unload` event listener. diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 082ebaeb34e..86b9cfde2ca 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -6,6 +6,7 @@ import uuid from collections.abc import AsyncGenerator from typing import Any +from unittest.mock import AsyncMock, Mock import pytest import pytest_asyncio @@ -31,6 +32,52 @@ class SubState2(RedisTestState): """A test substate for redis state manager tests.""" +@pytest.mark.asyncio +async def test_get_state_reads_tree_in_one_command(): + """Read persisted and missing states together, preserving their tree positions.""" + redis = mock_redis() + manager = StateManagerRedis(redis=redis) + token = BaseStateToken(ident="batched-read", cls=RedisTestState) + persisted = RedisTestState() + persisted.foo = "persisted" + classes = sorted( + manager._get_required_state_classes(RedisTestState, subclasses=True), + key=lambda cls: cls.get_full_name(), + ) + redis.mget = AsyncMock( + return_value=[ + persisted._serialize() if cls is RedisTestState else None for cls in classes + ] + ) + + state = await manager.get_state(token) + + assert isinstance(state, RedisTestState) + assert state.foo == "persisted" + assert state.count == 0 + assert set(state.substates) == {SubState1.get_name(), SubState2.get_name()} + assert all(child.parent_state is state for child in state.substates.values()) + redis.mget.assert_awaited_once_with([str(token.with_cls(cls)) for cls in classes]) + await manager.close() + + +@pytest.mark.asyncio +async def test_get_state_reuses_populated_tree_without_reading(): + """Fetching an already attached state must not contact Redis or replace it.""" + redis = mock_redis() + manager = StateManagerRedis(redis=redis) + token = BaseStateToken(ident="populated-read", cls=SubState1) + state = RedisTestState() + redis.mget = AsyncMock(return_value=[]) + redis.pipeline = Mock(side_effect=AssertionError("Unexpected Redis read")) + + child = await manager.get_state(token, top_level=False, for_state_instance=state) + + assert child is state.substates[SubState1.get_name()] + redis.mget.assert_not_awaited() + await manager.close() + + @pytest.fixture def root_state() -> type[RedisTestState]: diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 4f1bb6f81db..7c1229d6723 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -61,6 +61,20 @@ async def mock_get(key: KeyT): # noqa: RUF029 _expire_keys() return keys.get(_key_bytes(key)) + async def mock_mget(requested_keys: list[KeyT]): + """Read keys in request order, including missing values. + + Args: + requested_keys: The keys to read. + + Returns: + The stored values, or None for missing keys. + """ + # Let concurrent tasks run, as they would during the real Redis IO. + await asyncio.sleep(0) + _expire_keys() + return [keys.get(_key_bytes(key)) for key in requested_keys] + async def mock_set( # noqa: RUF029 key: KeyT, value: EncodableT, @@ -258,6 +272,7 @@ async def listen() -> AsyncGenerator[dict[str, Any] | None, None]: redis_mock = AsyncMock(spec=Redis) redis_mock.get = mock_get + redis_mock.mget = mock_mget redis_mock.set = mock_set redis_mock.delete = mock_delete redis_mock.getdel = mock_getdel diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 08cda78c1e4..a7e204c3c60 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -507,6 +507,25 @@ def test_get_parent_state(): assert GrandchildState.get_parent_state() == ChildState +def test_state_names_remain_cached_for_large_apps(mocker: MockerFixture): + """Walking more than 128 states must not evict their immutable class names. + + Args: + mocker: The mock fixture. + """ + states = [ + type(f"CachedNameState{i}", (BaseState,), {"__module__": __name__}) + for i in range(200) + ] + names = [state.get_full_name() for state in states] + snake_case = mocker.patch( + "reflex.state.format.to_snake_case", wraps=format.to_snake_case + ) + + assert [state.get_full_name() for state in states] == names + snake_case.assert_not_called() + + def test_get_substates(): """Test getting the substates.""" assert TestState.get_substates() == {ChildState, ChildState2, ChildState3}