diff --git a/news/+redis-lock-handoff.bugfix.md b/news/+redis-lock-handoff.bugfix.md new file mode 100644 index 00000000000..628c3f31a4c --- /dev/null +++ b/news/+redis-lock-handoff.bugfix.md @@ -0,0 +1 @@ +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. diff --git a/reflex/istate/manager/redis.py b/reflex/istate/manager/redis.py index 2d7d25ad244..1ef948bc933 100644 --- a/reflex/istate/manager/redis.py +++ b/reflex/istate/manager/redis.py @@ -44,6 +44,27 @@ "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 +""" + +# 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 @@ -65,6 +86,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. @@ -202,6 +243,18 @@ 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 + ) + # 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 + ) + # 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())) @@ -540,8 +593,16 @@ 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)) + # The lease keeps its lock alive for this use, within the + # lock expiration from now (see `_Lease.keep_lock_alive`). + self._hand_off_cached_state(token.lock_key, time.monotonic()) + try: + yield cached_state + finally: + 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)) return # Opportunistic locking is enabled, so try to hold the lock across multiple calls. @@ -553,8 +614,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) @@ -564,41 +630,43 @@ 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. + single_update = await self._extend_lock( + self._lock_key(token), lock_id, self.lock_expiration ) + 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 @@ -700,7 +768,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: @@ -709,66 +778,7 @@ async def _create_lease_break_task( self._ensure_lock_task() lock_key = token.lock_key - - async def do_flush() -> None: - 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 - 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: - 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=}." - ) - - async def lease_breaker(): - cancelled_error: asyncio.CancelledError | None = None - async with cleanup_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 - # 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 - finally: - # Re-raise any cancellation error after cleaning up. - if cancelled_error is not None: - raise cancelled_error + redis_lock_key = self._lock_key(token) 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. @@ -777,8 +787,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 @@ -796,17 +807,43 @@ async def lease_breaker(): # 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 ): - self._local_leases[lock_key] = task = asyncio.create_task( - lease_breaker(), - name=f"reflex_lease_breaker|{lock_key}|{lock_id.decode()}", + 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, 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 = _Lease( + self, token, lock_id, redis_lock_key, cleanup_ctx.pop_all(), context ) - # Fetch the requested state into the cache. - self._cached_states[lock_key] = await self.get_state(token) + self._local_leases[lock_key] = task = lease.start_task() 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. @@ -819,18 +856,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) + # 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. @@ -1027,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: @@ -1111,15 +1160,22 @@ 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: if self._debug_enabled: 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( @@ -1147,7 +1203,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 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: @@ -1156,24 +1222,47 @@ async def _lock( state_is_locked = False 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: - 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 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. - logger.warning( - f"{lock_key.decode()} was released by {lock_id.decode()}, but it belonged to {deleted_lock_id.decode()}. This is a bug." - ) - # To avoid race when a waiter is registered after the del message is processed. + 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. + + Returns: + The ID holding the lock before the release, None if nobody held it. + """ + # redis-py types EVAL arguments as str; the lock key and ID are bytes here. + release = cast("Callable[..., Awaitable[bytes | None]]", self.redis.eval) + 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. @@ -1194,5 +1283,237 @@ 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) + + +@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 5ed218e1ef4..69eb67ab5ed 100644 --- a/tests/units/istate/manager/test_redis.py +++ b/tests/units/istate/manager/test_redis.py @@ -3,26 +3,29 @@ import asyncio import contextlib import enum +import gc import os import time import uuid from collections import Counter -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 from unittest.mock import AsyncMock, Mock 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 ( + _RELEASE_LOCK_SCRIPT, StateManagerRedis, _default_lock_expiration, _default_oplock_hold_time_ms, ) 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 @@ -160,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. @@ -284,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, @@ -292,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. @@ -301,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 @@ -422,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]], ): @@ -429,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. """ @@ -437,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 @@ -519,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]], ): @@ -526,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. """ @@ -534,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 @@ -630,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, ): @@ -637,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. """ @@ -645,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 @@ -763,7 +791,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. @@ -771,6 +802,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 @@ -778,21 +810,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 @@ -1017,6 +1058,773 @@ async def test_set_state_discards_writes_when_lock_changes_hands( assert saved.count == 0 +@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_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, + 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_keeps_its_lock_for_queued_uses( + state_manager_redis: StateManagerRedis, + root_state: type[RedisTestState], + short_lock_expiration: int, +): + """A lease keeps its lock while the uses queued ahead of its flush run. + + 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_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, + 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. + 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 + await asyncio.sleep(short_lock_expiration * 1.25 / 1000) + + 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 == 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.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, + 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_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, +): + """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. + 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. + """ + 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 other_state_manager() + 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, + 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 1cbbd65ed95..c4f257511f2 100644 --- a/tests/units/mock_redis.py +++ b/tests/units/mock_redis.py @@ -11,7 +11,11 @@ from redis.asyncio import Redis from redis.typing import EncodableT, KeyT -from reflex.istate.manager.redis import _FENCED_SAVE_SCRIPT +from reflex.istate.manager.redis import ( + _EXTEND_LOCK_SCRIPT, + _FENCED_SAVE_SCRIPT, + _RELEASE_LOCK_SCRIPT, +) from reflex.utils import prerequisites WRONGTYPE_MESSAGE = "WRONGTYPE Operation against a key holding the wrong kind of value" @@ -103,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() @@ -191,25 +197,40 @@ def pipeline(transaction: bool = True): return _Pipeline() async def mock_eval(script: str, numkeys: int, *keys_and_args: Any) -> Any: - """Emulate the one script the redis state manager runs. + """Emulate the scripts the redis state manager runs. - It is the fenced save: check the lock held in KEYS[1] against ARGV[1], - write the remaining keys with the expiration in ARGV[2], and return - the lock's PTTL, or None without writing when the lock is not held. + The fenced save checks the lock held in KEYS[1] against ARGV[1], writes + the remaining keys with the expiration in ARGV[2], and returns the + lock's PTTL, or None without writing when the lock is not held. 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 the fenced save script. + 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 lock's PTTL after writing, or None when nothing was written. + The script's reply. Raises: - NotImplementedError: If the script is not the fenced save script. + NotImplementedError: If the script is not one of the state manager's. """ + if script == _RELEASE_LOCK_SCRIPT: + lock_key, lock_id = keys_and_args + owner = await redis_mock.get(lock_key) + if owner == lock_id: + await redis_mock.delete(lock_key) + return owner + if script == _EXTEND_LOCK_SCRIPT: + lock_key, lock_id, px = keys_and_args + if await redis_mock.get(lock_key) != lock_id: + return 0 + return int(await redis_mock.pexpire(lock_key, int(px))) if script != _FENCED_SAVE_SCRIPT: - msg = "mock_redis only emulates the state manager's fenced save script." + msg = "mock_redis only emulates the state manager's scripts." raise NotImplementedError(msg) lock_key, *state_keys = keys_and_args[:numkeys] lock_id, expiration, *payloads = keys_and_args[numkeys:]