From 24b95564da915f54a420cb2164967d8c040abb9f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 20:48:43 +0000 Subject: [PATCH 1/5] 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 09e56802650..27e4eac93da 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", @@ -739,6 +759,7 @@ def __init_subclass__(cls, mixin: bool = False, **kwargs): cls._refresh_interval_computed_vars() all_base_state_classes[cls.get_full_name()] = None + _bump_state_hierarchy_generation() @classmethod def _refresh_interval_computed_vars(cls) -> None: @@ -2755,3 +2776,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 d95422ba8323505d9a957a96273959a80399c082 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 10 Jul 2026 20:49:19 +0000 Subject: [PATCH 2/5] 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. From 591f7b854c698302d89b940a356b2a8e738f3fcf Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Fri, 17 Jul 2026 18:27:14 -0700 Subject: [PATCH 3/5] Expect a single held-too-long warning now that set_state is pipelined --- tests/units/test_state.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/units/test_state.py b/tests/units/test_state.py index 35fec2068f6..effbce7f833 100644 --- a/tests/units/test_state.py +++ b/tests/units/test_state.py @@ -2250,7 +2250,9 @@ async def _coro_blocker(): console_warn.assert_not_called() else: console_warn.assert_called() - assert console_warn.call_count == 7 + # set_state persists the whole tree in one pipelined call, so the + # held-too-long warning fires once instead of once per substate. + assert console_warn.call_count == 1 class CopyingAsyncMock(AsyncMock): From 66c2c99f00caa6aa3cca7f3fcffb24d3cb3d5d29 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Fri, 17 Jul 2026 18:32:48 -0700 Subject: [PATCH 4/5] Fix pipelined set_state test: distinct substate vars, self-contained mock getdel --- tests/units/istate/manager/test_redis.py | 12 ++++++++---- tests/units/mock_redis.py | 12 +++++++++--- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 9695dfa4736..5988e2876ea 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -26,10 +26,14 @@ class RedisTestState(BaseState): class SubState1(RedisTestState): """A test substate for redis state manager tests.""" + sub1_foo: str = "" + class SubState2(RedisTestState): """A test substate for redis state manager tests.""" + sub2_foo: str = "" + @pytest.fixture def root_state() -> type[RedisTestState]: @@ -772,8 +776,8 @@ async def counting_pttl(key, *args, **kwargs): 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" + (await state.get_state(SubState1)).sub1_foo = "sub1" + (await state.get_state(SubState2)).sub2_foo = "sub2" get_calls.clear() pttl_calls.clear() finally: @@ -789,8 +793,8 @@ async def counting_pttl(key, *args, **kwargs): 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" + assert (await final_state.get_state(SubState1)).sub1_foo == "sub1" + assert (await final_state.get_state(SubState2)).sub2_foo == "sub2" async def test_set_state_resets_touched_flag( diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 4eb0255ac41..8fd1b6f0a9d 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -127,9 +127,15 @@ async def mock_delete(key: KeyT) -> int: # noqa: RUF029 return 1 return 0 - async def mock_getdel(key: KeyT) -> Any: - value = await redis_mock.get(key) - await redis_mock.delete(key) + async def mock_getdel(key: KeyT) -> Any: # noqa: RUF029 + # Self-contained (not delegating to redis_mock.get/delete) so tests + # that instrument those methods don't count GETDEL as a GET. + _expire_keys() + key = _key_bytes(key) + value = keys.pop(key, None) + if value is not None: + expire_times.pop(key, None) + _keyspace_event(key, "del") return value async def mock_pexpire(key: KeyT, px: int, xx: bool = False) -> bool: # noqa: RUF029 From 4c6cc4f217d91bbebc0a9fe35aaefeee58735bfe Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Fri, 17 Jul 2026 18:34:36 -0700 Subject: [PATCH 5/5] Appease pyright in redis manager tests --- tests/units/istate/manager/test_redis.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 5988e2876ea..4ee01142d21 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -769,8 +769,8 @@ 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 + redis.get = counting_get # pyright: ignore[reportAttributeAccessIssue] + redis.pttl = counting_pttl # pyright: ignore[reportAttributeAccessIssue] try: async with state_manager_redis.modify_state( BaseStateToken(ident=token, cls=root_state) @@ -781,8 +781,8 @@ async def counting_pttl(key, *args, **kwargs): get_calls.clear() pttl_calls.clear() finally: - redis.get = orig_get - redis.pttl = orig_pttl + redis.get = orig_get # pyright: ignore[reportAttributeAccessIssue] + redis.pttl = orig_pttl # pyright: ignore[reportAttributeAccessIssue] # 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. @@ -792,6 +792,7 @@ async def counting_pttl(key, *args, **kwargs): final_state = await state_manager_redis.get_state( BaseStateToken(ident=token, cls=root_state) ) + assert isinstance(final_state, root_state) assert final_state.count == 1 assert (await final_state.get_state(SubState1)).sub1_foo == "sub1" assert (await final_state.get_state(SubState2)).sub2_foo == "sub2" @@ -836,6 +837,7 @@ def counting_pipeline(*args, **kwargs): assert pipelines_created == 0 final_state = await state_manager_redis.get_state(state_token) + assert isinstance(final_state, root_state) assert final_state.count == 5 @@ -866,6 +868,7 @@ async def test_set_state_offloads_large_pickles( assert to_thread.called final_state = await state_manager_redis.get_state(state_token) + assert isinstance(final_state, root_state) assert final_state.count == 7