diff --git a/packages/reflex-base/news/6181.performance.md b/packages/reflex-base/news/6181.performance.md new file mode 100644 index 00000000000..6c431a63ac9 --- /dev/null +++ b/packages/reflex-base/news/6181.performance.md @@ -0,0 +1 @@ +Reduce browser rendering overhead for state updates, especially in apps with many substates. diff --git a/packages/reflex-base/src/reflex_base/compiler/templates.py b/packages/reflex-base/src/reflex_base/compiler/templates.py index 23ad07d4fad..d4b3af8ebc5 100644 --- a/packages/reflex-base/src/reflex_base/compiler/templates.py +++ b/packages/reflex-base/src/reflex_base/compiler/templates.py @@ -388,18 +388,8 @@ def context_template( """ ) - state_reducer_str = "\n".join( - rf'const [{format_state_name(state_name)}, dispatch_{format_state_name(state_name)}] = useReducer(applyDelta, initialState["{state_name}"])' - for state_name in initial_state - ) - create_state_contexts_str = "\n".join( - rf"createElement(StateContexts.{format_state_name(state_name)},{{value: {format_state_name(state_name)}}}," - for state_name in initial_state - ) - - dispatchers_str = "\n".join( - f'"{state_name}": dispatch_{format_state_name(state_name)},' + rf"createElement(SubstateProvider, {{substateName: '{state_name}', contextName: '{format_state_name(state_name)}'}}," for state_name in initial_state ) @@ -431,7 +421,7 @@ def context_template( else "" ) - return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useReducer, useState, createElement, useEffect }} from "react" + return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useReducer, useRef, useState, createElement, useEffect, useLayoutEffect }} from "react" import {{ applyDelta, ReflexEvent, hydrateClientStorage, useEventLoop, refs }} from "$/utils/state" import {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, getStateContext, registerApp, eventLoop }} from "$/utils/context-registry" import {{ jsx }} from "@emotion/react"; @@ -518,19 +508,57 @@ def context_template( ); }} -export function StateProvider({{ children }}) {{ - {state_reducer_str} - const dispatchers = useMemo(() => {{ - return {{ - {dispatchers_str} - }} - }}, []) +// ``useLayoutEffect`` warns when rendered on the server, where no effect runs +// at all, so fall back to ``useEffect`` there. +const useIsomorphicLayoutEffect = + typeof document !== "undefined" ? useLayoutEffect : useEffect; + +// Holds the mutable substate -> dispatch registry that ``SubstateProvider`` +// writes into and ``EventLoopProvider`` reads. The registry object identity is +// stable for the lifetime of the tree, so neither adding a dispatcher nor +// updating a substate re-renders the consumers of ``DispatchContext``. +const DispatchProvider = ({{ children }}) => {{ + const dispatchers = useRef({{}}); + return useMemo( + () => + createElement(DispatchContext, {{ value: dispatchers.current }}, children), + [children], + ); +}}; + +// One provider per substate: each owns its own reducer, so a delta for one +// substate only re-renders its provider instead of recreating every provider. +const SubstateProvider = ({{ children, substateName, contextName }}) => {{ + const dispatchers = useContext(DispatchContext); + const [state, dispatchSubstate] = useReducer( + applyDelta, + initialState[substateName], + ); + // A layout effect, not a passive one: layout effects for the whole commit + // run before any passive effect, so every dispatcher is registered before + // ``EventLoopProvider`` (mounted below this provider) connects the socket. + // A delta naming an unregistered substate is a fatal state mismatch. + useIsomorphicLayoutEffect(() => {{ + dispatchers[substateName] = dispatchSubstate; + return () => {{ + delete dispatchers[substateName]; + }}; + }}, [dispatchers, dispatchSubstate, substateName]); + return useMemo( + () => createElement(StateContexts[contextName], {{ value: state }}, children), + [children, state, contextName], + ); +}}; - return ( - {create_state_contexts_str} - createElement(DispatchContext, {{value: dispatchers}}, children) +export function StateProvider({{ children }}) {{ + return useMemo( + () => ( + createElement(DispatchProvider, {{}}, + {create_state_contexts_str}children {")" * len(initial_state)} - ) + )), + [children], + ); }}""" diff --git a/tests/integration/tests_playwright/test_router_query.py b/tests/integration/tests_playwright/test_router_query.py index 50b83b94b25..96a70c27d34 100644 --- a/tests/integration/tests_playwright/test_router_query.py +++ b/tests/integration/tests_playwright/test_router_query.py @@ -24,6 +24,8 @@ import pytest from playwright.sync_api import Page, expect +from reflex_base.constants.state import FIELD_MARKER +from socketio.packet import Packet from reflex.testing import AppHarness @@ -32,6 +34,11 @@ def RouterQueryApp(): """App exercising replaceState, redirect, and redirect(replace=True).""" import reflex as rx + class ConnectedState(rx.State): + """A sibling updated by the first on-load event after connecting.""" + + connected_count: int = 0 + class RouterQueryState(rx.State): # Incremented by the page on_load handler; proves whether a navigation # (and thus on_load) actually fired. @@ -59,9 +66,11 @@ def name_param(self) -> str: return self.router.url.query_parameters.get("name", "") @rx.event - def on_load(self): - """Record that the page on_load handler fired.""" + async def on_load(self): + """Update both substates in the same on-load delta.""" self.load_count += 1 + connected = await self.get_state(ConnectedState) + connected.connected_count += 1 @rx.event def replace_via_script(self, query: str): @@ -141,6 +150,11 @@ def index(): read_only=True, id="load-count", ), + rx.input( + value=f"{ConnectedState.connected_count}", + read_only=True, + id="connected-count", + ), rx.input( value=f"{RouterQueryState.ping_count}", read_only=True, @@ -196,6 +210,44 @@ def _load(harness: AppHarness, page: Page) -> str: return base +def test_initial_connection_dispatches_both_substates( + router_query_app: AppHarness, page: Page +): + """The first on-load delta updates both sibling substates in the browser. + + Args: + router_query_app: Running application in dev or production mode. + page: Browser page receiving the initial connection's state updates. + """ + frames: list[str | bytes] = [] + page.on( + "websocket", + lambda socket: socket.on("framereceived", lambda frame: frames.append(frame)), + ) + + _load(router_query_app, page) + expect(page.locator("#connected-count")).to_have_value("1") + + # Socket.IO event packets use the Engine.IO message and event prefixes, 42. + for frame in frames: + if not isinstance(frame, str) or not frame.startswith("42"): + continue + packet = Packet(encoded_packet=frame[1:]) + assert packet.data is not None + event, update = packet.data + if event != "event": + continue + updated_substates = sum( + fields.get("load_count" + FIELD_MARKER) == 1 + or fields.get("connected_count" + FIELD_MARKER) == 1 + for fields in update.get("delta", {}).values() + ) + if updated_substates == 2: + break + else: + pytest.fail("The initial on-load update did not contain both substates") + + def test_replace_state_is_not_reactive_but_next_event_syncs( router_query_app: AppHarness, page: Page ): diff --git a/tests/units/compiler/test_compiler.py b/tests/units/compiler/test_compiler.py index 07307221829..446d6cea424 100644 --- a/tests/units/compiler/test_compiler.py +++ b/tests/units/compiler/test_compiler.py @@ -1669,6 +1669,39 @@ def test_context_template_owner_stack_pin(disable_owner_stacks: bool): assert "captureOwnerStack" in rendered +def test_context_template_one_provider_per_substate(): + """Each substate gets its own provider so one delta re-renders one context. + + A single provider owning every reducer means any delta recreates every + ``StateContexts`` element; nesting one ``SubstateProvider`` per substate + keeps the untouched providers memoized. + """ + from reflex_base.compiler.templates import context_template + + rendered = context_template( + is_dev_mode=True, + default_color_mode='"light"', + initial_state={ + "reflex___state____state": {}, + "reflex___state____state__sub": {}, + }, + state_name="reflex___state____state", + ) + + assert ( + "createElement(SubstateProvider, {substateName: 'reflex___state____state', " + "contextName: 'reflex___state____state'}," in rendered + ) + assert ( + "createElement(SubstateProvider, {substateName: 'reflex___state____state__sub', " + "contextName: 'reflex___state____state__sub'}," in rendered + ) + # The reducers moved into SubstateProvider; StateProvider only composes. + provider_body = rendered[rendered.index("export function StateProvider") :] + assert "useReducer" not in provider_body + assert "createElement(DispatchProvider, {}," in provider_body + + def test_context_template_client_side_component_is_named(): """``ClientSide`` returns a named component, not an anonymous arrow.""" from reflex_base.compiler.templates import context_template