From 9cd59d2f22ecc7527b2da56594395a34e878170a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 20:48:43 +0000 Subject: [PATCH 1/2] perf: pipeline redis set_state writes, cache required-state classes set_state recursed over the substate tree spawning a task per substate; every recursive call re-GETed the same lock key and re-PTTLed it, and each touched state got its own unpipelined SET -- roughly 2N+1 round trips per event flush for N states. The tree is now walked once, the lock is verified once, and all touched states are written in a single pipeline: 3 round trips per flush regardless of tree size (12->3 for a root with 3 substates, 39->3 for a 13-state tree). Also: - Pickling runs off the event loop (asyncio.to_thread) for states whose previous payload exceeded 64KiB, instead of stalling the loop. - The recursive required-state-classes computation (pure class-level data) is now lru-cached per registration context and state hierarchy generation; subclass creation and module hot reload bump the generation, so dynamic classes invalidate cleanly. - _was_touched is reset after a successful write (mirroring the disk manager), so a once-touched state held by the oplock cache is no longer re-pickled and re-SET on every subsequent flush. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Sns7EuBPYvfGCxRSRZ1bUx --- reflex/istate/manager/redis.py | 201 ++++++++++++++--------- reflex/state.py | 22 +++ tests/units/istate/manager/test_redis.py | 152 +++++++++++++++++ 3 files changed, 301 insertions(+), 74 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index c2dccc50c9b..48f3e449a92 100644 --- a/reflex/istate/manager/redis.py +++ b/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( diff --git a/reflex/state.py b/reflex/state.py index 3762ad0ac2a..d9460cf9b80 100644 --- a/reflex/state.py +++ b/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() diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 0b37389a337..9695dfa4736 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -736,3 +736,155 @@ async def modify(): ) assert isinstance(final_state, root_state) assert final_state.count == 2 + + +async def test_set_state_verifies_lock_once_and_pipelines_writes( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """set_state checks the lock once and writes all states in one pipeline. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + """ + state_manager_redis._oplock_enabled = False + redis = state_manager_redis.redis + token = str(uuid.uuid4()) + + get_calls: list[Any] = [] + pttl_calls: list[Any] = [] + orig_get = redis.get + orig_pttl = redis.pttl + + async def counting_get(key, *args, **kwargs): + get_calls.append(key) + return await orig_get(key, *args, **kwargs) + + async def counting_pttl(key, *args, **kwargs): + pttl_calls.append(key) + return await orig_pttl(key, *args, **kwargs) + + redis.get = counting_get + redis.pttl = counting_pttl + try: + async with state_manager_redis.modify_state( + BaseStateToken(ident=token, cls=root_state) + ) as state: + state.count = 1 + (await state.get_state(SubState1)).foo = "sub1" + (await state.get_state(SubState2)).foo = "sub2" + get_calls.clear() + pttl_calls.clear() + finally: + redis.get = orig_get + redis.pttl = orig_pttl + + # With three touched states, the old implementation issued one lock GET + # and one PTTL per state in the tree; now the lock is verified once. + assert len(get_calls) == 1 + assert len(pttl_calls) == 1 + + final_state = await state_manager_redis.get_state( + BaseStateToken(ident=token, cls=root_state) + ) + assert final_state.count == 1 + assert (await final_state.get_state(SubState1)).foo == "sub1" + assert (await final_state.get_state(SubState2)).foo == "sub2" + + +async def test_set_state_resets_touched_flag( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """A state written to redis is not re-written by a subsequent flush. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + """ + state_manager_redis._oplock_enabled = False + token = str(uuid.uuid4()) + state_token = BaseStateToken(ident=token, cls=root_state) + + state = await state_manager_redis.get_state(state_token) + state.count = 5 + state._clean() + assert state._get_was_touched() + await state_manager_redis.set_state(state_token, state) + assert not state._was_touched + + # Flushing the same instance again should not re-write unchanged states, + # so no write pipeline is even created. + pipelines_created = 0 + orig_pipeline = state_manager_redis.redis.pipeline + + def counting_pipeline(*args, **kwargs): + nonlocal pipelines_created + pipelines_created += 1 + return orig_pipeline(*args, **kwargs) + + state_manager_redis.redis.pipeline = counting_pipeline + try: + await state_manager_redis.set_state(state_token, state) + finally: + state_manager_redis.redis.pipeline = orig_pipeline + assert pipelines_created == 0 + + final_state = await state_manager_redis.get_state(state_token) + assert final_state.count == 5 + + +async def test_set_state_offloads_large_pickles( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + mocker, +): + """States whose previous pickle was large are serialized off the loop. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + mocker: The pytest-mock fixture. + """ + state_manager_redis._oplock_enabled = False + token = str(uuid.uuid4()) + state_token = BaseStateToken(ident=token, cls=root_state) + + state = await state_manager_redis.get_state(state_token) + state.count = 7 + state_manager_redis._last_serialized_size[root_state.get_full_name()] = 10**9 + + to_thread = mocker.patch( + "reflex.istate.manager.redis.asyncio.to_thread", wraps=asyncio.to_thread + ) + await state_manager_redis.set_state(state_token, state) + assert to_thread.called + + final_state = await state_manager_redis.get_state(state_token) + assert final_state.count == 7 + + +def test_required_state_classes_cache_invalidation( + root_state: type[RedisTestState], +): + """The cached required-state-classes set follows hierarchy changes. + + Args: + root_state: The root state class. + """ + from reflex_base.registry import RegistrationContext + + from reflex.istate.manager.redis import _required_state_classes + from reflex.state import _get_state_hierarchy_generation + + ctx = RegistrationContext.get() + before = _required_state_classes(root_state, ctx, _get_state_hierarchy_generation()) + assert SubState1 in before + + class LateSubState(root_state): # pyright: ignore[reportGeneralTypeIssues] + """A substate defined after the cache was warmed.""" + + after = _required_state_classes(root_state, ctx, _get_state_hierarchy_generation()) + assert LateSubState not in before + assert LateSubState in after From c4b7b468c674e1be42545f41eb797a5b90234335 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 20:49:19 +0000 Subject: [PATCH 2/2] chore: add news fragment for #6745 Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Sns7EuBPYvfGCxRSRZ1bUx --- news/6745.performance.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 news/6745.performance.md diff --git a/news/6745.performance.md b/news/6745.performance.md new file mode 100644 index 00000000000..e5ad1580920 --- /dev/null +++ b/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.