perf: pipeline redis set_state writes, cache required-state classes

news/6745.performance.md

@@ -0,0 +1 @@
+The redis state manager now verifies the lock once and writes all touched states in a single pipeline (3 round trips per event flush instead of ~2N+1), serializes large states off the event loop, caches the per-class required-state computation, and no longer re-writes unchanged states on subsequent flushes.

reflex/istate/manager/redis.py

@@ -3,6 +3,7 @@
 import asyncio
 import contextlib
 import dataclasses
+import functools
 import inspect
 import os
 import sys
@@ -15,6 +16,7 @@
 from redis.asyncio import Redis
 from reflex_base.config import get_config
 from reflex_base.environment import environment
+from reflex_base.registry import RegistrationContext
 from reflex_base.utils import console
 from reflex_base.utils.exceptions import (
     InvalidLockWarningThresholdError,
@@ -29,7 +31,7 @@
     _default_token_expiration,
 )
 from reflex.istate.manager.token import TOKEN_TYPE, BaseStateToken, StateToken
-from reflex.state import BaseState
+from reflex.state import BaseState, _get_state_hierarchy_generation
 from reflex.utils.tasks import ensure_task
 
 NOTIFY_KEYSPACE_EVENTS = (
@@ -98,6 +100,84 @@ def _default_oplock_hold_time_ms() -> int:
 SMR = f"[SMR:{os.getpid()}]"
 start = time.monotonic()
 
+# Serialize states off the event loop when their previous pickle for the same
+# state name exceeded this size (bytes).
+_OFFLOAD_SERIALIZE_THRESHOLD = 64 * 1024
+
+
+@functools.lru_cache(maxsize=1024)
+def _required_state_classes(
+    target_state_cls: type[BaseState],
+    registration_ctx: RegistrationContext,
+    generation: int,
+) -> frozenset[type[BaseState]]:
+    """Get the state classes required to fetch the target state, cached.
+
+    The result is derived purely from class-level data, so entries are keyed
+    on the active registration context and the state hierarchy generation:
+    dynamic state class creation or a module hot reload starts a new
+    generation, orphaning stale entries.
+
+    Args:
+        target_state_cls: The target state class being fetched.
+        registration_ctx: The registration context the hierarchy lives in.
+        generation: The state hierarchy generation the result is valid for.
+
+    Returns:
+        The set of state classes required to fetch the target state.
+    """
+    return frozenset(_collect_required_state_classes(target_state_cls, subclasses=True))
+
+
+def _collect_required_state_classes(
+    target_state_cls: type[BaseState],
+    subclasses: bool = False,
+    required_state_classes: set[type[BaseState]] | None = None,
+) -> set[type[BaseState]]:
+    """Recursively determine which states are required to fetch the target state.
+
+    This will always include potentially dirty substates that depend on vars
+    in the target_state_cls.
+
+    Args:
+        target_state_cls: The target state class being fetched.
+        subclasses: Whether to include subclasses of the target state.
+        required_state_classes: Recursive argument tracking state classes that have already been seen.
+
+    Returns:
+        The set of state classes required to fetch the target state.
+    """
+    if required_state_classes is None:
+        required_state_classes = set()
+    # Get the substates if requested.
+    if subclasses:
+        for substate in target_state_cls.get_substates():
+            _collect_required_state_classes(
+                substate,
+                subclasses=True,
+                required_state_classes=required_state_classes,
+            )
+    if target_state_cls in required_state_classes:
+        return required_state_classes
+    required_state_classes.add(target_state_cls)
+
+    # Get dependent substates.
+    for pd_substates in target_state_cls._get_potentially_dirty_states():
+        _collect_required_state_classes(
+            pd_substates,
+            subclasses=False,
+            required_state_classes=required_state_classes,
+        )
+
+    # Get the parent state if it exists.
+    if parent_state := target_state_cls.get_parent_state():
+        _collect_required_state_classes(
+            parent_state,
+            subclasses=False,
+            required_state_classes=required_state_classes,
+        )
+    return required_state_classes
+
 
 class RedisPubSubMessage(TypedDict):
     """A Redis Pub/Sub message."""
@@ -162,6 +242,12 @@ class StateManagerRedis(StateManager):
         default_factory=environment.REFLEX_OPLOCK_ENABLED.get, init=False
     )
 
