Faster first page load: one-lock hydrate from the connect packet, diffed against compiled defaults

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.

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.

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") {

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({{

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.

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)