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

[Alek99]

All Submissions:

• Have you followed the guidelines stated in CONTRIBUTING.md file?
• Have you checked to ensure there aren't any other open Pull Requests for the desired changed?

Type of change

• Bug fix (non-breaking change which fixes an issue)

Changes To Core Features:

• Have you added an explanation of what your changes do and why you'd like us to include them?
• Have you written new tests for your core changes, as applicable?

StateManagerRedis.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 redis round trips per event flush for N states. get_state already pipelines its reads; set_state now does the same for writes (ENG-10097):

• Verify the lock once, write once: the tree is walked iteratively, the lock GET + PTTL happen a single time, and all touched states' payloads go out in one pipeline. 3 round trips per flush regardless of tree size.
• Off-loop pickling: states whose previous payload exceeded 64 KiB are serialized via asyncio.to_thread instead of stalling the event loop (safe: the token's lock is held, so nothing mutates the tree mid-pickle — the old code awaited between pickles too).
• Cached _get_required_state_classes: the recursive class-level computation ran on every get_state. It's now an lru_cached module function keyed on (target class, registration context, state hierarchy generation); BaseState.__init_subclass__ and reload_state_module bump the generation, so dynamically created classes and hot reload invalidate cleanly (including same-name class redefinition, where registry sizes don't change).
• _was_touched reset after write (mirrors the disk manager's existing get_and_reset_touched_state): previously the flag was never cleared, so a once-touched state kept alive by the oplock cache was re-pickled and re-SET on every subsequent flush. Reset only happens after a successful pipeline execute, so failed writes retry.

Round trips per event flush (modeled on the exact before/after command flows; "3" = lock GET + PTTL + one pipelined write):

state tree	before	after	worst-case serialized latency @ 0.25 ms RTT
root + 3 substates	12	3	3.00 ms → 0.75 ms (4x)
root + 9 substates	30	3	7.50 ms → 0.75 ms (10x)
3 levels × 3 wide (13 states)	39	3	9.75 ms → 0.75 ms (13x)

New unit tests (run against both mock and real redis in CI): test_set_state_verifies_lock_once_and_pipelines_writes counts actual GET/PTTL commands during a 3-substate modify (asserts exactly one of each); test_set_state_resets_touched_flag asserts a second flush of an unchanged instance creates no write pipeline at all; test_set_state_offloads_large_pickles asserts the to_thread path engages above the threshold; test_required_state_classes_cache_invalidation asserts a subclass defined after cache warm-up appears in the cached result via the generation bump. Existing modify_state/oplock/expiration tests cover end-to-end behavior.

Linear: ENG-10097

🤖 Generated with Claude Code

https://claude.ai/code/session_01Sns7EuBPYvfGCxRSRZ1bUx

---

Generated by Claude Code

[greptile-apps]

Greptile Summary

This PR optimizes Redis-backed state reads and writes. The main changes are:

• Batched touched-state writes into one Redis pipeline.
• Added off-loop serialization for large state payloads.
• Cached required-state class discovery with generation invalidation.
• Reset touched flags after successful Redis writes.
• Added unit coverage for the new Redis manager behavior.

Confidence Score: 4/5

The Redis write path needs a fix for large-state serialization before merging.

• Most batching and cache invalidation behavior matches the old contracts.
• The large-payload path now serializes a live state object in another thread.
• A concurrent mutation can be missed after _was_touched is cleared.

reflex/istate/manager/redis.py

Important Files Changed

Filename	Overview
reflex/istate/manager/redis.py	Batches Redis state writes, adds large-payload offload, resets touched flags after writes, and caches required-state discovery; the off-thread serialization path can race with live state mutation.
reflex/state.py	Adds a state hierarchy generation counter and bumps it on subclass creation and reload to invalidate cached required-state results.
tests/units/istate/manager/test_redis.py	Adds focused tests for pipelined writes, lock checks, touched-flag reset, serialization offload, and cache invalidation.
news/6745.performance.md	Adds a performance news entry for the Redis state manager optimizations.

_{Reviews (1): Last reviewed commit: "chore: add news fragment for #6745" | Re-trigger Greptile}

[greptile-apps]

Live State Thread Race

When a state's previous pickle is larger than the threshold, this sends the live substate object to a worker thread and yields the event loop. If a background task mutates the same instance while _serialize() is running, Redis can receive a stale or partial snapshot and _was_touched is then cleared, so the interleaved change may not be written on the next flush.

[masenf]

i'm wondering if this to_thread change makes sense right now. pickle is a CPU-bound operation, moving it to a thread doesn't stop it from blocking the event loop, unless there's some kind of special pre-emption going on.

also, it might be better to pickle all of the states in a thread and return all the results at once to avoid overhead.

[greptile-apps]

You're right on both counts.

GIL issue: pickle.dumps on typical Python state objects is pure-Python and holds the GIL throughout. asyncio.to_thread schedules work on a thread pool, but the GIL means the event loop thread is still effectively blocked while pickling runs — there's no preemption benefit for CPU-bound work. (The only way to truly unblock the event loop for CPU-bound serialization is ProcessPoolExecutor, but that introduces the race condition raised in my earlier comment even more severely.)

Batching: Even if threading helped, spawning a separate to_thread call per large state is worse than one batched call. The overhead of each thread dispatch adds up.

The simplest and most correct fix is to drop the asyncio.to_thread path entirely and serialize all touched states synchronously before building the pipeline:

Suggested change

	pickle_state = await asyncio.to_thread(substate._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
	# Serialize all touched states synchronously, then pipeline the writes.
	serialized: list[tuple[BaseState, bytes]] = [
	(substate, substate._serialize()) for substate in touched_states
	]
	pipeline = self.redis.pipeline()
	written_states: list[BaseState] = []
	for substate, pickle_state in serialized:

This avoids the event loop yield (and thus the race), removes the threshold branching complexity, and is strictly equivalent for the GIL-bound case. If true off-loop serialization becomes a future goal, it would need a ProcessPoolExecutor with a copy/snapshot of the state — not a live reference.

[codspeed-hq]

Merging this PR will not alter performance

✅ 26 untouched benchmarks
⏩ 8 skipped benchmarks¹

---

_{Comparing claude/reflex-perf-optimizations-21kg8y-redis (c4b7b46) with main (9a5c4d3)}

Footnotes

1. 8 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9cd59d2f22

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[masenf]

the _was_touched flag was already being reset by virtue of BaseState.__setstate__ popping the flag out of the payload that gets saved in the pickle, so this is just added overhead now.