+    # Size of the last serialized payload per state full name, used to decide
+    # when to offload pickling to a thread.
+    _last_serialized_size: dict[str, int] = dataclasses.field(
+        default_factory=dict, init=False
+    )
+
     # Cached states
     _cached_states: dict[str, Any] = dataclasses.field(default_factory=dict, init=False)
     _cached_states_locks: dict[str, asyncio.Lock] = dataclasses.field(
@@ -208,56 +294,6 @@ def __post_init__(self):
             asyncio.get_running_loop()  # Check if we're in an event loop.
             self._ensure_lock_task()
 
-    def _get_required_state_classes(
-        self,
-        target_state_cls: type[BaseState],
-        subclasses: bool = False,
-        required_state_classes: set[type[BaseState]] | None = None,
-    ) -> set[type[BaseState]]:
-        """Recursively determine which states are required to fetch the target state.
-
-        This will always include potentially dirty substates that depend on vars
-        in the target_state_cls.
-
-        Args:
-            target_state_cls: The target state class being fetched.
-            subclasses: Whether to include subclasses of the target state.
-            required_state_classes: Recursive argument tracking state classes that have already been seen.
-
-        Returns:
-            The set of state classes required to fetch the target state.
-        """
-        if required_state_classes is None:
-            required_state_classes = set()
-        # Get the substates if requested.
-        if subclasses:
-            for substate in target_state_cls.get_substates():
-                self._get_required_state_classes(
-                    substate,
-                    subclasses=True,
-                    required_state_classes=required_state_classes,
-                )
-        if target_state_cls in required_state_classes:
-            return required_state_classes
-        required_state_classes.add(target_state_cls)
-
-        # Get dependent substates.
-        for pd_substates in target_state_cls._get_potentially_dirty_states():
-            self._get_required_state_classes(
-                pd_substates,
-                subclasses=False,
-                required_state_classes=required_state_classes,
-            )
-
-        # Get the parent state if it exists.
-        if parent_state := target_state_cls.get_parent_state():
-            self._get_required_state_classes(
-                parent_state,
-                subclasses=False,
-                required_state_classes=required_state_classes,
-            )
-        return required_state_classes
-
     def _get_populated_states(
         self,
         target_state: BaseState,
@@ -322,7 +358,11 @@ async def get_state(
 
         # Determine which states from the tree need to be fetched.
         required_state_classes = sorted(
-            self._get_required_state_classes(requested_state_cls, subclasses=True)
+            _required_state_classes(
+                requested_state_cls,
+                RegistrationContext.get(),
+                _get_state_hierarchy_generation(),
+            )
             - {type(s) for s in flat_state_tree.values()},
             key=lambda x: x.get_full_name(),
         )
@@ -441,32 +481,45 @@ async def set_state(
                     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()
+        # Walk the substate tree and collect the states that need writing.
+        touched_states: list[BaseState] = []
+        stack = [base_state]
+        while stack:
+            substate = stack.pop()
+            stack.extend(substate.substates.values())
+            if substate._get_was_touched():
+                touched_states.append(substate)
+        if not touched_states:
+            return
+
+        # Persist only each touched state (parents and substates are excluded
+        # by BaseState.__getstate__) in a single pipelined round trip.
+        pipeline = self.redis.pipeline()
+        written_states: list[BaseState] = []
+        for substate in touched_states:
+            full_name = substate.get_full_name()
+            if (
+                self._last_serialized_size.get(full_name, 0)
+                > _OFFLOAD_SERIALIZE_THRESHOLD
+            ):
+                # Pickling large states would stall the event loop.
+                pickle_state = await asyncio.to_thread(substate._serialize)
+            else:
+                pickle_state = substate._serialize()
             if pickle_state:
-                await self.redis.set(
-                    str(token.with_cls(type(base_state))),
+                self._last_serialized_size[full_name] = len(pickle_state)
+                pipeline.set(
+                    str(token.with_cls(type(substate))),
                     pickle_state,
                     ex=self.token_expiration,
                 )
-
-        # Wait for substates to be persisted.
-        for t in tasks:
-            await t
+                written_states.append(substate)
+        if written_states:
+            await pipeline.execute()
+            # Reset the touched flag so an unchanged state is not re-pickled
+            # and re-written by subsequent flushes of the same instance.
+            for substate in written_states:
+                substate._was_touched = False
 
     @contextlib.asynccontextmanager
     async def _try_modify_state(

reflex/state.py

@@ -357,6 +357,26 @@ def _is_user_descriptor(value: Any) -> bool:
 
 all_base_state_classes: dict[str, None] = {}
 
+# Bumped whenever the state class hierarchy changes (subclass creation or
+# module reload), invalidating caches derived from class-level data.
+_state_hierarchy_generation: int = 0
+
+
+def _get_state_hierarchy_generation() -> int:
+    """Get the current state hierarchy generation.
+
+    Returns:
+        A counter that increases whenever the state class hierarchy changes.
+    """
+    return _state_hierarchy_generation
+
+
+def _bump_state_hierarchy_generation() -> None:
+    """Invalidate caches keyed on the state class hierarchy."""
+    global _state_hierarchy_generation
+    _state_hierarchy_generation += 1
+
+
 CLASS_VAR_NAMES = frozenset({
     "vars",
     "base_vars",
@@ -720,6 +740,7 @@ def __init_subclass__(cls, mixin: bool = False, **kwargs):
         cls._init_var_dependency_dicts()
 
         all_base_state_classes[cls.get_full_name()] = None
+        _bump_state_hierarchy_generation()
 
     @classmethod
     def _add_event_handler(
@@ -2682,3 +2703,4 @@ def reload_state_module(
             state._var_dependencies = {}
             state._init_var_dependency_dicts()
     state.get_class_substate.cache_clear()
+    _bump_state_hierarchy_generation()