diff --git a/news/+first-load-hydrate.performance.md b/news/+first-load-hydrate.performance.md new file mode 100644 index 00000000000..5a034a31790 --- /dev/null +++ b/news/+first-load-hydrate.performance.md @@ -0,0 +1 @@ +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 new file mode 100644 index 00000000000..cb3e8a81216 --- /dev/null +++ b/packages/reflex-base/news/+first-load-hydrate.performance.md @@ -0,0 +1 @@ +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 8ba6d00509c..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. * @@ -395,7 +463,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 +503,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,15 +665,29 @@ export const connect = async ( const endpoint = getBackendURL(EVENTURL); const on_hydrated_queue = []; - // Create the socket. - socket.current = io(endpoint.href, { - path: endpoint["pathname"], - transports: transports, - protocols: [reflexEnvironment.version], - autoUnref: false, - query: { token: getToken() }, - reconnection: false, // Reconnection will be handled manually. + // 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. The key is read by + // the backend as CompileVars.CONNECT_AUTH_EVENT. + const bootAuth = (first) => ({ + event: withRouterData(initialEvents(first)[0], params), }); + + // Create the socket. + 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); @@ -623,8 +713,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 +766,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); @@ -790,6 +877,7 @@ export const connect = async ( }); document.addEventListener("visibilitychange", checkVisibility); + socket.current.connect(); }; /** @@ -1062,14 +1150,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..c16b82223f6 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,18 +319,26 @@ 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}', - {{vars: client_storage_vars}}, + {{{constants.CompileVars.PAYLOAD_VARS}: client_storage_vars}}, ), ); }} @@ -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["{constants.CompileVars.PAYLOAD_VARS}"] = client_storage_vars; + }} + if (first) {{ + payload["{constants.CompileVars.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..977451b7456 100644 --- a/packages/reflex-base/src/reflex_base/constants/compiler.py +++ b/packages/reflex-base/src/reflex_base/constants/compiler.py @@ -57,6 +57,17 @@ 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 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/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/base_state_processor.py b/packages/reflex-base/src/reflex_base/event/processor/base_state_processor.py index de04bac1e93..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 @@ -36,10 +37,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 (CompileVars.HYDRATE, CompileVars.HYDRATE_AND_LOAD) + ) def _check_valid_yield(events: Any, handler_name: str = "unknown") -> Any: @@ -410,7 +414,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/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/app.py b/reflex/app.py index ff18d3bda0c..6694651ddd2 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 ( @@ -2003,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. @@ -2026,6 +2028,20 @@ 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(constants.CompileVars.CONNECT_AUTH_EVENT)) + is not None + ): + 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. @@ -2277,11 +2293,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/compiler/compiler.py b/reflex/compiler/compiler.py index 3ba8746316e..867453d8eb0 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,21 +218,21 @@ def _compile_contexts(state: type[BaseState] | None, theme: Component | None) -> not is_prod_mode() and not environment.REFLEX_REACT_OWNER_STACKS.get() ) - return ( - templates.context_template( - initial_state=utils.compile_state(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, ) 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 033998d1ef5..09e4dd25542 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 +from collections.abc import AsyncIterator, Awaitable, Callable from typing import Any, TypedDict, cast from redis import ResponseError @@ -114,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.""" @@ -329,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 @@ -397,46 +418,61 @@ async def set_state( RuntimeError: If the state instance doesn't match the state name in the token. """ token = self._coerce_token(token) - # 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 - ): + if isinstance(token, BaseStateToken): + # 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. + pickle_state = token.serialize(state) + writes = [(str(token), pickle_state)] if pickle_state else [] + + 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 + + event_suffix = ( + f" Happened in event: {event.name}" + if (event := context.get("event")) is not None + else "" + ) + lock_key = self._lock_key(token) + # 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. + # 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) 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 "" - ) + + event_suffix ) 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 - - base_state = cast(BaseState, state) - - lock_key = token.lock_key - - if lock_id is not None and lock_key not in self._local_leases: - time_taken = ( - self.lock_expiration - (await self.redis.pttl(self._lock_key(token))) - ) / 1000 + if ( + 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: - 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 " @@ -444,32 +480,27 @@ async def set_state( extra={"dedupe": True}, ) - # Recursively set_state on all known substates. - 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()}", - ) - 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, - ) + 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. - # Wait for substates to be persisted. - for t in tasks: - await t + Args: + token: The token (any state class) identifying the client. + base_state: The state instance whose tree to persist. + + 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/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..dd851054313 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 @@ -326,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`. @@ -411,6 +451,7 @@ def _is_user_descriptor(value: Any) -> bool: }) CLASS_VAR_NAMES = frozenset({ + "_reflex_internal_class_cache", "_fast_attr_names", "vars", "base_vars", @@ -429,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]] = {} @@ -606,6 +650,8 @@ def __init_subclass__(cls, mixin: bool = False, **kwargs): super().__init_subclass__(**kwargs) + cls._reflex_internal_class_cache = {} + if cls._mixin: return @@ -1130,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. @@ -1158,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. @@ -1178,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. @@ -1189,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. @@ -2474,10 +2520,238 @@ def set_is_hydrated(self, value: bool) -> None: """ self.is_hydrated = value + @event(supersedes=constants.CompileVars.ON_LOAD_SUPERSEDE_GROUP) + 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 _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 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 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 [ + _short_digest("\n".join(names)), + *( + _short_digest(format.json_dumps(snapshot[state_name], sort_keys=True)) + for state_name in names + ), + ] + + +@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]] + # The digest of the sorted state names. + names_digest: 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. + """ + 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() + }, + names_digest=names_digest, + hashes=dict(zip(sorted(snapshot), state_hashes, strict=True)), + ) + + +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. + + 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. + 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 + defaults match the backend's, and left untouched for the others. + """ + cached = _initial_snapshot_cache.get(root_cls) + 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] + 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 = {} + for state_name, state_vars in delta.items(): + default_vars = cached.serialized.get(state_name) + if default_vars is None or frontend_hashes.get(state_name) != cached.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 _serialize_var(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: State, +) -> 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 + 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 +2868,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): @@ -2609,29 +2877,29 @@ 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. 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, - ), - State.set_is_hydrated(True), - ] + return _load_events_for_page(self) + + @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): 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 5ec8de77be6..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]: @@ -736,3 +783,70 @@ async def modify(): ) assert isinstance(final_state, root_state) assert final_state.count == 2 + + +async def test_set_state_saves_tree_in_one_round_trip( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """Saving a state tree checks the lock and writes every touched state in one command. + + 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_eval = redis.eval + evals: list[tuple[Any, ...]] = [] + + 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.eval = counting_eval # 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.eval = real_eval # pyright: ignore[reportAttributeAccessIssue] + + # 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_discards_writes_when_lock_changes_hands( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """A save whose lock expired or was re-acquired before 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..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, @@ -142,39 +156,61 @@ async def mock_pexpire(key: KeyT, px: int, xx: bool = False) -> bool: # noqa: R return True return False - def pipeline(): - pipeline_mock = Mock() - results = [] - - def get_pipeline(key: KeyT): - results.append(redis_mock.get(key=key)) + class _Pipeline: + def __init__(self): + self.results = [] - def set_pipeline( + 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)) - - def sadd_pipeline(key: KeyT, value: EncodableT): - 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)) - - async def execute(): - _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 + self.results.append( + 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)) + + 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): + results = await asyncio.gather(*self.results) + self.results = [] + return results + + 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() @@ -236,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 @@ -244,6 +281,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 a70130f70c1..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 @@ -745,3 +745,166 @@ 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:]] == [ + {state_name: {hydrated_key: False}}, + {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[2][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 + 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() + compiled = compile_state(State) + hashes = state_snapshot_hashes(compiled) + + 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" + } + + # 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 (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": ["0" * 16, *hashes[1:]]}) + ) + 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()] diff --git a/tests/units/test_app.py b/tests/units/test_app.py index ac7fca79696..3770de9bf04 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -4344,3 +4344,56 @@ 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() + + +@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" + ) diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 37c0ed2fc4c..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} @@ -2171,7 +2190,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): @@ -5234,6 +5254,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