diff --git a/news/6906.performance.md b/news/6906.performance.md new file mode 100644 index 00000000000..9e9f5db5f8d --- /dev/null +++ b/news/6906.performance.md @@ -0,0 +1 @@ +State deltas no longer re-ship the session block and request headers on every navigation; only the router fields that changed (page, url, route_id) are sent, cutting the per-navigation delta by ~72%. diff --git a/packages/reflex-base/news/6906.performance.md b/packages/reflex-base/news/6906.performance.md new file mode 100644 index 00000000000..9e9f5db5f8d --- /dev/null +++ b/packages/reflex-base/news/6906.performance.md @@ -0,0 +1 @@ +State deltas no longer re-ship the session block and request headers on every navigation; only the router fields that changed (page, url, route_id) are sent, cutting the per-navigation delta by ~72%. 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..c66b2ccb4d1 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 @@ -156,13 +156,33 @@ export const isStateful = () => { return event_queue.some((event) => event.name.startsWith("reflex___state")); }; +// Root-state field carrying RouterData; must match `constants.ROUTER + +// FIELD_MARKER` on the backend (see reflex.istate.data). +const ROUTER_FIELD = "router_rx_state_"; + /** * Apply a delta to the state. * @param state The state to apply the delta to. * @param delta The delta to apply. */ export const applyDelta = (state, delta) => { - return { ...state, ...delta }; + const new_state = { ...state, ...delta }; + // Once the connection-scoped router fields (session, headers) have been + // sent, the backend elides them from subsequent deltas; merge partial + // router payloads over the previously received value so they carry + // forward. The merge is field-agnostic, so adding RouterData fields on the + // backend needs no change here. + const router = delta[ROUTER_FIELD]; + const prev_router = state[ROUTER_FIELD]; + if ( + router !== null && + prev_router !== null && + typeof router === "object" && + typeof prev_router === "object" + ) { + new_state[ROUTER_FIELD] = { ...prev_router, ...router }; + } + return new_state; }; /** 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 4185cef9e99..36c10df48cf 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 @@ -12,7 +12,7 @@ from importlib.util import find_spec from typing import TYPE_CHECKING, Any -from reflex.istate.data import RouterData +from reflex.istate.data import RouterData, router_connection_scope from reflex.istate.manager.token import BaseStateToken from reflex.istate.proxy import StateProxy from reflex.utils import types @@ -358,11 +358,19 @@ async def _execute_event( # re-assign only when the value is set and different if router_data and state.router_data != router_data: + previous_router = state.router # assignment will recurse into substates and force recalculation of # dependent ComputedVar (dynamic route variables) state.router_data = router_data if state.router != (router := RouterData.from_router_data(router_data)): state.router = router + # When only the per-navigation fields changed, the delta + # can elide the connection-scoped fields the client + # already holds. Direct router writes elsewhere reset + # this flag (see BaseState.__setattr__). + state._router_static_unchanged = router_connection_scope( + previous_router + ) == router_connection_scope(router) # Preprocess the event. if ( diff --git a/packages/reflex-base/src/reflex_base/utils/exceptions.py b/packages/reflex-base/src/reflex_base/utils/exceptions.py index bbf29239edf..62f2c5c77ec 100644 --- a/packages/reflex-base/src/reflex_base/utils/exceptions.py +++ b/packages/reflex-base/src/reflex_base/utils/exceptions.py @@ -72,6 +72,10 @@ class VarNameError(ReflexError, NameError): """Custom NameError for when a state var has been shadowed by a substate var.""" +class ReservedStateFieldError(ReflexError, NameError): + """Raised when a state class declares a field name reserved for internal use.""" + + class VarTypeError(ReflexError, TypeError): """Custom TypeError for var related errors.""" diff --git a/packages/reflex-base/src/reflex_base/utils/types.py b/packages/reflex-base/src/reflex_base/utils/types.py index 1298287ec6e..3cd66ad03c7 100644 --- a/packages/reflex-base/src/reflex_base/utils/types.py +++ b/packages/reflex-base/src/reflex_base/utils/types.py @@ -163,7 +163,14 @@ def __call__( dict: Dict, # noqa: UP006 } -RESERVED_BACKEND_VAR_NAMES = {"_abc_impl", "_backend_vars", "_was_touched", "_mixin"} +RESERVED_BACKEND_VAR_NAMES = { + "_abc_impl", + "_backend_vars", + "_partial_router_capable", + "_router_static_unchanged", + "_was_touched", + "_mixin", +} class Unset: diff --git a/reflex/app.py b/reflex/app.py index 63ec53a4c75..eae8c8f3c25 100644 --- a/reflex/app.py +++ b/reflex/app.py @@ -1990,17 +1990,25 @@ async def on_connect(self, sid: str, environ: dict): self._token_manager.ensure_lost_and_found_task(self.emit_update) query_params = urllib.parse.parse_qs(environ.get("QUERY_STRING", "")) token_list = query_params.get("token", []) - if token_list: - await self.link_token_to_sid(sid, token_list[0]) - else: - logger.warning(f"No token provided in connection for session {sid}") - subprotocol = environ.get("HTTP_SEC_WEBSOCKET_PROTOCOL") if subprotocol and subprotocol != constants.Reflex.VERSION: logger.warning( f"Frontend version {subprotocol} for session {sid} does not match the backend version {constants.Reflex.VERSION}." ) + if token_list: + await self.link_token_to_sid( + sid, + token_list[0], + # Only a frontend compiled by this exact backend version is + # known to merge partial router deltas; anything else (e.g. a + # cached bundle from before a rolling deployment) gets the + # full router in every delta. + partial_router_capable=subprotocol == constants.Reflex.VERSION, + ) + else: + logger.warning(f"No token provided in connection for session {sid}") + def on_disconnect(self, sid: str) -> asyncio.Task | None: """Event for when the websocket disconnects. @@ -2232,12 +2240,17 @@ async def on_client_error(self, sid: str, data: Any): # handlers (e.g. error trackers) receive client errors too. self.app.frontend_exception_handler(Exception(report)) - async def link_token_to_sid(self, sid: str, token: str): + async def link_token_to_sid( + self, sid: str, token: str, partial_router_capable: bool = False + ): """Link a token to a session id. Args: sid: The Socket.IO session id. token: The client token. + partial_router_capable: Whether the connected frontend advertised + (via exact version match on the websocket subprotocol) that it + merges partial router deltas. """ # Use TokenManager for duplicate detection and Redis support new_token = await self._token_manager.link_token_to_sid(token, sid) @@ -2253,3 +2266,4 @@ async def link_token_to_sid(self, sid: str, token: str): ) as state: state.router_data[constants.RouteVar.SESSION_ID] = sid state.router = RouterData.from_router_data(state.router_data) + state._partial_router_capable = partial_router_capable diff --git a/reflex/istate/data.py b/reflex/istate/data.py index d0f044e54d0..d853d44682b 100644 --- a/reflex/istate/data.py +++ b/reflex/istate/data.py @@ -461,6 +461,26 @@ def from_router_data(cls, router_data: dict) -> "RouterData": ) +# RouterData fields that are fixed for the lifetime of a client connection. +# State deltas may omit these once the client has them (see +# serialize_partial_router_data); the event processor decides that by +# comparing exactly these fields, so adding one here is all that is needed to +# keep the two in step. +CONNECTION_SCOPED_ROUTER_FIELDS = ("session", "headers") + + +def router_connection_scope(obj: RouterData) -> tuple: + """Get the connection-scoped router values used to detect client-visible changes. + + Args: + obj: the RouterData object. + + Returns: + A tuple of the connection-scoped field values, in declaration order. + """ + return tuple(getattr(obj, name) for name in CONNECTION_SCOPED_ROUTER_FIELDS) + + @serializer(to=dict) def serialize_router_data(obj: RouterData) -> dict: """Serialize a RouterData object to a dict. @@ -472,8 +492,31 @@ def serialize_router_data(obj: RouterData) -> dict: A dict representation of the RouterData object. """ return { - "session": obj.session, - "headers": obj.headers, + **dict( + zip( + CONNECTION_SCOPED_ROUTER_FIELDS, + router_connection_scope(obj), + strict=True, + ) + ), + **serialize_partial_router_data(obj), + } + + +def serialize_partial_router_data(obj: RouterData) -> dict: + """Serialize only the per-navigation fields of a RouterData object. + + Used for state deltas once the connection-scoped fields (session and + headers) have already been sent to the client; the frontend merges this + partial payload over its previously received router value. + + Args: + obj: the RouterData object. + + Returns: + A dict with the per-navigation fields of the RouterData object. + """ + return { "page": obj._page, # ReflexURL is a str subclass, so json.dumps handles it natively and # never invokes the `default=serialize` hook. Call the URL serializer diff --git a/reflex/state.py b/reflex/state.py index a24b6f376d3..dafbc96c676 100644 --- a/reflex/state.py +++ b/reflex/state.py @@ -46,6 +46,7 @@ DynamicRouteArgShadowsStateVarError, EventHandlerShadowsBuiltInStateMethodError, ReflexRuntimeError, + ReservedStateFieldError, SetUndefinedStateVarError, StateMismatchError, StateSchemaMismatchError, @@ -73,7 +74,7 @@ import reflex.istate.dynamic from reflex import event from reflex.istate import HANDLED_PICKLE_ERRORS, debug_failed_pickles -from reflex.istate.data import RouterData +from reflex.istate.data import RouterData, serialize_partial_router_data from reflex.istate.proxy import ImmutableMutableProxy as ImmutableMutableProxy from reflex.istate.proxy import MutableProxy, is_mutable_type from reflex.istate.storage import ClientStorageBase @@ -441,6 +442,19 @@ class BaseState(EvenMoreBasicBaseState): # Whether the state has ever been touched since instantiation. _was_touched: bool = field(default=False, is_var=False) + # Whether the event processor's last router reassignment left the + # connection-scoped fields (session, headers) unchanged. Transient: + # recomputed on every reassignment, cleared by any direct router write, + # and never pickled. + _router_static_unchanged: bool = field(default=False, is_var=False) + + # Whether the connected client advertised, via an exact version match on + # the websocket subprotocol, that its applyDelta merges partial router + # payloads. Set on connect; False for older frontends (e.g. cached + # bundles during a rolling deployment), which then receive the full + # router in every delta. + _partial_router_capable: bool = field(default=False, is_var=False) + # A special event handler for setting base vars. setvar: ClassVar[EventHandler] @@ -547,6 +561,12 @@ def __init_subclass__(cls, mixin: bool = False, **kwargs): super().__init_subclass__(**kwargs) + # Internal router bookkeeping fields must not be redefined by user + # states: a shadowing value would silently control whether router + # deltas are sent partially. Checked before the mixin early-return so + # a mixin cannot smuggle the field into concrete states. + cls._check_reserved_internal_fields() + if cls._mixin: return @@ -962,6 +982,29 @@ def _check_overridden_methods(cls): msg = f"The event handler name `{method_name}` shadows a builtin State method; use a different name instead" raise EventHandlerShadowsBuiltInStateMethodError(msg) + _RESERVED_INTERNAL_FIELD_NAMES = frozenset({ + "_partial_router_capable", + "_router_static_unchanged", + }) + + @classmethod + def _check_reserved_internal_fields(cls): + """Check that internal bookkeeping fields are not redefined. + + Raises: + ReservedStateFieldError: When a state class declares a field + reserved for internal use. + """ + declared = set(inspect.get_annotations(cls)) | { + name for name in cls._RESERVED_INTERNAL_FIELD_NAMES if name in cls.__dict__ + } + for name in cls._RESERVED_INTERNAL_FIELD_NAMES & declared: + msg = ( + f"The field name `{name}` in {cls.__module__}.{cls.__name__} is " + "reserved for internal use; use a different name instead" + ) + raise ReservedStateFieldError(msg) + @classmethod def _check_overridden_basevars(cls): """Check for shadow base vars and raise error if any. @@ -1560,6 +1603,12 @@ def __setattr__(self, name: str, value: Any): self.dirty_vars.add(name) self._mark_dirty() + # Any direct router write invalidates the partial-router-delta + # optimization; the event processor re-arms it after comparing the + # connection-scoped fields (see BaseStateEventProcessor). + if name == constants.ROUTER: + object.__setattr__(self, "_router_static_unchanged", False) + def reset(self): """Reset all the base vars to their default values.""" # Reset the base vars. @@ -1885,6 +1934,20 @@ def get_delta(self) -> Delta: if not types.is_backend_base_variable(prop, type(self)) } + if ( + self.parent_state is None + and (router_field := constants.ROUTER + FIELD_MARKER) in subdelta + and self._router_static_unchanged + and self._partial_router_capable + ): + # The connection-scoped router fields (session, headers) this + # client already received are unchanged, so ship only the + # per-navigation fields; the frontend merges the partial payload + # over its previously received router value. + subdelta[router_field] = serialize_partial_router_data( + subdelta[router_field] + ) + if len(subdelta) > 0: delta[self.get_full_name()] = subdelta @@ -2075,6 +2138,10 @@ def __getstate__(self): state.pop("parent_state", None) state.pop("substates", None) state.pop("_was_touched", None) + # Transient, request-scoped: recomputed by the event processor on every + # router reassignment. Persisting it could arm a partial router delta + # for a client that never received the connection-scoped fields. + state.pop("_router_static_unchanged", None) # Remove all inherited vars. for inherited_var_name in self.inherited_vars: state.pop(inherited_var_name, None) diff --git a/tests/units/test_app.py b/tests/units/test_app.py index 484796337a5..47c56e83705 100644 --- a/tests/units/test_app.py +++ b/tests/units/test_app.py @@ -21,6 +21,7 @@ import reflex_base from pytest_mock import MockerFixture from reflex_base.components.component import Component +from reflex_base.constants import RouteVar from reflex_base.constants.state import FIELD_MARKER from reflex_base.event import Event from reflex_base.event.context import EventContext @@ -586,6 +587,122 @@ def set_int_val(self, value: int): ] +@pytest.mark.asyncio +async def test_router_delta_partial_only_when_connection_scope_unchanged( + test_state: type[ATestState], + mock_base_state_event_processor: BaseStateEventProcessor, + mock_root_event_context: EventContext, + emitted_deltas: list[tuple[str, dict[str, dict[str, Any]]]], + token: str, + clean_registration_context: RegistrationContext, + router_data: dict[str, str | dict], +): + """The processor must elide session/headers only while they are unchanged. + + Drives the real comparison in the event processor rather than setting the + internal flag directly, so arming a partial payload after the session + actually changed would fail here. + + Args: + test_state: State Fixture. + mock_base_state_event_processor: BaseStateEventProcessor Fixture. + mock_root_event_context: The mock event context (for direct state access). + emitted_deltas: List to store emitted deltas. + token: a Token. + clean_registration_context: The registration context fixture. + router_data: The router data fixture. + """ + clean_registration_context.register_base_state(test_state) + router_field = constants.ROUTER + FIELD_MARKER + + state = test_state() # pyright: ignore [reportCallIssue] + state.add_var("nav_val", int, 0) + + def set_nav_val(self, value: int): + self.nav_val = value + + state._add_event_handler("set_nav_val", set_nav_val) + + def _event(rd: dict) -> Event: + return Event( + name=f"{test_state.get_name()}.set_nav_val", + payload={"value": 1}, + router_data=rd, + ) + + async def _router_delta(rd: dict): + emitted_deltas.clear() + async with mock_base_state_event_processor as processor: + await processor.enqueue(token, _event(rd)) + await processor.join() + for _tok, delta in emitted_deltas: + for substate in delta.values(): + if router_field in substate: + return substate[router_field] + return None + + # Before the connect handler has verified the client version, deltas are + # always full — even across repeated navigations with an unchanged + # session. This is what a cached pre-upgrade frontend receives during a + # rolling deployment. + legacy_first = await _router_delta(router_data) + assert isinstance(legacy_first, RouterData) + legacy_second = await _router_delta({**router_data, RouteVar.PATH: "/legacy"}) + assert isinstance(legacy_second, RouterData) + + # The connect handler verified an exact version match: partial deltas + # allowed from here on. + root = await mock_root_event_context.state_manager.get_state( + BaseStateToken(ident=token, cls=test_state) + ) + root._partial_router_capable = True + + # First navigation after capability: session/headers changed relative to + # nothing? No — they were already sent above; unchanged, so partial. + first = await _router_delta({**router_data, RouteVar.PATH: "/armed"}) + assert not isinstance(first, RouterData) + + # Same session and headers, different page: connection-scoped fields are + # elided. + same_connection = await _router_delta({**router_data, RouteVar.PATH: "/second"}) + assert same_connection is not None + assert not isinstance(same_connection, RouterData) + assert set(same_connection) == {"page", "url", "route_id"} + + # A new session id must fall back to the full payload, otherwise the + # client would keep serving the old token forever. + reconnected_sid = "a-different-session-id" + new_session = await _router_delta({ + **router_data, + RouteVar.PATH: "/third", + RouteVar.SESSION_ID: reconnected_sid, + }) + assert isinstance(new_session, RouterData), ( + "session changed but the delta omitted it; the client would keep the stale value" + ) + + # Changed headers with an unchanged session must also fall back to the + # full payload — headers are the other half of the connection scope. The + # state now holds `reconnected_sid` (assigned by the previous event), so + # reusing it here means the session compares equal and ONLY the headers + # differ; the mutation check for this leg is dropping "headers" from + # CONNECTION_SCOPED_ROUTER_FIELDS, which makes this assertion fail. + original_headers = router_data[RouteVar.HEADERS] + assert isinstance(original_headers, dict) + new_headers = await _router_delta({ + **router_data, + RouteVar.PATH: "/fourth", + RouteVar.SESSION_ID: reconnected_sid, + RouteVar.HEADERS: { + **original_headers, + "user-agent": "A Different Agent", + }, + }) + assert isinstance(new_headers, RouterData), ( + "headers changed but the delta omitted them; the client would keep stale headers" + ) + + @pytest.fixture def list_mutation_state(): """Create a state with list mutation features. @@ -1919,6 +2036,8 @@ def _dynamic_state_event(name, val, **kwargs): name=f"{OnLoadInternalState.get_full_name()}.{constants.CompileVars.ON_LOAD_INTERNAL.rpartition('.')[2]}", val=exp_val, ) + # This client never advertises partial-router capability (no connect + # handshake in this flow), so deltas always carry the full router. exp_router = RouterData.from_router_data(on_load_internal.router_data) async with mock_base_state_event_processor as processor: await processor.enqueue( diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 416e09ade65..94a4020ef43 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -32,6 +32,7 @@ InvalidLockWarningThresholdError, LockExpiredError, ReflexRuntimeError, + ReservedStateFieldError, SetUndefinedStateVarError, StateSerializationError, UnretrievableVarValueError, @@ -42,7 +43,14 @@ import reflex as rx from reflex.app import App from reflex.environment import environment -from reflex.istate.data import HeaderData, RouterData, _FrozenDictStrStr +from reflex.istate.data import ( + HeaderData, + RouterData, + _FrozenDictStrStr, + router_connection_scope, + serialize_partial_router_data, + serialize_router_data, +) from reflex.istate.manager import StateManager from reflex.istate.manager.disk import StateManagerDisk from reflex.istate.manager.memory import StateManagerMemory @@ -1086,6 +1094,118 @@ def test_get_client_ip(test_state, router_data): assert test_state.router.session.client_ip == "127.0.0.1" +def test_partial_router_delta(test_state, router_data): + """get_delta ships a partial router payload only when flag and capability agree. + + The comparison that arms `_router_static_unchanged` is driven end to end + through the event processor in + ``test_router_delta_partial_only_when_connection_scope_unchanged``; this + test treats the flag and the client capability as givens and pins down + ``get_delta``'s serialization for each combination. + + Args: + test_state: A state. + router_data: The router data fixture. + """ + router_field = constants.ROUTER + FIELD_MARKER + + def router_delta_value(*, unchanged: bool, capable: bool, path: str): + test_state.router = RouterData.from_router_data({ + **router_data, + RouteVar.PATH: path, + }) + # The direct write above cleared the flag; apply the scenario. + test_state._router_static_unchanged = unchanged + test_state._partial_router_capable = capable + value = test_state.get_delta()[test_state.get_full_name()][router_field] + test_state._clean() + return value + + # Unchanged connection scope + capable client: partial payload. + value = router_delta_value(unchanged=True, capable=True, path="/partial") + assert set(value) == {"page", "url", "route_id"} + assert value["route_id"] == "/partial" + + # Capable client, but the connection scope changed: full payload. + value = router_delta_value(unchanged=False, capable=True, path="/changed") + assert isinstance(value, RouterData) + + # Unchanged scope, but the client never advertised support (e.g. a cached + # pre-upgrade bundle during a rolling deployment): full payload. + value = router_delta_value(unchanged=True, capable=False, path="/legacy") + assert isinstance(value, RouterData) + + # Any direct router write clears the flag, so a write after arming falls + # back to the full payload without help from the processor. + test_state._router_static_unchanged = True + test_state._partial_router_capable = True + test_state.router = RouterData.from_router_data({ + **router_data, + RouteVar.PATH: "/direct", + }) + value = test_state.get_delta()[test_state.get_full_name()][router_field] + assert isinstance(value, RouterData) + test_state._clean() + + +def test_reserved_internal_router_fields_cannot_be_redefined(): + """User states must not shadow the internal router bookkeeping fields. + + A shadowing value would silently control whether router deltas are sent + partially, so redefinition raises instead. + """ + for name in ("_router_static_unchanged", "_partial_router_capable"): + with pytest.raises(ReservedStateFieldError): + type( + "ShadowingState", + (BaseState,), + { + "__module__": __name__, + "__qualname__": "ShadowingState", + "__annotations__": {name: bool}, + name: True, + }, + ) + # Mixins are checked at definition too, so the field cannot be + # smuggled into concrete states through a mixin base. + with pytest.raises(ReservedStateFieldError): + type( + "ShadowingMixin", + (BaseState,), + { + "__module__": __name__, + "__qualname__": "ShadowingMixin", + "__annotations__": {name: bool}, + name: True, + }, + mixin=True, + ) + + +def test_partial_router_delta_covers_every_elided_field(router_data): + """Every field a navigation delta omits must be one that is compared. + + A field that is elided but not compared would leave a stale value on the + client forever, so changing any omitted field must be visible to the + processor's comparison. + + Args: + router_data: The router data fixture. + """ + router = RouterData.from_router_data(router_data) + omitted = set(serialize_router_data(router)) - set( + serialize_partial_router_data(router) + ) + assert omitted, "partial payload must omit something, else there is no saving" + + for omitted_field in omitted: + changed = dataclasses.replace(router, **{omitted_field: None}) + assert router_connection_scope(changed) != router_connection_scope(router), ( + f"router field {omitted_field!r} is omitted from navigation deltas but is" + " not compared, so a change to it would never reach the client" + ) + + def test_get_current_page(test_state): assert test_state.router._page.path == ""