From fbbb9ce2ba5cd492c7042c3184c7cf1a05c5a506 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Fri, 9 Oct 2026 01:39:43 +0200 Subject: [PATCH 01/11] fix(state-manager): release the redis lock on cancelled or late handoffs, keep a lease's lock while its state is in use and wake waiters a release missed --- news/+redis-lock-handoff.bugfix.md | 1 + reflex/istate/manager/redis.py | 274 ++++++++++----- tests/units/istate/manager/test_redis.py | 402 ++++++++++++++++++++++- tests/units/mock_redis.py | 12 + 4 files changed, 598 insertions(+), 91 deletions(-) create mode 100644 news/+redis-lock-handoff.bugfix.md diff --git a/news/+redis-lock-handoff.bugfix.md b/news/+redis-lock-handoff.bugfix.md new file mode 100644 index 00000000000..3941737ad3e --- /dev/null +++ b/news/+redis-lock-handoff.bugfix.md @@ -0,0 +1 @@ +Fixed the redis state manager keeping a token's lock until it expired after an event was cancelled while taking it (e.g. by `supersedes`), a late release deleting the next holder's lock, a lease's updates being lost when its holders outlasted the lock expiration, and waiters sitting out the whole lock expiration when a release did not wake them. diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 814435fe433..e1675f421c5 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -45,6 +45,17 @@ "e" # For evicted events (i.e. maxmemory exceeded) ) +# Deletes a lock only while it holds the releasing ID, so a release that comes +# after the lock expired leaves the next holder's lock alone. Replies with the +# ID that held the lock before. +_RELEASE_LOCK_SCRIPT = """ +local owner = redis.call("get", KEYS[1]) +if owner == ARGV[1] then + redis.call("del", KEYS[1]) +end +return owner +""" + async def enable_keyspace_notifications( redis: Redis, events: str = NOTIFY_KEYSPACE_EVENTS @@ -188,6 +199,10 @@ class StateManagerRedis(StateManager): _local_leases: dict[str, asyncio.Task] = dataclasses.field( default_factory=dict, init=False ) + # Flushes of leases cancelled before they started. + _lease_cleanups: set[asyncio.Task] = dataclasses.field( + default_factory=set, init=False + ) # The unique ID for this state manager, the domain for _local_leases. _instance_id: str = dataclasses.field(default_factory=lambda: str(uuid.uuid4())) @@ -510,8 +525,12 @@ async def _try_modify_state( # Opportunistically reuse existing lock. async with self._get_state_cached(token) as cached_state: if cached_state is not None: - yield cached_state - self._notify_next_waiter(self._lock_key(token)) + try: + yield cached_state + finally: + # A waiter queued on the redis lock only learns about the + # lease from here, also when the holder failed. + self._notify_next_waiter(self._lock_key(token)) return # Opportunistic locking is enabled, so try to hold the lock across multiple calls. @@ -523,8 +542,13 @@ async def _try_modify_state( ) except OplockFound: # While waiting for the lock, another process has acquired it, but we can piggy back. - pass - else: + yield None + return + + # Unless a lease takes the lock over (emptying lock_held_ctx) or the + # single update below holds it, it is released here, so a cancelled + # modify cannot leave it behind until it expires. + try: # Do not create a lease break task when multiple instances are waiting. if ( not await self._get_local_lease(lock_key) @@ -534,41 +558,44 @@ async def _try_modify_state( logger.debug( f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} has contention, not leasing" ) - async with lock_held_ctx: - state = await self.get_state(token) - yield state - await self.set_state(token, state, lock_id=lock_id, **context) - return - - # Create the lease break task since we got the lock. - if ( - new_lease_task := await self._create_lease_break_task( + single_update = True + else: + lease_task = await self._create_lease_break_task( token, lock_id, cleanup_ctx=lock_held_ctx, **context ) - ) is ( - current_lease_task := await self._get_local_lease(lock_key) - ) and new_lease_task is not None: - if self._debug_enabled: - logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} obtained lock {lock_id.decode()}." - ) - elif current_lease_task is None: - # Check if we still have the redis lock, then just try to send this one update and release it. - await self._try_extend_lock(self._lock_key(token)) - if await self.redis.get(self._lock_key(token)) == lock_id: + single_update = False + if lease_task is not None: if self._debug_enabled: logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} holding lock {lock_id.decode()}, {new_lease_task=} already exited, doing single update..." + f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} obtained lock {lock_id.decode()}." ) - async with lock_held_ctx: - state = await self.get_state(token) - yield state - await self.set_state(token, state, lock_id=lock_id, **context) - return - elif self._debug_enabled: - logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lock {lock_id.decode()} expired while waiting for lease task to exit..." + elif await self._get_local_lease(lock_key) is None: + # Check if we still have the redis lock, then just try to send this one update and release it. + await self._try_extend_lock(self._lock_key(token)) + single_update = ( + await self.redis.get(self._lock_key(token)) == lock_id ) + if self._debug_enabled: + outcome = ( + "still held, doing single update" + if single_update + else "expired while waiting for lease task to exit" + ) + logger.debug( + f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lock {lock_id.decode()} {outcome}..." + ) + except BaseException: + await lock_held_ctx.aclose() + raise + + if single_update: + async with lock_held_ctx: + state = await self.get_state(token) + yield state + await self.set_state(token, state, lock_id=lock_id, **context) + return + # Nothing left to release once a lease took the lock over. + await lock_held_ctx.aclose() # Have to retry getting the state, but now it's probably cached. yield None @@ -670,7 +697,8 @@ async def _create_lease_break_task( Args: token: The token to create the lease break task for. lock_id: The ID of the lock. - cleanup_ctx: Enter this context while running the lease break task. + cleanup_ctx: The context holding the lock; a returned lease task + takes it over and leaves it empty. context: The state modification context. Returns: @@ -679,6 +707,7 @@ async def _create_lease_break_task( self._ensure_lock_task() lock_key = token.lock_key + lease_started = False async def do_flush() -> None: if (state_lock := self._cached_states_locks.get(lock_key)) is None: @@ -687,28 +716,58 @@ async def do_flush() -> None: f"State lock for {lock_key} missing while finalizing lease." ) return - async with state_lock: - # Write the state to redis while no one else can modify the cached copy. - state = self._cached_states.pop(lock_key, None) - try: - if state: - if self._debug_enabled: + # Whoever holds or waits for the cached state goes before the flush, + # and the write takes its time as well, so keep the lock alive + # meanwhile: it could expire under them, and the flush would then + # drop their changes. + keep_alive = asyncio.create_task(self._keep_lock_alive(token)) + try: + async with state_lock: + # Write the state to redis while no one else can modify the + # cached copy. A lease cancelled before it started can flush + # after the next lease took over, whose copy is not its to write. + state = ( + self._cached_states.pop(lock_key, None) + if self._local_leases.get(lock_key) is task + else None + ) + try: + if state: + if self._debug_enabled: + logger.debug( + f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {lock_id.decode()} flushing state" + ) + await self.set_state( + token, state, lock_id=lock_id, **context + ) + finally: + if (current_lease := self._local_leases.get(lock_key)) is task: + self._local_leases.pop(lock_key, None) + # TODO: clean up the cached states locks periodically + elif self._debug_enabled: logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {lock_id.decode()} flushing state" + f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {lock_id.decode()} cleanup of {task=} found different task in _local_leases {current_lease=}." ) - await self.set_state(token, state, lock_id=lock_id, **context) - finally: - if (current_lease := self._local_leases.get(lock_key)) is task: - self._local_leases.pop(lock_key, None) - # TODO: clean up the cached states locks periodically - elif self._debug_enabled: - logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {lock_id.decode()} cleanup of {task=} found different task in _local_leases {current_lease=}." - ) + finally: + keep_alive.cancel() + + async def flush() -> None: + try: + # Shield the flush from cancellation to ensure it always runs to completion. + await asyncio.shield(do_flush()) + except Exception as e: + # Propagate exception to the main loop, since we have nowhere to catch it. + asyncio.get_running_loop().call_exception_handler({ + "message": "Exception in Redis State Manager lease breaker", + "exception": e, + }) + raise async def lease_breaker(): + nonlocal lease_started + lease_started = True cancelled_error: asyncio.CancelledError | None = None - async with cleanup_ctx: + async with lease_ctx: lease_break_time = self.oplock_hold_time_ms / 1000 if self._debug_enabled: logger.debug( @@ -718,28 +777,29 @@ async def lease_breaker(): await asyncio.sleep(lease_break_time) except asyncio.CancelledError as err: cancelled_error = err - # We got cancelled so if someone is holding the lock, - # extend the timeout so they get the full time to complete. - if ( - state_lock := self._cached_states_locks[lock_key] - ) is not None and state_lock.locked(): - await self._try_extend_lock(self._lock_key(token)) try: - # Shield the flush from cancellation to ensure it always runs to completion. - await asyncio.shield(do_flush()) - except Exception as e: - # Propagate exception to the main loop, since we have nowhere to catch it. - if not isinstance(e, asyncio.CancelledError): - asyncio.get_running_loop().call_exception_handler({ - "message": "Exception in Redis State Manager lease breaker", - "exception": e, - }) - raise + await flush() finally: # Re-raise any cancellation error after cleaning up. if cancelled_error is not None: raise cancelled_error + async def release_unstarted_lease() -> None: + async with lease_ctx: + await flush() + + def on_lease_done(lease: asyncio.Task) -> None: + # A lease cancelled before its first step never ran lease_breaker, + # so nothing else flushes its state and releases its lock. + if lease_started: + return + cleanup = asyncio.create_task( + release_unstarted_lease(), + name=f"reflex_lease_cleanup|{lock_key}|{lock_id.decode()}", + ) + self._lease_cleanups.add(cleanup) + cleanup.add_done_callback(self._lease_cleanups.discard) + if (state_lock := self._cached_states_locks.get(lock_key)) is not None: # We have an existing lock, so lets see if we have an existing lease to cancel. async with state_lock: @@ -747,8 +807,9 @@ async def lease_breaker(): # There's already a lease break task, so cancel it to clear it out. existing_task.cancel() if existing_task is not None: - with contextlib.suppress(asyncio.CancelledError): - await existing_task + # Wait for its flush; its outcome is its own (it reports a + # failure itself), while a cancellation of this task must stop it. + await asyncio.wait({existing_task}) if not self._lock_updates_subscribed.is_set(): # Only a contention notification breaks the lease early, so the @@ -768,12 +829,19 @@ async def lease_breaker(): lock_key not in self._local_leases and await self._n_lock_contenders(self._lock_key(token)) == 0 ): + # Fetch the requested state into the cache. + self._cached_states[lock_key] = await self.get_state(token) + # The lease may start long after the lock was taken (behind the + # previous lease's flush), so it gets a full expiration of its own. + await self._try_extend_lock(self._lock_key(token)) + # The lease owns the lock from here, with no await left that + # could fail before it does. + lease_ctx = cleanup_ctx.pop_all() self._local_leases[lock_key] = task = asyncio.create_task( lease_breaker(), name=f"reflex_lease_breaker|{lock_key}|{lock_id.decode()}", ) - # Fetch the requested state into the cache. - self._cached_states[lock_key] = await self.get_state(token) + task.add_done_callback(on_lease_done) return task return None @@ -802,6 +870,17 @@ async def _try_extend_lock(self, lock_key: bytes) -> bool | None: """ return await self.redis.pexpire(lock_key, self.lock_expiration, xx=True) + async def _keep_lock_alive(self, token: StateToken[Any]) -> None: + """Extend a token's lock every third of its expiration until cancelled. + + Args: + token: The token whose lock to keep. + """ + lock_key = self._lock_key(token) + while True: + await self._try_extend_lock(lock_key) + await asyncio.sleep(self.lock_expiration / 3000) + async def _try_get_lock(self, lock_key: bytes, lock_id: bytes) -> bool | None: """Try to get a redis lock for a token. @@ -1081,9 +1160,12 @@ async def _wait_lock(self, lock_key: bytes, lock_id: bytes) -> None: f"{SMR} [{time.monotonic() - start:.3f}] {lock_key.decode()} waiting for {lock_id.decode()}" ) try: + # Re-check the lock at least every tenth of its expiration: + # a release whose notification got lost would otherwise + # hold every waiter until the lock expiration. await asyncio.wait_for( lock_released_event.wait(), - timeout=max(self.lock_expiration / 1000, 0), + timeout=max(self.lock_expiration / 10_000, 0), ) except (TimeoutError, asyncio.TimeoutError): if self._debug_enabled: @@ -1117,7 +1199,17 @@ async def _lock( f"{event_name}:{uuid.uuid4().hex}" if event_name else uuid.uuid4().hex ).encode() - await self._wait_lock(lock_key, lock_id) + try: + await self._wait_lock(lock_key, lock_id) + except OplockFound: + raise + except BaseException: + # A cancelled acquisition can have set the lock without receiving + # the reply; release it rather than leave it until it expires. + await asyncio.shield(self._release_lock(lock_key, lock_id)) + # It may also have taken the wakeup meant for the next waiter. + self._notify_next_waiter(lock_key) + raise state_is_locked = True try: @@ -1127,22 +1219,36 @@ async def _lock( raise finally: if state_is_locked: - # only delete our lock - deleted_lock_id = cast( - "bytes | None", await self.redis.getdel(lock_key) - ) - if deleted_lock_id == lock_id: + owner = await asyncio.shield(self._release_lock(lock_key, lock_id)) + if owner == lock_id: if self._debug_enabled: logger.debug( f"{SMR} [{time.monotonic() - start:.3f}] {lock_key.decode()} released by {lock_id.decode()}" ) - elif deleted_lock_id is not None: - # This can happen if the caller never tried to `set_state` before the lock expired and is a pretty bad bug. + elif owner is not None: + # The caller held the lock past its expiration without a + # `set_state` noticing, and another holder took it since. logger.warning( - f"{lock_key.decode()} was released by {lock_id.decode()}, but it belonged to {deleted_lock_id.decode()}. This is a bug." + f"{lock_key.decode()} expired while held by {lock_id.decode()} and now belongs to {owner.decode()}, leaving it to them." ) - # To avoid race when a waiter is registered after the del message is processed. - self._notify_next_waiter(lock_key) + # To avoid race when a waiter is registered after the del message is + # processed, and to let the waiters re-check a lock that expired. + self._notify_next_waiter(lock_key) + + async def _release_lock(self, lock_key: bytes, lock_id: bytes) -> bytes | None: + """Delete the redis lock if it is still held by the given ID. + + Args: + lock_key: The redis key for the lock. + lock_id: The ID of the lock to release. + + Returns: + The ID holding the lock before the release, None if nobody held it. + """ + res = self.redis.eval(_RELEASE_LOCK_SCRIPT, 1, lock_key, lock_id) + if inspect.isawaitable(res): + res = await res + return cast("bytes | None", res) async def close(self): """Explicitly close the redis connection and connection_pool. @@ -1164,5 +1270,7 @@ async def close(self): for lease_task in self._local_leases.values(): lease_task.cancel() await asyncio.gather(*self._local_leases.values(), return_exceptions=True) + # Leases cancelled before they started flush in their own tasks. + await asyncio.gather(*self._lease_cleanups, return_exceptions=True) finally: await self.redis.aclose(close_connection_pool=True) diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 2b544fc01b7..9471cd84c2a 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -12,7 +12,7 @@ import pytest import pytest_asyncio -from reflex_base.utils.exceptions import EnvironmentVarValueError +from reflex_base.utils.exceptions import EnvironmentVarValueError, LockExpiredError from reflex.istate.manager.redis import ( StateManagerRedis, @@ -697,7 +697,10 @@ async def test_oplock_immediate_cancel( root_state: type[RedisTestState], event_log: list[dict[str, Any]], ): - """Test that immediate cancellation of modify releases oplock. + """Test that immediate cancellation of the lease releases oplock. + + The lease may be cancelled before its task ever ran, which still has to + flush the state and release the lock. Args: state_manager_redis: The StateManagerRedis to test. @@ -705,6 +708,7 @@ async def test_oplock_immediate_cancel( event_log: The redis event log. """ token = str(uuid.uuid4()) + state_token = BaseStateToken(ident=token, cls=root_state) state_manager_redis._debug_enabled = True state_manager_redis._oplock_enabled = True @@ -712,21 +716,30 @@ async def test_oplock_immediate_cancel( # subscription a lease requires is unavailable. await _subscribed(state_manager_redis) - async def canceller(): + async def canceller() -> asyncio.Task: while (lease_task := state_manager_redis._local_leases.get(token)) is None: # noqa: ASYNC110 await asyncio.sleep(0) lease_task.cancel() + return lease_task task = asyncio.create_task(canceller()) - async with state_manager_redis.modify_state( - BaseStateToken(ident=token, cls=root_state), - ) as new_state: - assert await state_manager_redis._get_local_lease(token) is None + async with state_manager_redis.modify_state(state_token) as new_state: assert isinstance(new_state, root_state) new_state.count += 1 - await task + lease_task = await task + await asyncio.wait({lease_task}) + await asyncio.gather(*state_manager_redis._lease_cleanups) + + assert await state_manager_redis._get_local_lease(token) is None + assert ( + await state_manager_redis.redis.get(state_manager_redis._lock_key(state_token)) + is None + ) + final_state = await state_manager_redis.get_state(state_token) + assert isinstance(final_state, root_state) + assert final_state.count == 1 @pytest.mark.asyncio @@ -884,6 +897,379 @@ async def modify(): assert final_state.count == 2 +@pytest.mark.asyncio +async def test_oplock_cancel_after_lock_acquired_releases_lock( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + monkeypatch: pytest.MonkeyPatch, +): + """Cancelling a modify between taking the lock and leasing it releases the lock. + + A superseded event is cancelled at whatever await it is in, and a lock left + behind keeps every other event of the token waiting until it expires. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + state_manager_redis._oplock_enabled = True + await _subscribed(state_manager_redis) + + lock_acquired = asyncio.Event() + + async def block_after_acquiring(lock_key: bytes) -> int: + lock_acquired.set() + await asyncio.Event().wait() + return 0 + + monkeypatch.setattr( + state_manager_redis, "_n_lock_contenders", block_after_acquiring + ) + + async def modify(): + async with state_manager_redis.modify_state(token): + pass + + task = asyncio.create_task(modify()) + await lock_acquired.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert ( + await state_manager_redis.redis.get(state_manager_redis._lock_key(token)) + is None + ) + + +@pytest.mark.asyncio +async def test_cancel_while_acquiring_lock_releases_it( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + monkeypatch: pytest.MonkeyPatch, +): + """A lock set in redis whose reply never reached the cancelled caller is released. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_set = asyncio.Event() + redis_set = state_manager_redis.redis.set + + async def set_then_hang(*args: Any, **kwargs: Any) -> Any: + result = await redis_set(*args, **kwargs) + if kwargs.get("nx"): + lock_set.set() + await asyncio.Event().wait() + return result + + monkeypatch.setattr(state_manager_redis.redis, "set", set_then_hang) + + async def modify(): + async with state_manager_redis.modify_state(token): + pass + + task = asyncio.create_task(modify()) + await lock_set.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert ( + await state_manager_redis.redis.get(state_manager_redis._lock_key(token)) + is None + ) + + +@pytest.mark.asyncio +async def test_lock_release_after_expiration_keeps_the_next_owner( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """Releasing a lock that already expired must not delete its next owner's lock. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + + async with state_manager_redis._lock(token): + # The lock expires and another holder takes it. + await state_manager_redis.redis.delete(lock_key) + await state_manager_redis.redis.set( + lock_key, b"next-owner", px=state_manager_redis.lock_expiration + ) + + assert await state_manager_redis.redis.get(lock_key) == b"next-owner" + + +@pytest.mark.asyncio +async def test_oplock_lease_break_extends_the_lock_for_a_cached_holder( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, +): + """A lease that breaks while the cached state is in use extends its lock. + + Otherwise the flush that follows the holder finds the lock expired and the + changes made under the lease are lost. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 2 + await _subscribed(state_manager_redis) + + async with state_manager_redis.modify_state(token) as state: + assert isinstance(state, root_state) + state.count += 1 + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + # Still in use past the lock's original expiration, within the + # expiration it gets from the lease break. + await asyncio.sleep(short_lock_expiration * 1.25 / 1000) + + await lease_task + + final_state = await state_manager_redis.get_state(token) + assert isinstance(final_state, root_state) + assert final_state.count == 1 + + +@pytest.mark.asyncio +async def test_oplock_waiter_woken_when_the_cached_holder_raises( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + monkeypatch: pytest.MonkeyPatch, +): + """A local waiter joins the lease when the event ahead of it raises. + + The waiter queued on the redis lock before the lease existed, so only the + holder can tell it about the lease; otherwise it sits out the whole lease. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = 2000 + await _subscribed(state_manager_redis) + + holder_locked = asyncio.Event() + waiter_queued = asyncio.Event() + + async def lease_despite_the_waiter(lock_key: bytes) -> int: + # Another local waiter's SREM can empty the shared waiters set, so the + # holder leases although this one is still queued. + if not holder_locked.is_set(): + holder_locked.set() + await waiter_queued.wait() + return 0 + + monkeypatch.setattr( + state_manager_redis, "_n_lock_contenders", lease_despite_the_waiter + ) + + class HandlerError(Exception): + """Raised by the event holding the cached state.""" + + async def failing_holder(): + async with state_manager_redis.modify_state(token): + raise HandlerError + + async def waiter() -> float: + started = time.monotonic() + async with state_manager_redis.modify_state(token): + return time.monotonic() - started + + holder = asyncio.create_task(failing_holder()) + await holder_locked.wait() + waiting = asyncio.create_task(waiter()) + lock_key = state_manager_redis._lock_key(token) + while not state_manager_redis._n_lock_waiters(lock_key): # noqa: ASYNC110 + await asyncio.sleep(0) + # Its contention notice arrives before the lease exists, so it breaks nothing. + await asyncio.sleep(0.1) + waiter_queued.set() + + with pytest.raises(HandlerError): + await holder + # Woken well before its next re-check of the lock. + assert await waiting < state_manager_redis.lock_expiration / 10_000 / 2 + + +@pytest.mark.asyncio +async def test_cancelled_waiter_passes_its_wakeup_on( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + monkeypatch: pytest.MonkeyPatch, +): + """A waiter cancelled after its wakeup hands the wakeup to the next waiter. + + Otherwise the next waiter, queued after the release was announced, waits + until its wait times out. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + state_manager_redis._oplock_enabled = False + await _subscribed(state_manager_redis) + + first_retry = asyncio.Event() + try_get_lock = state_manager_redis._try_get_lock + + async def first_waiter_stalls_on_retry(key: bytes, lock_id: bytes) -> bool | None: + if lock_id.startswith(b"first:") and state_manager_redis._n_lock_waiters(key): + first_retry.set() + await asyncio.Event().wait() + return await try_get_lock(key, lock_id) + + monkeypatch.setattr( + state_manager_redis, "_try_get_lock", first_waiter_stalls_on_retry + ) + + async def take_lock(event_name: str) -> float: + started = time.monotonic() + async with state_manager_redis._lock(token, event_name=event_name): + return time.monotonic() - started + + release_holder = asyncio.Event() + + async def holder(): + async with state_manager_redis._lock(token, event_name="holder"): + await release_holder.wait() + + holding = asyncio.create_task(holder()) + while await state_manager_redis.redis.get(lock_key) is None: # noqa: ASYNC110 + await asyncio.sleep(0) + first = asyncio.create_task(take_lock("first")) + while not state_manager_redis._n_lock_waiters(lock_key): # noqa: ASYNC110 + await asyncio.sleep(0) + + release_holder.set() + await holding + await first_retry.wait() + # The release is announced before the second waiter queues. + await asyncio.sleep(0.1) + second = asyncio.create_task(take_lock("second")) + while state_manager_redis._n_lock_waiters(lock_key) < 2: # noqa: ASYNC110 + await asyncio.sleep(0) + + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + # Woken well before its next re-check of the lock. + assert await second < state_manager_redis.lock_expiration / 10_000 / 2 + + +@pytest.mark.asyncio +async def test_waiter_rechecks_a_lock_released_without_notice( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + monkeypatch: pytest.MonkeyPatch, +): + """A waiter whose release notification got lost still takes the free lock. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + state_manager_redis.lock_expiration = 2000 + await _subscribed(state_manager_redis) + # Neither the release's own notice nor its keyspace event reach the waiter. + monkeypatch.setattr(state_manager_redis, "_notify_next_waiter", lambda key: None) + + release_holder = asyncio.Event() + + async def holder(): + async with state_manager_redis._lock(token): + await release_holder.wait() + + async def waiter() -> float: + started = time.monotonic() + async with state_manager_redis._lock(token): + return time.monotonic() - started + + holding = asyncio.create_task(holder()) + while await state_manager_redis.redis.get(lock_key) is None: # noqa: ASYNC110 + await asyncio.sleep(0) + waiting = asyncio.create_task(waiter()) + while not state_manager_redis._n_lock_waiters(lock_key): # noqa: ASYNC110 + await asyncio.sleep(0) + release_holder.set() + await holding + + assert await waiting < state_manager_redis.lock_expiration / 1000 / 2 + + +@pytest.mark.asyncio +async def test_waiter_woken_when_the_holder_lock_expired( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], +): + """A holder whose lock expired under it still wakes the next waiter. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + """ + internals = getattr(state_manager_redis.redis, "_internals", {}) + if "keys" not in internals: + pytest.skip("Expiring a lock without its keyspace event needs the mock redis.") + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + await _subscribed(state_manager_redis) + + waiter_queued = asyncio.Event() + + async def holder(): + async with state_manager_redis._lock(token): + await waiter_queued.wait() + # The lock expires, and its expired event was already handled. + internals["keys"].pop(lock_key) + internals["expire_times"].pop(lock_key, None) + msg = "the lock expired" + raise LockExpiredError(msg) + + async def waiter() -> float: + started = time.monotonic() + async with state_manager_redis._lock(token): + return time.monotonic() - started + + holding = asyncio.create_task(holder()) + while await state_manager_redis.redis.get(lock_key) is None: # noqa: ASYNC110 + await asyncio.sleep(0) + waiting = asyncio.create_task(waiter()) + while not state_manager_redis._n_lock_waiters(lock_key): # noqa: ASYNC110 + await asyncio.sleep(0) + waiter_queued.set() + with pytest.raises(LockExpiredError): + await holding + + # Woken well before its next re-check of the lock. + assert await waiting < state_manager_redis.lock_expiration / 10_000 / 2 + + def test_oplock_hold_time_below_one_millisecond( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 572ea4f8653..3351e6a5e59 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -11,6 +11,7 @@ from redis.asyncio import Redis from redis.typing import EncodableT, KeyT +from reflex.istate.manager.redis import _RELEASE_LOCK_SCRIPT from reflex.utils import prerequisites WRONGTYPE_MESSAGE = "WRONGTYPE Operation against a key holding the wrong kind of value" @@ -132,6 +133,16 @@ async def mock_getdel(key: KeyT) -> Any: await redis_mock.delete(key) return value + async def mock_eval(script: str, numkeys: int, *keys_and_args: Any) -> Any: + # Only the lock release script runs on redis. + assert script == _RELEASE_LOCK_SCRIPT + assert numkeys == 1 + key, lock_id = keys_and_args + owner = await redis_mock.get(key) + if owner == lock_id: + await redis_mock.delete(key) + return owner + async def mock_pexpire(key: KeyT, px: int, xx: bool = False) -> bool: # noqa: RUF029 _expire_keys() key = _key_bytes(key) @@ -250,6 +261,7 @@ async def listen() -> AsyncGenerator[dict[str, Any], None]: redis_mock.set = mock_set redis_mock.delete = mock_delete redis_mock.getdel = mock_getdel + redis_mock.eval = mock_eval redis_mock.sadd = mock_sadd redis_mock.srem = mock_srem redis_mock.scard = mock_scard From 9534a444b77bc8a3c77e4a1881b369d92a3f25e2 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Fri, 9 Oct 2026 09:51:34 +0200 Subject: [PATCH 02/11] fix(state-manager): renew a lease's lock only while it owns it and within each use's lock expiration, and report a failed renewal --- news/+redis-lock-handoff.bugfix.md | 2 +- reflex/istate/manager/redis.py | 100 +++++++++++------ tests/units/istate/manager/test_redis.py | 131 +++++++++++++++++++++-- tests/units/mock_redis.py | 40 +++++-- 4 files changed, 224 insertions(+), 49 deletions(-) diff --git a/news/+redis-lock-handoff.bugfix.md b/news/+redis-lock-handoff.bugfix.md index 3941737ad3e..6e6c2f3d6b0 100644 --- a/news/+redis-lock-handoff.bugfix.md +++ b/news/+redis-lock-handoff.bugfix.md @@ -1 +1 @@ -Fixed the redis state manager keeping a token's lock until it expired after an event was cancelled while taking it (e.g. by `supersedes`), a late release deleting the next holder's lock, a lease's updates being lost when its holders outlasted the lock expiration, and waiters sitting out the whole lock expiration when a release did not wake them. +Fixed the redis state manager keeping a token's lock until it expired after an event was cancelled while taking it (e.g. by `supersedes`), a late release deleting the next holder's lock, a lease's updates being lost when the events queued ahead of its flush outlasted the lock expiration together, and waiters sitting out the whole lock expiration when a release did not wake them. diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index e1675f421c5..1a414900a13 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -9,7 +9,7 @@ import sys import time import uuid -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Awaitable, Callable from datetime import timedelta from typing import Any, TypedDict, cast @@ -56,6 +56,16 @@ return owner """ +# Sets a lock's expiration only while it holds the given ID (ARGV[1]) to ARGV[2] +# milliseconds, so a renewal never keeps another holder's lock alive. Replies +# with 1 when it did, 0 otherwise. +_EXTEND_LOCK_SCRIPT = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("pexpire", KEYS[1], ARGV[2]) +end +return 0 +""" + async def enable_keyspace_notifications( redis: Redis, events: str = NOTIFY_KEYSPACE_EVENTS @@ -203,6 +213,10 @@ class StateManagerRedis(StateManager): _lease_cleanups: set[asyncio.Task] = dataclasses.field( default_factory=set, init=False ) + # When the current user of a token's cached state got it (token -> monotonic time). + _cached_state_users_since: dict[str, float] = dataclasses.field( + default_factory=dict, init=False + ) # The unique ID for this state manager, the domain for _local_leases. _instance_id: str = dataclasses.field(default_factory=lambda: str(uuid.uuid4())) @@ -525,9 +539,13 @@ async def _try_modify_state( # Opportunistically reuse existing lock. async with self._get_state_cached(token) as cached_state: if cached_state is not None: + # The lease keeps its lock alive for this use, within the + # lock expiration from now (see `_create_lease_break_task`). + self._cached_state_users_since[token.lock_key] = time.monotonic() try: yield cached_state finally: + self._cached_state_users_since.pop(token.lock_key, None) # A waiter queued on the redis lock only learns about the # lease from here, also when the holder failed. self._notify_next_waiter(self._lock_key(token)) @@ -571,9 +589,8 @@ async def _try_modify_state( ) elif await self._get_local_lease(lock_key) is None: # Check if we still have the redis lock, then just try to send this one update and release it. - await self._try_extend_lock(self._lock_key(token)) - single_update = ( - await self.redis.get(self._lock_key(token)) == lock_id + single_update = await self._extend_lock( + self._lock_key(token), lock_id, self.lock_expiration ) if self._debug_enabled: outcome = ( @@ -707,7 +724,32 @@ async def _create_lease_break_task( self._ensure_lock_task() lock_key = token.lock_key + redis_lock_key = self._lock_key(token) lease_started = False + # When a contention notice broke the lease before its hold time. + broken_at: float | None = None + + async def keep_lock_alive() -> None: + # Whoever uses the cached state ahead of the flush gets the lock + # expiration from when they got it (or from an early lease break, if + # later), as a handler holding the lock itself would; the flush gets + # it from now. A use that runs over loses the lock, and the flush + # then fails instead of writing. + expiration = self.lock_expiration / 1000 + while True: + now = time.monotonic() + if (since := self._cached_state_users_since.get(lock_key)) is None: + deadline = now + expiration + else: + deadline = ( + since if broken_at is None else max(since, broken_at) + ) + expiration + ttl_ms = int((deadline - now) * 1000) + if ttl_ms <= 0 or not await self._extend_lock( + redis_lock_key, lock_id, ttl_ms + ): + return + await asyncio.sleep(expiration / 3) async def do_flush() -> None: if (state_lock := self._cached_states_locks.get(lock_key)) is None: @@ -720,7 +762,7 @@ async def do_flush() -> None: # and the write takes its time as well, so keep the lock alive # meanwhile: it could expire under them, and the flush would then # drop their changes. - keep_alive = asyncio.create_task(self._keep_lock_alive(token)) + keep_alive = asyncio.create_task(keep_lock_alive()) try: async with state_lock: # Write the state to redis while no one else can modify the @@ -750,6 +792,9 @@ async def do_flush() -> None: ) finally: keep_alive.cancel() + # A renewal that failed fails the flush, which reports it. + with contextlib.suppress(asyncio.CancelledError): + await keep_alive async def flush() -> None: try: @@ -764,7 +809,7 @@ async def flush() -> None: raise async def lease_breaker(): - nonlocal lease_started + nonlocal lease_started, broken_at lease_started = True cancelled_error: asyncio.CancelledError | None = None async with lease_ctx: @@ -777,6 +822,7 @@ async def lease_breaker(): await asyncio.sleep(lease_break_time) except asyncio.CancelledError as err: cancelled_error = err + broken_at = time.monotonic() try: await flush() finally: @@ -829,13 +875,17 @@ def on_lease_done(lease: asyncio.Task) -> None: lock_key not in self._local_leases and await self._n_lock_contenders(self._lock_key(token)) == 0 ): - # Fetch the requested state into the cache. - self._cached_states[lock_key] = await self.get_state(token) + state = await self.get_state(token) # The lease may start long after the lock was taken (behind the - # previous lease's flush), so it gets a full expiration of its own. - await self._try_extend_lock(self._lock_key(token)) + # previous lease's flush), so it gets a full expiration of its + # own, unless the lock already expired. + if not await self._extend_lock( + redis_lock_key, lock_id, self.lock_expiration + ): + return None # The lease owns the lock from here, with no await left that # could fail before it does. + self._cached_states[lock_key] = state lease_ctx = cleanup_ctx.pop_all() self._local_leases[lock_key] = task = asyncio.create_task( lease_breaker(), @@ -857,29 +907,20 @@ def _lock_key(token: StateToken[Any]) -> bytes: """ return f"{token.lock_key}_lock".encode() - async def _try_extend_lock(self, lock_key: bytes) -> bool | None: - """Extends the current lock for another lock_expiration period. - - Does not change ownership of the lock! + async def _extend_lock(self, lock_key: bytes, lock_id: bytes, px: int) -> bool: + """Set the lock's expiration, if it is still held by the given ID. Args: lock_key: The redis key for the lock. + lock_id: The ID of the lock. + px: The new expiration in milliseconds. Returns: True if the lock was extended. """ - return await self.redis.pexpire(lock_key, self.lock_expiration, xx=True) - - async def _keep_lock_alive(self, token: StateToken[Any]) -> None: - """Extend a token's lock every third of its expiration until cancelled. - - Args: - token: The token whose lock to keep. - """ - lock_key = self._lock_key(token) - while True: - await self._try_extend_lock(lock_key) - await asyncio.sleep(self.lock_expiration / 3000) + # redis-py types EVAL arguments as str; the lock key and ID are bytes here. + extend = cast("Callable[..., Awaitable[int]]", self.redis.eval) + return bool(await extend(_EXTEND_LOCK_SCRIPT, 1, lock_key, lock_id, px)) async def _try_get_lock(self, lock_key: bytes, lock_id: bytes) -> bool | None: """Try to get a redis lock for a token. @@ -1245,10 +1286,9 @@ async def _release_lock(self, lock_key: bytes, lock_id: bytes) -> bytes | None: Returns: The ID holding the lock before the release, None if nobody held it. """ - res = self.redis.eval(_RELEASE_LOCK_SCRIPT, 1, lock_key, lock_id) - if inspect.isawaitable(res): - res = await res - return cast("bytes | None", res) + # redis-py types EVAL arguments as str; the lock key and ID are bytes here. + release = cast("Callable[..., Awaitable[bytes | None]]", self.redis.eval) + return await release(_RELEASE_LOCK_SCRIPT, 1, lock_key, lock_id) async def close(self): """Explicitly close the redis connection and connection_pool. diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 9471cd84c2a..f6b0d2e8536 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -1011,15 +1011,55 @@ async def test_lock_release_after_expiration_keeps_the_next_owner( @pytest.mark.asyncio -async def test_oplock_lease_break_extends_the_lock_for_a_cached_holder( +async def test_oplock_lease_keeps_its_lock_for_queued_uses( state_manager_redis: StateManagerRedis, root_state: type[RedisTestState], short_lock_expiration: int, ): - """A lease that breaks while the cached state is in use extends its lock. + """A lease keeps its lock while the uses queued ahead of its flush run. - Otherwise the flush that follows the holder finds the lock expired and the - changes made under the lease are lost. + Each use stays within the lock expiration, but together they outlast it, so + without renewal the flush would find the lock expired and drop their changes. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 2 + await _subscribed(state_manager_redis) + + async with state_manager_redis.modify_state(token): + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + + async def use(): + async with state_manager_redis.modify_state(token) as state: + assert isinstance(state, root_state) + state.count += 1 + await asyncio.sleep(short_lock_expiration * 0.6 / 1000) + + # Both queue up on the cached state before the lease breaks. + await asyncio.gather(use(), use()) + await lease_task + + final_state = await state_manager_redis.get_state(token) + assert isinstance(final_state, root_state) + assert final_state.count == 2 + + +@pytest.mark.asyncio +async def test_oplock_lease_lets_a_use_past_the_lock_expiration_expire( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, +): + """A use of the cached state longer than the lock expiration loses the lock. + + Its lease ends in LockExpiredError without writing, as a handler holding + the lock itself would. Args: state_manager_redis: The StateManagerRedis to test. @@ -1036,15 +1076,88 @@ async def test_oplock_lease_break_extends_the_lock_for_a_cached_holder( state.count += 1 lease_task = await state_manager_redis._get_local_lease(token.lock_key) assert lease_task is not None - # Still in use past the lock's original expiration, within the - # expiration it gets from the lease break. await asyncio.sleep(short_lock_expiration * 1.25 / 1000) - await lease_task - + with pytest.raises(LockExpiredError): + await lease_task final_state = await state_manager_redis.get_state(token) assert isinstance(final_state, root_state) - assert final_state.count == 1 + assert final_state.count == 0 + + +@pytest.mark.asyncio +async def test_oplock_lease_renewal_leaves_another_holders_lock_alone( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, +): + """A lease whose lock another holder took never extends that holder's lock. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 2 + await _subscribed(state_manager_redis) + + async with state_manager_redis.modify_state(token): + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + + async with state_manager_redis.modify_state(token): + # The lease's lock expires and another holder takes it, for less time + # than this use keeps the cached state past the lease break. + await state_manager_redis.redis.delete(lock_key) + await state_manager_redis.redis.set( + lock_key, b"other", px=short_lock_expiration * 6 // 10 + ) + await asyncio.sleep(short_lock_expiration * 0.9 / 1000) + assert await state_manager_redis.redis.get(lock_key) is None + + with pytest.raises(LockExpiredError): + await lease_task + + +@pytest.mark.asyncio +async def test_oplock_lease_reports_a_failed_renewal( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, + monkeypatch: pytest.MonkeyPatch, +): + """A lock renewal that fails ends the lease with its error. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 2 + await _subscribed(state_manager_redis) + + async with state_manager_redis.modify_state(token): + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + + class RenewalError(Exception): + """Raised by the failing renewal.""" + + async def failing_renewal(lock_key: bytes, lock_id: bytes, px: int) -> bool: # noqa: RUF029 + raise RenewalError + + monkeypatch.setattr(state_manager_redis, "_extend_lock", failing_renewal) + # In use past the lease break, so the flush renews the lock while it waits. + async with state_manager_redis.modify_state(token): + await asyncio.sleep(short_lock_expiration * 0.75 / 1000) + with pytest.raises(RenewalError): + await lease_task @pytest.mark.asyncio diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 3351e6a5e59..5dba78663a4 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -11,7 +11,7 @@ from redis.asyncio import Redis from redis.typing import EncodableT, KeyT -from reflex.istate.manager.redis import _RELEASE_LOCK_SCRIPT +from reflex.istate.manager.redis import _EXTEND_LOCK_SCRIPT, _RELEASE_LOCK_SCRIPT from reflex.utils import prerequisites WRONGTYPE_MESSAGE = "WRONGTYPE Operation against a key holding the wrong kind of value" @@ -134,14 +134,36 @@ async def mock_getdel(key: KeyT) -> Any: return value async def mock_eval(script: str, numkeys: int, *keys_and_args: Any) -> Any: - # Only the lock release script runs on redis. - assert script == _RELEASE_LOCK_SCRIPT - assert numkeys == 1 - key, lock_id = keys_and_args - owner = await redis_mock.get(key) - if owner == lock_id: - await redis_mock.delete(key) - return owner + """Emulate the lock scripts the redis state manager runs. + + The lock release deletes the lock in KEYS[1] only while it holds + ARGV[1], and returns the ID that held it. The lock extension sets the + lock's expiration to ARGV[2] milliseconds only while it holds ARGV[1], + and returns 1 when it did. + + Args: + script: The Lua source, which must be one of the state manager's. + numkeys: How many leading entries of keys_and_args are keys. + keys_and_args: The keys followed by the arguments. + + Returns: + The script's reply. + + Raises: + NotImplementedError: If the script is not one of the state manager's. + """ + if script not in (_RELEASE_LOCK_SCRIPT, _EXTEND_LOCK_SCRIPT): + msg = "mock_redis only emulates the state manager's scripts." + raise NotImplementedError(msg) + lock_key, lock_id, *args = keys_and_args + owner = await redis_mock.get(lock_key) + if script == _RELEASE_LOCK_SCRIPT: + if owner == lock_id: + await redis_mock.delete(lock_key) + return owner + if owner != lock_id: + return 0 + return int(await redis_mock.pexpire(lock_key, int(args[0]))) async def mock_pexpire(key: KeyT, px: int, xx: bool = False) -> bool: # noqa: RUF029 _expire_keys() From 45a60698c1ddd24274cde4cc65d26af0838e1d81 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Fri, 9 Oct 2026 10:27:42 +0200 Subject: [PATCH 03/11] perf(state-manager): shield a redis lock release only when a cancellation interrupts it --- reflex/istate/manager/redis.py | 44 +++++++++++++--------- tests/units/istate/manager/test_redis.py | 47 +++++++++++++++++++++++- 2 files changed, 72 insertions(+), 19 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 1a414900a13..d4813ea130e 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -1247,7 +1247,7 @@ async def _lock( except BaseException: # A cancelled acquisition can have set the lock without receiving # the reply; release it rather than leave it until it expires. - await asyncio.shield(self._release_lock(lock_key, lock_id)) + await self._release_lock(lock_key, lock_id) # It may also have taken the wakeup meant for the next waiter. self._notify_next_waiter(lock_key) raise @@ -1259,26 +1259,30 @@ async def _lock( state_is_locked = False raise finally: - if state_is_locked: - owner = await asyncio.shield(self._release_lock(lock_key, lock_id)) - if owner == lock_id: - if self._debug_enabled: - logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key.decode()} released by {lock_id.decode()}" + try: + if state_is_locked: + owner = await self._release_lock(lock_key, lock_id) + if owner == lock_id: + if self._debug_enabled: + logger.debug( + f"{SMR} [{time.monotonic() - start:.3f}] {lock_key.decode()} released by {lock_id.decode()}" + ) + elif owner is not None: + # The caller held the lock past its expiration without a + # `set_state` noticing, and another holder took it since. + logger.warning( + f"{lock_key.decode()} expired while held by {lock_id.decode()} and now belongs to {owner.decode()}, leaving it to them." ) - elif owner is not None: - # The caller held the lock past its expiration without a - # `set_state` noticing, and another holder took it since. - logger.warning( - f"{lock_key.decode()} expired while held by {lock_id.decode()} and now belongs to {owner.decode()}, leaving it to them." - ) - # To avoid race when a waiter is registered after the del message is - # processed, and to let the waiters re-check a lock that expired. - self._notify_next_waiter(lock_key) + finally: + # To avoid race when a waiter is registered after the del message is + # processed, and to let the waiters re-check a lock that expired. + self._notify_next_waiter(lock_key) async def _release_lock(self, lock_key: bytes, lock_id: bytes) -> bytes | None: """Delete the redis lock if it is still held by the given ID. + A cancellation waits for the release to finish before it propagates. + Args: lock_key: The redis key for the lock. lock_id: The ID of the lock to release. @@ -1288,7 +1292,13 @@ async def _release_lock(self, lock_key: bytes, lock_id: bytes) -> bytes | None: """ # redis-py types EVAL arguments as str; the lock key and ID are bytes here. release = cast("Callable[..., Awaitable[bytes | None]]", self.redis.eval) - return await release(_RELEASE_LOCK_SCRIPT, 1, lock_key, lock_id) + try: + return await release(_RELEASE_LOCK_SCRIPT, 1, lock_key, lock_id) + except asyncio.CancelledError: + # The release may not have reached redis, so run it again where the + # cancellation cannot stop it; only this path pays for the extra task. + await asyncio.shield(release(_RELEASE_LOCK_SCRIPT, 1, lock_key, lock_id)) + raise async def close(self): """Explicitly close the redis connection and connection_pool. diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index f6b0d2e8536..d6d43402b12 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -6,15 +6,16 @@ import os import time import uuid -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable from types import ModuleType -from typing import Any +from typing import Any, cast import pytest import pytest_asyncio from reflex_base.utils.exceptions import EnvironmentVarValueError, LockExpiredError from reflex.istate.manager.redis import ( + _RELEASE_LOCK_SCRIPT, StateManagerRedis, _default_lock_expiration, _default_oplock_hold_time_ms, @@ -986,6 +987,48 @@ async def modify(): ) +@pytest.mark.asyncio +async def test_cancel_while_releasing_lock_releases_it( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + monkeypatch: pytest.MonkeyPatch, +): + """A release cut short by a cancellation still deletes the lock. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + releasing = asyncio.Event() + # redis-py types EVAL replies as str; the lock release replies with bytes. + redis_eval = cast("Callable[..., Awaitable[Any]]", state_manager_redis.redis.eval) + + async def first_release_hangs(script: str, *args: Any) -> Any: + if script == _RELEASE_LOCK_SCRIPT and not releasing.is_set(): + releasing.set() + await asyncio.Event().wait() + return await redis_eval(script, *args) + + monkeypatch.setattr(state_manager_redis.redis, "eval", first_release_hangs) + + async def hold_lock(): + async with state_manager_redis._lock(token): + pass + + task = asyncio.create_task(hold_lock()) + await releasing.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert ( + await state_manager_redis.redis.get(state_manager_redis._lock_key(token)) + is None + ) + + @pytest.mark.asyncio async def test_lock_release_after_expiration_keeps_the_next_owner( state_manager_redis: StateManagerRedis, From fb6b4bcd37a44392dffc05a8aea7238f3ffed269 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Fri, 9 Oct 2026 11:15:54 +0200 Subject: [PATCH 04/11] fix(state-manager): renew a flushing lease's lock as soon as its cached state changes hands --- reflex/istate/manager/redis.py | 62 ++++++++++++++++++------ tests/units/istate/manager/test_redis.py | 43 ++++++++++++++++ 2 files changed, 89 insertions(+), 16 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index d4813ea130e..8099b4222ec 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -217,6 +217,10 @@ class StateManagerRedis(StateManager): _cached_state_users_since: dict[str, float] = dataclasses.field( default_factory=dict, init=False ) + # Wakes a flushing lease's lock renewal when its cached state changes hands. + _cached_state_handoffs: dict[str, asyncio.Event] = dataclasses.field( + default_factory=dict, init=False + ) # The unique ID for this state manager, the domain for _local_leases. _instance_id: str = dataclasses.field(default_factory=lambda: str(uuid.uuid4())) @@ -541,11 +545,11 @@ async def _try_modify_state( if cached_state is not None: # The lease keeps its lock alive for this use, within the # lock expiration from now (see `_create_lease_break_task`). - self._cached_state_users_since[token.lock_key] = time.monotonic() + self._hand_off_cached_state(token.lock_key, time.monotonic()) try: yield cached_state finally: - self._cached_state_users_since.pop(token.lock_key, None) + self._hand_off_cached_state(token.lock_key, None) # A waiter queued on the redis lock only learns about the # lease from here, also when the holder failed. self._notify_next_waiter(self._lock_key(token)) @@ -736,20 +740,29 @@ async def keep_lock_alive() -> None: # it from now. A use that runs over loses the lock, and the flush # then fails instead of writing. expiration = self.lock_expiration / 1000 - while True: - now = time.monotonic() - if (since := self._cached_state_users_since.get(lock_key)) is None: - deadline = now + expiration - else: - deadline = ( - since if broken_at is None else max(since, broken_at) - ) + expiration - ttl_ms = int((deadline - now) * 1000) - if ttl_ms <= 0 or not await self._extend_lock( - redis_lock_key, lock_id, ttl_ms - ): - return - await asyncio.sleep(expiration / 3) + handoff = self._cached_state_handoffs[lock_key] = asyncio.Event() + try: + while True: + handoff.clear() + now = time.monotonic() + if (since := self._cached_state_users_since.get(lock_key)) is None: + deadline = now + expiration + else: + deadline = ( + since if broken_at is None else max(since, broken_at) + ) + expiration + ttl_ms = int((deadline - now) * 1000) + if ttl_ms <= 0 or not await self._extend_lock( + redis_lock_key, lock_id, ttl_ms + ): + return + # The lock may now expire with the current user's deadline, + # so renew it as soon as the next user takes over. + with contextlib.suppress(TimeoutError, asyncio.TimeoutError): + await asyncio.wait_for(handoff.wait(), expiration / 3) + finally: + if self._cached_state_handoffs.get(lock_key) is handoff: + del self._cached_state_handoffs[lock_key] async def do_flush() -> None: if (state_lock := self._cached_states_locks.get(lock_key)) is None: @@ -895,6 +908,23 @@ def on_lease_done(lease: asyncio.Task) -> None: return task return None + def _hand_off_cached_state(self, lock_key: str, since: float | None) -> None: + """Record when the current user of a cached state got it. + + Wakes the lock renewal of the lease flushing that state, if any, which + then renews the lock for the new user. + + Args: + lock_key: The token's lock key. + since: The monotonic time the user got the state, None once it is done. + """ + if since is None: + self._cached_state_users_since.pop(lock_key, None) + else: + self._cached_state_users_since[lock_key] = since + if (handoff := self._cached_state_handoffs.get(lock_key)) is not None: + handoff.set() + @staticmethod def _lock_key(token: StateToken[Any]) -> bytes: """Get the redis key for a token's lock. diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index d6d43402b12..bb015003b5a 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -1093,6 +1093,49 @@ async def use(): assert final_state.count == 2 +@pytest.mark.asyncio +async def test_oplock_lease_renews_its_lock_when_the_cached_state_changes_hands( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, +): + """A lease renews its lock for the next use as soon as that use starts. + + The renewal for the first use lets the lock expire with that use's + deadline, and the next periodic renewal would only come after it. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + state_manager_redis._oplock_enabled = True + # The flush starts at a tenth of the expiration and renews every third, + # so no periodic renewal falls between the handoff and the first deadline. + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 10 + await _subscribed(state_manager_redis) + + async with state_manager_redis.modify_state(token): + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + + async def use(seconds: float): + async with state_manager_redis.modify_state(token) as state: + assert isinstance(state, root_state) + state.count += 1 + await asyncio.sleep(seconds) + + await asyncio.gather( + use(short_lock_expiration * 0.9 / 1000), use(short_lock_expiration * 0.6 / 1000) + ) + await lease_task + + final_state = await state_manager_redis.get_state(token) + assert isinstance(final_state, root_state) + assert final_state.count == 2 + + @pytest.mark.asyncio async def test_oplock_lease_lets_a_use_past_the_lock_expiration_expire( state_manager_redis: StateManagerRedis, From 50f4b28066be97f2753c4e1b22132e22226ab1fb Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Fri, 9 Oct 2026 12:22:13 +0200 Subject: [PATCH 05/11] fix(state-manager): wake a lease's lock renewal with a timer instead of asyncio.wait_for, which can swallow its cancellation on Python 3.11 --- reflex/istate/manager/redis.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 8099b4222ec..b2d1d5db402 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -757,9 +757,16 @@ async def keep_lock_alive() -> None: ): return # The lock may now expire with the current user's deadline, - # so renew it as soon as the next user takes over. - with contextlib.suppress(TimeoutError, asyncio.TimeoutError): - await asyncio.wait_for(handoff.wait(), expiration / 3) + # so renew it as soon as the next user takes over. Not + # `asyncio.wait_for`, which can swallow the cancellation + # that ends this loop on Python 3.11. + renew_later = asyncio.get_running_loop().call_later( + expiration / 3, handoff.set + ) + try: + await handoff.wait() + finally: + renew_later.cancel() finally: if self._cached_state_handoffs.get(lock_key) is handoff: del self._cached_state_handoffs[lock_key] From f3fcf6ae9f06f77dc28476679daf82984ec01300 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sun, 11 Oct 2026 13:25:58 +0200 Subject: [PATCH 06/11] fix(state-manager): keep a lease cancelled during its flush holding its lock until the write lands --- reflex/istate/manager/redis.py | 17 ++++++- tests/units/istate/manager/test_redis.py | 60 ++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index b2d1d5db402..1d15ce934d4 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -817,9 +817,20 @@ async def do_flush() -> None: await keep_alive async def flush() -> None: + # The flush runs to completion even when this task is cancelled, + # and this task waits for it: returning early would release the + # lock under the write, which the fenced write then fails, and let + # the next lease read the state before the write lands. + flushing = asyncio.create_task(do_flush()) + cancelled_error: asyncio.CancelledError | None = None + while not flushing.done(): + try: + # Unlike awaiting the task, this neither cancels it nor raises its error. + await asyncio.wait({flushing}) + except asyncio.CancelledError as err: # noqa: PERF203 # once per cancellation + cancelled_error = err try: - # Shield the flush from cancellation to ensure it always runs to completion. - await asyncio.shield(do_flush()) + flushing.result() except Exception as e: # Propagate exception to the main loop, since we have nowhere to catch it. asyncio.get_running_loop().call_exception_handler({ @@ -827,6 +838,8 @@ async def flush() -> None: "exception": e, }) raise + if cancelled_error is not None: + raise cancelled_error async def lease_breaker(): nonlocal lease_started, broken_at diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index bb015003b5a..80dfe419750 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -1246,6 +1246,66 @@ async def failing_renewal(lock_key: bytes, lock_id: bytes, px: int) -> bool: # await lease_task +@pytest.mark.asyncio +async def test_oplock_lease_cancelled_while_flushing_keeps_its_lock_until_written( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, + monkeypatch: pytest.MonkeyPatch, +): + """A lease cancelled during its flush keeps its lock until the write lands. + + Ending at the cancellation would release the lock under the write, which + the fenced write then fails, losing the change. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 10 + await _subscribed(state_manager_redis) + + write_started = asyncio.Event() + write_continue = asyncio.Event() + set_state = state_manager_redis.set_state + + async def held_set_state(*args: Any, **kwargs: Any) -> None: + write_started.set() + await write_continue.wait() + await set_state(*args, **kwargs) + + monkeypatch.setattr(state_manager_redis, "set_state", held_set_state) + + async with state_manager_redis.modify_state(token) as state: + assert isinstance(state, root_state) + state.count += 1 + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + lock_id = await state_manager_redis.redis.get(lock_key) + assert lock_id is not None + + await write_started.wait() + lease_task.cancel() + for _ in range(10): + await asyncio.sleep(0) + assert not lease_task.done() + assert await state_manager_redis.redis.get(lock_key) == lock_id + + write_continue.set() + with pytest.raises(asyncio.CancelledError): + await lease_task + assert await state_manager_redis.redis.get(lock_key) is None + monkeypatch.undo() + final_state = await state_manager_redis.get_state(token) + assert isinstance(final_state, root_state) + assert final_state.count == 1 + + @pytest.mark.asyncio async def test_oplock_waiter_woken_when_the_cached_holder_raises( state_manager_redis: StateManagerRedis, From fd329380c95d21515ee94f49de589d621ec843f5 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sun, 11 Oct 2026 14:02:59 +0200 Subject: [PATCH 07/11] refactor(state-manager): wait out a lease's flush through a helper instead of a try in the loop --- reflex/istate/manager/redis.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 1d15ce934d4..a9490debcec 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -87,6 +87,26 @@ async def enable_keyspace_notifications( raise +async def _wait_unless_cancelled( + task: asyncio.Task, +) -> asyncio.CancelledError | None: + """Wait for a task to finish, unless the waiting task is cancelled first. + + Unlike awaiting the task, this neither cancels it nor raises its error. + + Args: + task: The task to wait for. + + Returns: + The cancellation that interrupted the wait, if any. + """ + try: + await asyncio.wait({task}) + except asyncio.CancelledError as err: + return err + return None + + def _default_lock_expiration() -> int: """Get the default lock expiration time. @@ -824,10 +844,7 @@ async def flush() -> None: flushing = asyncio.create_task(do_flush()) cancelled_error: asyncio.CancelledError | None = None while not flushing.done(): - try: - # Unlike awaiting the task, this neither cancels it nor raises its error. - await asyncio.wait({flushing}) - except asyncio.CancelledError as err: # noqa: PERF203 # once per cancellation + if (err := await _wait_unless_cancelled(flushing)) is not None: cancelled_error = err try: flushing.result() From 06029bed0d8b6147eda036b595cbaf1c72f10373 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sun, 11 Oct 2026 14:07:19 +0200 Subject: [PATCH 08/11] docs(state-manager): document the lease break task's helpers --- reflex/istate/manager/redis.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index a9490debcec..7b3660e76ff 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -754,6 +754,11 @@ async def _create_lease_break_task( broken_at: float | None = None async def keep_lock_alive() -> None: + """Renew the lease's lock for each use of the cached state and for the flush. + + Stops once the lease no longer owns the lock or a use runs past its + lock expiration. + """ # Whoever uses the cached state ahead of the flush gets the lock # expiration from when they got it (or from an early lease break, if # later), as a handler holding the lock itself would; the flush gets @@ -792,6 +797,11 @@ async def keep_lock_alive() -> None: del self._cached_state_handoffs[lock_key] async def do_flush() -> None: + """Write the cached state to redis and end the local lease. + + Keeps the lock alive while the uses queued ahead of the write and + the write itself run. + """ if (state_lock := self._cached_states_locks.get(lock_key)) is None: # If we lost the lock, we can't write the state, something went wrong. logger.warning( @@ -837,6 +847,11 @@ async def do_flush() -> None: await keep_alive async def flush() -> None: + """Run the flush to completion and report a failure to the event loop. + + A cancellation of this task waits for the flush and is re-raised + once the write is done. + """ # The flush runs to completion even when this task is cancelled, # and this task waits for it: returning early would release the # lock under the write, which the fenced write then fails, and let @@ -859,6 +874,7 @@ async def flush() -> None: raise cancelled_error async def lease_breaker(): + """Hold the lease until its hold time ends or it is broken, then flush it.""" nonlocal lease_started, broken_at lease_started = True cancelled_error: asyncio.CancelledError | None = None @@ -881,10 +897,16 @@ async def lease_breaker(): raise cancelled_error async def release_unstarted_lease() -> None: + """Flush the state of a lease cancelled before it started and release its lock.""" async with lease_ctx: await flush() def on_lease_done(lease: asyncio.Task) -> None: + """Schedule the cleanup of a lease that was cancelled before it started. + + Args: + lease: The finished lease task. + """ # A lease cancelled before its first step never ran lease_breaker, # so nothing else flushes its state and releases its lock. if lease_started: From 471e5403a8472eb5b437470ad8f0d0388b40bb81 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sun, 25 Oct 2026 15:23:26 +0100 Subject: [PATCH 09/11] Keep remote lock waiters registered, report lease failures once, and extract the lease A remote waiter now re-registers as a lock contender on each re-check, so a lease that outlasts lock_expiration no longer leases again ahead of it; the mock redis no longer reports an sadd for an existing member. A failed lease is reported once: the done callback only reports what flush() did not already report. The lease closures move into a module-level _Lease dataclass, with no change in behavior. The news fragment is split into short sentences. --- news/+redis-lock-handoff.bugfix.md | 2 +- reflex/istate/manager/redis.py | 438 +++++++++++++---------- tests/units/istate/manager/test_redis.py | 135 +++++++ tests/units/mock_redis.py | 6 +- 4 files changed, 396 insertions(+), 185 deletions(-) diff --git a/news/+redis-lock-handoff.bugfix.md b/news/+redis-lock-handoff.bugfix.md index 6e6c2f3d6b0..aaac82662cf 100644 --- a/news/+redis-lock-handoff.bugfix.md +++ b/news/+redis-lock-handoff.bugfix.md @@ -1 +1 @@ -Fixed the redis state manager keeping a token's lock until it expired after an event was cancelled while taking it (e.g. by `supersedes`), a late release deleting the next holder's lock, a lease's updates being lost when the events queued ahead of its flush outlasted the lock expiration together, and waiters sitting out the whole lock expiration when a release did not wake them. +Fixed the redis state manager stalling a token's events until its lock expired after an event was cancelled while taking the lock (for example by `supersedes`). A late release no longer deletes the next holder's lock. A lease now keeps its lock while the events queued ahead of its flush run, so their updates are no longer lost. Other instances waiting for the lock meanwhile stay registered, so its holder does not lease it again ahead of them. Waiters no longer wait out the whole lock expiration when the event ahead of them fails, is cancelled, or loses its lock, or when a release notification is lost. diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index a3e41da2b1e..1ef948bc933 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -594,7 +594,7 @@ async def _try_modify_state( async with self._get_state_cached(token) as cached_state: if cached_state is not None: # The lease keeps its lock alive for this use, within the - # lock expiration from now (see `_create_lease_break_task`). + # lock expiration from now (see `_Lease.keep_lock_alive`). self._hand_off_cached_state(token.lock_key, time.monotonic()) try: yield cached_state @@ -779,174 +779,6 @@ async def _create_lease_break_task( lock_key = token.lock_key redis_lock_key = self._lock_key(token) - lease_started = False - # When a contention notice broke the lease before its hold time. - broken_at: float | None = None - - async def keep_lock_alive() -> None: - """Renew the lease's lock for each use of the cached state and for the flush. - - Stops once the lease no longer owns the lock or a use runs past its - lock expiration. - """ - # Whoever uses the cached state ahead of the flush gets the lock - # expiration from when they got it (or from an early lease break, if - # later), as a handler holding the lock itself would; the flush gets - # it from now. A use that runs over loses the lock, and the flush - # then fails instead of writing. - expiration = self.lock_expiration / 1000 - handoff = self._cached_state_handoffs[lock_key] = asyncio.Event() - try: - while True: - handoff.clear() - now = time.monotonic() - if (since := self._cached_state_users_since.get(lock_key)) is None: - deadline = now + expiration - else: - deadline = ( - since if broken_at is None else max(since, broken_at) - ) + expiration - ttl_ms = int((deadline - now) * 1000) - if ttl_ms <= 0 or not await self._extend_lock( - redis_lock_key, lock_id, ttl_ms - ): - return - # The lock may now expire with the current user's deadline, - # so renew it as soon as the next user takes over. Not - # `asyncio.wait_for`, which can swallow the cancellation - # that ends this loop on Python 3.11. - renew_later = asyncio.get_running_loop().call_later( - expiration / 3, handoff.set - ) - try: - await handoff.wait() - finally: - renew_later.cancel() - finally: - if self._cached_state_handoffs.get(lock_key) is handoff: - del self._cached_state_handoffs[lock_key] - - async def do_flush() -> None: - """Write the cached state to redis and end the local lease. - - Keeps the lock alive while the uses queued ahead of the write and - the write itself run. - """ - if (state_lock := self._cached_states_locks.get(lock_key)) is None: - # If we lost the lock, we can't write the state, something went wrong. - logger.warning( - f"State lock for {lock_key} missing while finalizing lease." - ) - return - # Whoever holds or waits for the cached state goes before the flush, - # and the write takes its time as well, so keep the lock alive - # meanwhile: it could expire under them, and the flush would then - # drop their changes. - keep_alive = asyncio.create_task(keep_lock_alive()) - try: - async with state_lock: - # Write the state to redis while no one else can modify the - # cached copy. A lease cancelled before it started can flush - # after the next lease took over, whose copy is not its to write. - state = ( - self._cached_states.pop(lock_key, None) - if self._local_leases.get(lock_key) is task - else None - ) - try: - if state: - if self._debug_enabled: - logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {lock_id.decode()} flushing state" - ) - await self.set_state( - token, state, lock_id=lock_id, **context - ) - finally: - if (current_lease := self._local_leases.get(lock_key)) is task: - self._local_leases.pop(lock_key, None) - # TODO: clean up the cached states locks periodically - elif self._debug_enabled: - logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {lock_id.decode()} cleanup of {task=} found different task in _local_leases {current_lease=}." - ) - finally: - keep_alive.cancel() - # A renewal that failed fails the flush, which reports it. - with contextlib.suppress(asyncio.CancelledError): - await keep_alive - - async def flush() -> None: - """Run the flush to completion and report a failure to the event loop. - - A cancellation of this task waits for the flush and is re-raised - once the write is done. - """ - # The flush runs to completion even when this task is cancelled, - # and this task waits for it: returning early would release the - # lock under the write, which the fenced write then fails, and let - # the next lease read the state before the write lands. - flushing = asyncio.create_task(do_flush()) - cancelled_error: asyncio.CancelledError | None = None - while not flushing.done(): - if (err := await _wait_unless_cancelled(flushing)) is not None: - cancelled_error = err - try: - flushing.result() - except Exception as e: - # Propagate exception to the main loop, since we have nowhere to catch it. - asyncio.get_running_loop().call_exception_handler({ - "message": "Exception in Redis State Manager lease breaker", - "exception": e, - }) - raise - if cancelled_error is not None: - raise cancelled_error - - async def lease_breaker(): - """Hold the lease until its hold time ends or it is broken, then flush it.""" - nonlocal lease_started, broken_at - lease_started = True - cancelled_error: asyncio.CancelledError | None = None - async with lease_ctx: - lease_break_time = self.oplock_hold_time_ms / 1000 - if self._debug_enabled: - logger.debug( - f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {lock_id.decode()} started, sleeping for {lease_break_time}s" - ) - try: - await asyncio.sleep(lease_break_time) - except asyncio.CancelledError as err: - cancelled_error = err - broken_at = time.monotonic() - try: - await flush() - finally: - # Re-raise any cancellation error after cleaning up. - if cancelled_error is not None: - raise cancelled_error - - async def release_unstarted_lease() -> None: - """Flush the state of a lease cancelled before it started and release its lock.""" - async with lease_ctx: - await flush() - - def on_lease_done(lease: asyncio.Task) -> None: - """Schedule the cleanup of a lease that was cancelled before it started. - - Args: - lease: The finished lease task. - """ - # A lease cancelled before its first step never ran lease_breaker, - # so nothing else flushes its state and releases its lock. - if lease_started: - return - cleanup = asyncio.create_task( - release_unstarted_lease(), - name=f"reflex_lease_cleanup|{lock_key}|{lock_id.decode()}", - ) - self._lease_cleanups.add(cleanup) - cleanup.add_done_callback(self._lease_cleanups.discard) if (state_lock := self._cached_states_locks.get(lock_key)) is not None: # We have an existing lock, so lets see if we have an existing lease to cancel. @@ -975,7 +807,7 @@ def on_lease_done(lease: asyncio.Task) -> None: # Create the task now if one didn't sneak past us. if ( lock_key not in self._local_leases - and await self._n_lock_contenders(self._lock_key(token)) == 0 + and await self._n_lock_contenders(redis_lock_key) == 0 ): state = await self.get_state(token) # The lease may start long after the lock was taken (behind the @@ -988,12 +820,10 @@ def on_lease_done(lease: asyncio.Task) -> None: # The lease owns the lock from here, with no await left that # could fail before it does. self._cached_states[lock_key] = state - lease_ctx = cleanup_ctx.pop_all() - self._local_leases[lock_key] = task = asyncio.create_task( - lease_breaker(), - name=f"reflex_lease_breaker|{lock_key}|{lock_id.decode()}", + lease = _Lease( + self, token, lock_id, redis_lock_key, cleanup_ctx.pop_all(), context ) - task.add_done_callback(on_lease_done) + self._local_leases[lock_key] = task = lease.start_task() return task return None @@ -1236,19 +1066,29 @@ async def _request_lock_release( yield return - lock_waiter_key = lock_key + b"_waiters" - pipeline = self.redis.pipeline() - # Signal intention to request oplock for this process. - pipeline.sadd(lock_waiter_key, self._instance_id) - pipeline.pexpire(lock_waiter_key, self.lock_expiration) - await pipeline.execute() + await self._register_lock_contender(lock_key) try: yield # Waiting for redis/oplock to be acquired. finally: - res = self.redis.srem(lock_waiter_key, self._instance_id) + res = self.redis.srem(lock_key + b"_waiters", self._instance_id) if inspect.isawaitable(res): await res + async def _register_lock_contender(self, lock_key: bytes) -> None: + """Register this instance as a contender for a lock for one lock expiration. + + Adding it to the contenders breaks the lock holder's lease; registering + it again while it is still a contender does not. + + Args: + lock_key: The redis key for the lock. + """ + lock_waiter_key = lock_key + b"_waiters" + pipeline = self.redis.pipeline() + pipeline.sadd(lock_waiter_key, self._instance_id) + pipeline.pexpire(lock_waiter_key, self.lock_expiration) + await pipeline.execute() + async def _get_local_lease( self, token: str, raise_when_found: bool = False ) -> asyncio.Task | None: @@ -1332,6 +1172,10 @@ async def _wait_lock(self, lock_key: bytes, lock_id: bytes) -> None: logger.debug( f"{SMR} [{time.monotonic() - start:.3f}] {lock_key.decode()} wait timeout for {lock_id.decode()}" ) + if self._oplock_enabled: + # A lease can keep the lock past one expiration, which + # would outlast the registration as a contender. + await self._register_lock_contender(lock_key) lock_released_event.set() # to re-check the lock if self._debug_enabled: logger.debug( @@ -1443,3 +1287,233 @@ async def close(self): await asyncio.gather(*self._lease_cleanups, return_exceptions=True) finally: await self.redis.aclose(close_connection_pool=True) + + +@dataclasses.dataclass(eq=False) +class _Lease: + """A local lease on a token's redis lock, which serves the token's state from memory. + + Its task holds the lock until the hold time ends or a contender breaks the + lease, then writes the cached state back and releases the lock. + """ + + manager: StateManagerRedis + token: StateToken[Any] + lock_id: bytes + redis_lock_key: bytes + # Holds the lock, which the lease releases when it ends. + lease_ctx: contextlib.AsyncExitStack + # The state modification context of the event that took the lock. + context: StateModificationContext + # Whether the task ran; one cancelled before its first step did not. + started: bool = False + # When a contention notice broke the lease before its hold time. + broken_at: float | None = None + # The failure the flush reported, which the task raises as well. + reported_error: BaseException | None = None + task: asyncio.Task = dataclasses.field(init=False) + + def start_task(self) -> asyncio.Task: + """Start the lease task, which takes the lock over. + + Returns: + The lease task. + """ + self.task = asyncio.create_task( + self.run(), + name=f"reflex_lease_breaker|{self.token.lock_key}|{self.lock_id.decode()}", + ) + self.task.add_done_callback(self._on_done) + return self.task + + async def run(self) -> None: + """Hold the lease until its hold time ends or it is broken, then flush it.""" + self.started = True + cancelled_error: asyncio.CancelledError | None = None + async with self.lease_ctx: + lease_break_time = self.manager.oplock_hold_time_ms / 1000 + if self.manager._debug_enabled: + logger.debug( + f"{SMR} [{time.monotonic() - start:.3f}] {self.token.lock_key} lease breaker {self.lock_id.decode()} started, sleeping for {lease_break_time}s" + ) + try: + await asyncio.sleep(lease_break_time) + except asyncio.CancelledError as err: + cancelled_error = err + self.broken_at = time.monotonic() + try: + await self.flush() + finally: + # Re-raise any cancellation error after cleaning up. + if cancelled_error is not None: + raise cancelled_error + + async def release_unstarted(self) -> None: + """Flush the state of a lease cancelled before it started and release its lock.""" + async with self.lease_ctx: + await self.flush() + + async def flush(self) -> None: + """Run the write-back to completion and report a failure to the event loop. + + A cancellation of this task waits for the write-back and is re-raised + once it is done. + """ + # The write-back runs to completion even when this task is cancelled, + # and this task waits for it: returning early would release the lock + # under the write, which the fenced write then fails, and let the next + # lease read the state before the write lands. + flushing = asyncio.create_task(self._write_back()) + cancelled_error: asyncio.CancelledError | None = None + while not flushing.done(): + if (err := await _wait_unless_cancelled(flushing)) is not None: + cancelled_error = err + try: + flushing.result() + except Exception as e: + # Report it here: a cancelled lease raises its cancellation instead. + self._report(e) + raise + if cancelled_error is not None: + raise cancelled_error + + async def _write_back(self) -> None: + """Write the cached state to redis and end the local lease. + + Keeps the lock alive while the uses queued ahead of the write and the + write itself run. + """ + manager = self.manager + lock_key = self.token.lock_key + if (state_lock := manager._cached_states_locks.get(lock_key)) is None: + # If we lost the lock, we can't write the state, something went wrong. + logger.warning(f"State lock for {lock_key} missing while finalizing lease.") + return + # Whoever holds or waits for the cached state goes before the flush, and + # the write takes its time as well, so keep the lock alive meanwhile: it + # could expire under them, and the flush would then drop their changes. + keep_alive = asyncio.create_task(self.keep_lock_alive()) + try: + async with state_lock: + # Write the state to redis while no one else can modify the + # cached copy. A lease cancelled before it started can flush + # after the next lease took over, whose copy is not its to write. + state = ( + manager._cached_states.pop(lock_key, None) + if manager._local_leases.get(lock_key) is self.task + else None + ) + try: + if state: + if manager._debug_enabled: + logger.debug( + f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {self.lock_id.decode()} flushing state" + ) + await manager.set_state( + self.token, state, lock_id=self.lock_id, **self.context + ) + finally: + if ( + current_lease := manager._local_leases.get(lock_key) + ) is self.task: + manager._local_leases.pop(lock_key, None) + # TODO: clean up the cached states locks periodically + elif manager._debug_enabled: + logger.debug( + f"{SMR} [{time.monotonic() - start:.3f}] {lock_key} lease breaker {self.lock_id.decode()} cleanup of task={self.task!r} found different task in _local_leases {current_lease=}." + ) + finally: + keep_alive.cancel() + # A renewal that failed fails the flush, which reports it. + with contextlib.suppress(asyncio.CancelledError): + await keep_alive + + async def keep_lock_alive(self) -> None: + """Renew the lease's lock for each use of the cached state and for the flush. + + Stops once the lease no longer owns the lock or a use runs past its + lock expiration. + """ + # Whoever uses the cached state ahead of the flush gets the lock + # expiration from when they got it (or from an early lease break, if + # later), as a handler holding the lock itself would; the flush gets it + # from now. A use that runs over loses the lock, and the flush then + # fails instead of writing. + manager = self.manager + lock_key = self.token.lock_key + expiration = manager.lock_expiration / 1000 + handoff = manager._cached_state_handoffs[lock_key] = asyncio.Event() + try: + while True: + handoff.clear() + now = time.monotonic() + if (since := manager._cached_state_users_since.get(lock_key)) is None: + deadline = now + expiration + else: + deadline = ( + since if self.broken_at is None else max(since, self.broken_at) + ) + expiration + ttl_ms = int((deadline - now) * 1000) + if ttl_ms <= 0 or not await manager._extend_lock( + self.redis_lock_key, self.lock_id, ttl_ms + ): + return + # The lock may now expire with the current user's deadline, so + # renew it as soon as the next user takes over. Not + # `asyncio.wait_for`, which can swallow the cancellation that + # ends this loop on Python 3.11. + renew_later = asyncio.get_running_loop().call_later( + expiration / 3, handoff.set + ) + try: + await handoff.wait() + finally: + renew_later.cancel() + finally: + if manager._cached_state_handoffs.get(lock_key) is handoff: + del manager._cached_state_handoffs[lock_key] + + def _on_done(self, task: asyncio.Task) -> None: + """Report the finished lease task's failure, or release a lease that never started. + + Args: + task: The finished lease task. + """ + if self.started: + self._report_unreported(task) + return + # A lease cancelled before its first step never ran, so nothing else + # flushes its state and releases its lock. + cleanup = asyncio.create_task( + self.release_unstarted(), + name=f"reflex_lease_cleanup|{self.token.lock_key}|{self.lock_id.decode()}", + ) + cleanups = self.manager._lease_cleanups + cleanups.add(cleanup) + cleanup.add_done_callback(cleanups.discard) + cleanup.add_done_callback(self._report_unreported) + + def _report(self, error: BaseException) -> None: + """Report a failure of the lease to the event loop. + + Args: + error: The failure. + """ + self.reported_error = error + asyncio.get_running_loop().call_exception_handler({ + "message": "Exception in Redis State Manager lease breaker", + "exception": error, + }) + + def _report_unreported(self, task: asyncio.Task) -> None: + """Retrieve a finished task's failure and report it, unless the flush did. + + Args: + task: The finished lease or cleanup task. + """ + if ( + not task.cancelled() + and (error := task.exception()) is not None + and error is not self.reported_error + ): + self._report(error) diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 1129662d9a5..2ae1b78a558 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -3,6 +3,7 @@ import asyncio import contextlib import enum +import gc import os import time import uuid @@ -1379,6 +1380,61 @@ async def failing_renewal(lock_key: bytes, lock_id: bytes, px: int) -> bool: # await lease_task +@pytest.mark.parametrize("failing", ["set_state", "_release_lock"]) +@pytest.mark.asyncio +async def test_oplock_lease_reports_its_failure_once( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, + monkeypatch: pytest.MonkeyPatch, + failing: str, +): + """A failed lease is reported once, not again when its unawaited task is collected. + + The flush reports its own failure; a failure after it, releasing the lock, + must still be reported. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + monkeypatch: Pytest monkeypatch fixture. + failing: The state manager method that fails. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 10 + await _subscribed(state_manager_redis) + + class LeaseError(Exception): + """Raised by the failing method.""" + + async def fail(*args: Any, **kwargs: Any) -> None: # noqa: RUF029 + raise LeaseError + + monkeypatch.setattr(state_manager_redis, failing, fail) + loop = asyncio.get_running_loop() + # Only the types: holding a reported error would keep its task alive. + reported: list[type[BaseException] | None] = [] + + def record(_loop: asyncio.AbstractEventLoop, context: dict[str, Any]) -> None: + reported.append(type(error) if (error := context.get("exception")) else None) + + previous_handler = loop.get_exception_handler() + loop.set_exception_handler(record) + try: + async with state_manager_redis.modify_state(token): + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + await asyncio.wait({lease_task}) + del lease_task + gc.collect() + finally: + loop.set_exception_handler(previous_handler) + + assert reported == [LeaseError] + + @pytest.mark.asyncio async def test_oplock_lease_cancelled_while_flushing_keeps_its_lock_until_written( state_manager_redis: StateManagerRedis, @@ -1439,6 +1495,85 @@ async def held_set_state(*args: Any, **kwargs: Any) -> None: assert final_state.count == 1 +@pytest.mark.asyncio +async def test_oplock_remote_waiter_outlasting_a_lock_expiration_blocks_the_next_lease( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, + monkeypatch: pytest.MonkeyPatch, +): + """A waiter on another instance stays a contender while a lease keeps the lock past one expiration. + + Otherwise its registration expires under the lease, the next local event + leases the lock again, and the waiter sits out that lease as well. + + Args: + state_manager_redis: The StateManagerRedis to test. + root_state: The root state class. + short_lock_expiration: The lock expiration time in milliseconds. + monkeypatch: Pytest monkeypatch fixture. + """ + token = BaseStateToken(ident=str(uuid.uuid4()), cls=root_state) + lock_key = state_manager_redis._lock_key(token) + state_manager_redis._oplock_enabled = True + state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 2 + await _subscribed(state_manager_redis) + remote = await _subscribed(StateManagerRedis(redis=state_manager_redis.redis)) + remote._oplock_enabled = True + remote.lock_expiration = short_lock_expiration + + # The remote waiter loses every race for the lock until it may take it. + may_take_lock = asyncio.Event() + try_get_lock = remote._try_get_lock + + async def lose_until_allowed(key: bytes, lock_id: bytes) -> bool | None: + if not may_take_lock.is_set(): + return None + return await try_get_lock(key, lock_id) + + monkeypatch.setattr(remote, "_try_get_lock", lose_until_allowed) + + async with state_manager_redis.modify_state(token): + lease_task = await state_manager_redis._get_local_lease(token.lock_key) + assert lease_task is not None + + use_started = asyncio.Event() + + async def use(): + async with state_manager_redis.modify_state(token) as state: + assert isinstance(state, root_state) + state.count += 1 + use_started.set() + await asyncio.sleep(short_lock_expiration * 0.6 / 1000) + + async def remote_modify(): + async with remote.modify_state(token) as state: + assert isinstance(state, root_state) + state.count += 1 + + # Both uses queue on the cached state before the remote waiter shows up, + # so the lease keeps the lock for both, past one lock expiration. + uses = asyncio.gather(use(), use()) + await use_started.wait() + waiting = asyncio.create_task(remote_modify()) + while not remote._n_lock_waiters(lock_key): # noqa: ASYNC110 + await asyncio.sleep(0) + await uses + await asyncio.wait({lease_task}) + + async with state_manager_redis.modify_state(token) as state: + assert isinstance(state, root_state) + state.count += 1 + assert await state_manager_redis._get_local_lease(token.lock_key) is None + + may_take_lock.set() + await asyncio.wait_for(waiting, timeout=short_lock_expiration / 1000) + await remote.close() + final_state = await state_manager_redis.get_state(token) + assert isinstance(final_state, root_state) + assert final_state.count == 4 + + @pytest.mark.asyncio async def test_oplock_waiter_woken_when_the_cached_holder_raises( state_manager_redis: StateManagerRedis, diff --git a/tests/units/mock_redis.py b/tests/units/mock_redis.py index 8e038c13d06..c4f257511f2 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -107,10 +107,12 @@ async def mock_sadd(key: KeyT, value: EncodableT) -> int: # noqa: RUF029 keyset = keys.setdefault(key, set()) if not isinstance(keyset, set): raise TypeError(WRONGTYPE_MESSAGE) - before = len(keyset) + if value in keyset: + return 0 keyset.add(value) + # Like redis, only announce a member that was added. _keyspace_event(key, "sadd") - return len(keyset) - before + return 1 async def mock_srem(key: KeyT, value: EncodableT) -> int: _expire_keys() From 262a3b33992f64f2ca1dff13688d226a842e0162 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sun, 25 Oct 2026 16:12:24 +0100 Subject: [PATCH 10/11] Give extra state managers in the redis tests their own client Closing a second manager on the shared client disconnected the pool the main manager still used. --- tests/units/istate/manager/test_redis.py | 60 +++++++++++++++++------- 1 file changed, 44 insertions(+), 16 deletions(-) diff --git a/tests/units/istate/manager/test_redis.py b/tests/units/istate/manager/test_redis.py index 2ae1b78a558..69eb67ab5ed 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -25,6 +25,7 @@ ) from reflex.istate.manager.token import BaseStateToken from reflex.state import BaseState +from reflex.utils import prerequisites from tests.units.mock_redis import mock_redis, real_redis @@ -162,6 +163,33 @@ async def state_manager_redis( await state_manager.close() +@pytest_asyncio.fixture(loop_scope="function") +async def other_state_manager( + state_manager_redis: StateManagerRedis, +) -> AsyncGenerator[Callable[[], Awaitable[StateManagerRedis]]]: + """Get a factory for subscribed managers standing in for other app instances. + + On a real Redis each gets its own client, so closing one leaves the shared + client usable; all of them are closed after the test. + + Args: + state_manager_redis: The StateManagerRedis whose Redis they share. + + Yields: + The factory. + """ + async with contextlib.AsyncExitStack() as stack: + + async def create() -> StateManagerRedis: + manager = StateManagerRedis( + redis=prerequisites.get_redis() or state_manager_redis.redis + ) + stack.push_async_callback(manager.close) + return await _subscribed(manager) + + yield create + + @pytest.fixture def event_log(state_manager_redis: StateManagerRedis) -> list[dict[str, Any]]: """Get the redis event log from the state manager. @@ -286,6 +314,7 @@ async def test_get_state_discards_unpicklable_state( async def test_modify_oplock( state_manager_redis: StateManagerRedis, + other_state_manager: Callable[[], Awaitable[StateManagerRedis]], root_state: type[RedisTestState], event_log: list[dict[str, Any]], event_log_on_update: asyncio.Event, @@ -294,6 +323,7 @@ async def test_modify_oplock( Args: state_manager_redis: The StateManagerRedis to test. + other_state_manager: Creates the managers of other app instances. root_state: The root state class. event_log: The redis event log. event_log_on_update: The event for new event records being added to the redis event log. @@ -303,9 +333,7 @@ async def test_modify_oplock( state_manager_redis._debug_enabled = True state_manager_redis._oplock_enabled = True - state_manager_2 = await _subscribed( - StateManagerRedis(redis=state_manager_redis.redis) - ) + state_manager_2 = await other_state_manager() state_manager_2._debug_enabled = True state_manager_2._oplock_enabled = True @@ -424,6 +452,7 @@ async def test_modify_oplock( async def test_oplock_contention_queue( state_manager_redis: StateManagerRedis, + other_state_manager: Callable[[], Awaitable[StateManagerRedis]], root_state: type[RedisTestState], event_log: list[dict[str, Any]], ): @@ -431,6 +460,7 @@ async def test_oplock_contention_queue( Args: state_manager_redis: The StateManagerRedis to test. + other_state_manager: Creates the managers of other app instances. root_state: The root state class. event_log: The redis event log. """ @@ -439,9 +469,7 @@ async def test_oplock_contention_queue( state_manager_redis._debug_enabled = True state_manager_redis._oplock_enabled = True - state_manager_2 = await _subscribed( - StateManagerRedis(redis=state_manager_redis.redis) - ) + state_manager_2 = await other_state_manager() state_manager_2._debug_enabled = True state_manager_2._oplock_enabled = True @@ -521,6 +549,7 @@ async def modify_3(): async def test_oplock_contention_no_lease( state_manager_redis: StateManagerRedis, + other_state_manager: Callable[[], Awaitable[StateManagerRedis]], root_state: type[RedisTestState], event_log: list[dict[str, Any]], ): @@ -528,6 +557,7 @@ async def test_oplock_contention_no_lease( Args: state_manager_redis: The StateManagerRedis to test. + other_state_manager: Creates the managers of other app instances. root_state: The root state class. event_log: The redis event log. """ @@ -536,16 +566,12 @@ async def test_oplock_contention_no_lease( state_manager_redis._debug_enabled = True state_manager_redis._oplock_enabled = True - state_manager_2 = await _subscribed( - StateManagerRedis(redis=state_manager_redis.redis) - ) + state_manager_2 = await other_state_manager() state_manager_2._debug_enabled = True state_manager_2._oplock_enabled = True - state_manager_3 = await _subscribed( - StateManagerRedis(redis=state_manager_redis.redis) - ) + state_manager_3 = await other_state_manager() state_manager_3._debug_enabled = True state_manager_3._oplock_enabled = True @@ -632,6 +658,7 @@ async def modify_3(): @pytest.mark.asyncio async def test_oplock_contention_racers( state_manager_redis: StateManagerRedis, + other_state_manager: Callable[[], Awaitable[StateManagerRedis]], root_state: type[RedisTestState], racer_delay: float | None, ): @@ -639,6 +666,7 @@ async def test_oplock_contention_racers( Args: state_manager_redis: The StateManagerRedis to test. + other_state_manager: Creates the managers of other app instances. root_state: The root state class. racer_delay: The delay before the second racer starts. """ @@ -647,9 +675,7 @@ async def test_oplock_contention_racers( state_manager_redis._debug_enabled = True state_manager_redis._oplock_enabled = True - state_manager_2 = await _subscribed( - StateManagerRedis(redis=state_manager_redis.redis) - ) + state_manager_2 = await other_state_manager() state_manager_2._debug_enabled = True state_manager_2._oplock_enabled = True lease_1 = None @@ -1498,6 +1524,7 @@ async def held_set_state(*args: Any, **kwargs: Any) -> None: @pytest.mark.asyncio async def test_oplock_remote_waiter_outlasting_a_lock_expiration_blocks_the_next_lease( state_manager_redis: StateManagerRedis, + other_state_manager: Callable[[], Awaitable[StateManagerRedis]], root_state: type[RedisTestState], short_lock_expiration: int, monkeypatch: pytest.MonkeyPatch, @@ -1509,6 +1536,7 @@ async def test_oplock_remote_waiter_outlasting_a_lock_expiration_blocks_the_next Args: state_manager_redis: The StateManagerRedis to test. + other_state_manager: Creates the managers of other app instances. root_state: The root state class. short_lock_expiration: The lock expiration time in milliseconds. monkeypatch: Pytest monkeypatch fixture. @@ -1518,7 +1546,7 @@ async def test_oplock_remote_waiter_outlasting_a_lock_expiration_blocks_the_next state_manager_redis._oplock_enabled = True state_manager_redis.oplock_hold_time_ms = short_lock_expiration // 2 await _subscribed(state_manager_redis) - remote = await _subscribed(StateManagerRedis(redis=state_manager_redis.redis)) + remote = await other_state_manager() remote._oplock_enabled = True remote.lock_expiration = short_lock_expiration From b4320ce4b224ed7f838e19f8d6a390b7631e4393 Mon Sep 17 00:00:00 2001 From: Benedikt Bartscher Date: Sun, 25 Oct 2026 16:12:30 +0100 Subject: [PATCH 11/11] Condense the redis lock handoff fragment --- news/+redis-lock-handoff.bugfix.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/news/+redis-lock-handoff.bugfix.md b/news/+redis-lock-handoff.bugfix.md index aaac82662cf..628c3f31a4c 100644 --- a/news/+redis-lock-handoff.bugfix.md +++ b/news/+redis-lock-handoff.bugfix.md @@ -1 +1 @@ -Fixed the redis state manager stalling a token's events until its lock expired after an event was cancelled while taking the lock (for example by `supersedes`). A late release no longer deletes the next holder's lock. A lease now keeps its lock while the events queued ahead of its flush run, so their updates are no longer lost. Other instances waiting for the lock meanwhile stay registered, so its holder does not lease it again ahead of them. Waiters no longer wait out the whole lock expiration when the event ahead of them fails, is cancelled, or loses its lock, or when a release notification is lost. +Fixed the redis state manager stalling a token's events until its lock expired after an event was cancelled (for example by `supersedes`) or failed, or after a lock release notification was lost. With `REFLEX_OPLOCK_ENABLED`, events queued on a lease no longer lose their updates when together they outlast the lock expiration, and events waiting on other instances are no longer passed over by the next lease.