diff --git a/docs/features/distributed-locking.md b/docs/features/distributed-locking.md index bec2efa..1c0b134 100644 --- a/docs/features/distributed-locking.md +++ b/docs/features/distributed-locking.md @@ -162,6 +162,13 @@ Three behavioural edges to design around: # the lock itself self-expires after 30 s (lock_timeout) as the safety net. ``` +On `RedisBackend`, cancelling the task mid-`acquire_lock` does not orphan the +lock: the in-flight `SET NX` and the release both run to completion — however +many cancellations land — before the `CancelledError` propagates. Only Redis +failing the release leaves the key, until the same 30 s TTL as the crash case +above. `CachekitIOBackend` does not yet drain cancellation this way: a cancel +mid-request can leave a server-granted lock held until its server-side timeout. + ### TTL Shorter Than Compute Time ```python @cache(ttl=1) # 1 second TTL diff --git a/src/cachekit/backends/redis/provider.py b/src/cachekit/backends/redis/provider.py index 707b738..f97f316 100644 --- a/src/cachekit/backends/redis/provider.py +++ b/src/cachekit/backends/redis/provider.py @@ -11,14 +11,18 @@ from __future__ import annotations +import asyncio +import functools import logging +import uuid from collections.abc import AsyncIterator from contextlib import asynccontextmanager from contextvars import ContextVar -from typing import Any, Optional +from typing import Any, Optional, TypeVar from urllib.parse import quote as url_encode import redis +from redis.exceptions import LockNotOwnedError from cachekit.backends.base import BaseBackend from cachekit.backends.errors import BackendError @@ -29,6 +33,32 @@ # Module-level ContextVar for async-safe tenant isolation tenant_context: ContextVar[Optional[str]] = ContextVar("tenant_context", default=None) +T = TypeVar("T") + + +async def _await_uninterrupted(fut: asyncio.Future[T]) -> T: + """Await ``fut`` to completion even if the current task is cancelled meanwhile. + + ``asyncio.to_thread`` work is uninterruptible once an executor thread picks it up, and a + still-queued work item is dropped if its future is cancelled first — so a cancelled awaiter + either loses the outcome of a round-trip that still completes, or loses the round-trip + itself. ``asyncio.wait`` never cancels its inputs and never unwraps their result, so keep + waiting on ``fut`` until it is really done, absorbing every cancellation, then re-raise the + last one: callers read ``fut`` for the real outcome before letting it propagate. + + Pass a plain future (``loop.run_in_executor``), never a Task: ``all_tasks()`` sweeps such as + ``asyncio.run`` teardown cancel Tasks out from under the drain, and the outcome is lost again. + """ + cancelled: Optional[asyncio.CancelledError] = None + while not fut.done(): + try: + await asyncio.wait({fut}) + except asyncio.CancelledError as exc: + cancelled = exc + if cancelled is not None: + raise cancelled + return fut.result() + class PerRequestRedisBackend: """Per-request Redis backend wrapper with tenant isolation. @@ -353,24 +383,28 @@ async def acquire_lock( Note: Each acquisition attempt is one non-blocking ``SET NX`` round-trip run via - ``asyncio.to_thread()``; the wait between attempts is an ``asyncio.sleep`` on - the event loop, never a sleep inside an executor thread. A blocking - ``Lock.acquire`` run via ``to_thread`` would pin one executor thread per waiter for - up to ``blocking_timeout``. The default executor has only ``min(32, cpu_count + 4)`` - threads (8 when ``cpu_count`` is 4), so once concurrent misses on one key reach that size - the holder's own ``get``/``set``/``release`` — also ``to_thread`` calls — queue behind - the waiters, every waiter times out, and all of them recompute. + ``loop.run_in_executor()`` and drained through ``_await_uninterrupted`` (so a + cancellation cannot drop a round-trip that still completes); the wait between + attempts is an ``asyncio.sleep`` on the event loop, never a sleep inside an + executor thread. A blocking ``Lock.acquire`` run in the executor would pin one + executor thread per waiter for up to ``blocking_timeout``. The default executor + has only ``min(32, cpu_count + 4)`` threads (8 when ``cpu_count`` is 4), so once + concurrent misses on one key reach that size the holder's own + ``get``/``set``/``release`` — also executor calls — queue behind the waiters, + every waiter times out, and all of them recompute. Sets thread_local=False because attempts and release may run on different executor threads. + Cancellation is drained, not raced: an in-flight attempt or release round-trip + always runs to completion, a lock the attempt wins is released, and only then is + the ``CancelledError`` re-raised. """ - import asyncio - import uuid - # Derive the on-wire Redis lock name from the bare cache key: ``:lock``. # Keeping this suffix on the wire preserves compatibility with existing Redis # deployments — the lock identity didn't change, only the protocol boundary # (the wrapper no longer pollutes the cache_key passed in). scoped_key = f"{self._scoped_key(key)}:lock" + from cachekit.cache_handler import redact_cache_key # local: cache_handler imports the backends package + try: from redis.lock import Lock @@ -384,8 +418,41 @@ async def acquire_lock( loop = asyncio.get_running_loop() deadline = None if blocking_timeout is None else loop.time() + blocking_timeout token = uuid.uuid4().hex # one token for the whole acquisition, however many attempts + + def _release_sync() -> None: + # Catch inside the executor callable, not around _release(): once a cancellation has + # landed, _await_uninterrupted re-raises it and an error left on the future would only + # surface as asyncio's "exception was never retrieved" at GC. + try: + lock.release() + except LockNotOwnedError as e: + logger.debug("Redis lock already expired or taken over before release: %s", e) # nothing to orphan + except redis.RedisError as e: + logger.warning( + "Redis lock release for %s failed; the key lives until its TTL", redact_cache_key(key), exc_info=e + ) + + async def _release() -> None: + # Drained: a cancel landing while this still queues for a thread must not drop the release. + await _await_uninterrupted(loop.run_in_executor(None, _release_sync)) + while True: - acquired = await asyncio.to_thread(lock.acquire, blocking=False, token=token) + # Drained: a cancel cannot stop the thread's SET NX from winning, only hide that it did. + attempt = loop.run_in_executor(None, functools.partial(lock.acquire, blocking=False, token=token)) + try: + acquired = await _await_uninterrupted(attempt) + except asyncio.CancelledError: + # The attempt has finished. One that failed (e.g. a Redis ConnectionError) cannot + # have won; log it rather than let it mask the cancellation. + if (err := attempt.exception()) is not None: + logger.warning( + "Redis lock attempt for %s failed while acquire_lock was being cancelled", + redact_cache_key(key), + exc_info=err, + ) + elif attempt.result(): + await _release() + raise # Same give-up rule as redis-py's Lock.acquire: stop once the next attempt # would land past the deadline. blocking_timeout=None means a single attempt. if acquired or deadline is None or loop.time() + lock.sleep > deadline: @@ -396,11 +463,7 @@ async def acquire_lock( finally: # Release lock if acquired (also run in thread pool) if acquired: - try: - await asyncio.to_thread(lock.release) - except Exception as e: - # Lock may have expired - log but don't fail - logger.debug("Error releasing Redis lock (may have expired): %s", e) + await _release() except Exception as exc: raise classify_redis_error(exc, operation="acquire_lock", key=key) from exc diff --git a/tests/unit/backends/test_redis_backend.py b/tests/unit/backends/test_redis_backend.py index 69ca320..3f999e5 100644 --- a/tests/unit/backends/test_redis_backend.py +++ b/tests/unit/backends/test_redis_backend.py @@ -16,6 +16,7 @@ from __future__ import annotations import asyncio +import logging import threading import time from concurrent.futures import ThreadPoolExecutor @@ -24,6 +25,8 @@ import pytest from redis.commands.core import Script from redis.connection import Encoder +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import LockNotOwnedError from redis.lock import Lock from cachekit.backends.redis import RedisBackend @@ -274,6 +277,15 @@ def __init__(self) -> None: self._store: dict[str, bytes] = {} self._mutex = threading.Lock() self.nx_attempts: list[float] = [] # monotonic time of every SET NX, i.e. every acquire attempt + # Test hooks for the cancellation-mid-attempt race: when set, an NX SET call + # signals nx_entered (so the test knows the executor thread is inside the call), + # blocks on block_nx until the test releases it, then signals nx_done once the + # store write has actually landed — independent of whatever asyncio did with the + # coroutine that was awaiting it. + self.nx_entered: threading.Event | None = None + self.block_nx: threading.Event | None = None + self.nx_done: threading.Event | None = None + self.nx_error: Exception | None = None # raised by the NX SET once unblocked, in place of a result def get_encoder(self) -> Encoder: return Encoder("utf-8", "strict", False) @@ -282,13 +294,23 @@ def register_script(self, script: str) -> Script: return Script(self, script) def set(self, name: str, value: bytes, nx: bool = False, px: int | None = None) -> bool | None: + if nx and self.nx_entered is not None: + self.nx_entered.set() + if nx and self.block_nx is not None: + self.block_nx.wait() + if nx and self.nx_error is not None: + raise self.nx_error with self._mutex: if nx: self.nx_attempts.append(time.monotonic()) if nx and name in self._store: - return None - self._store[name] = value - return True + result = None + else: + self._store[name] = value + result = True + if nx and self.nx_done is not None: + self.nx_done.set() + return result def evalsha(self, _sha: str, _numkeys: int, name: str, token: bytes) -> int: """The only script ``Lock`` runs here is LUA_RELEASE: delete iff the token still matches.""" @@ -299,6 +321,19 @@ def evalsha(self, _sha: str, _numkeys: int, name: str, token: bytes) -> int: return 1 +async def _entered(event: threading.Event, what: str) -> None: + """Poll a thread-side event from the loop; ``to_thread(event.wait)`` would take the very executor thread under test.""" + deadline = time.monotonic() + 2.0 + while not event.is_set(): + assert time.monotonic() < deadline, what + await asyncio.sleep(0.01) + + +async def _acquire_once(backend: PerRequestRedisBackend) -> None: + async with backend.acquire_lock("k", timeout=30.0, blocking_timeout=None): + pass + + @pytest.mark.unit class TestRedisLockWaitersDoNotPinExecutorThreads: """A lock waiter must not hold an executor thread while it waits. @@ -375,3 +410,177 @@ async def test_non_blocking_acquire_makes_exactly_one_attempt(self): assert contended is False assert len(fake.nx_attempts) == 2, "blocking_timeout=None must be a single SET NX per acquire_lock" + + async def test_cancellation_mid_attempt_releases_a_lock_it_goes_on_to_win(self): + """Cancelling the awaiter while the SET NX round-trip is in flight must not orphan the key. + + ``asyncio.to_thread`` can't be interrupted once the executor thread starts the + round-trip, so cancellation only stops the awaiting coroutine from seeing the + result — not the thread from winning the lock. Red on the pre-fix code (the + `try`/`finally` release block is never reached because the cancellation + propagates straight out of the `while True` loop); green once the attempt is awaited uninterrupted. + """ + fake = _FakeRedis() + fake.nx_entered = threading.Event() + fake.block_nx = threading.Event() + fake.nx_done = threading.Event() + backend = PerRequestRedisBackend(fake, tenant_id="t") + + task = asyncio.create_task(_acquire_once(backend)) + await _entered(fake.nx_entered, "executor thread never entered the SET NX call") + + task.cancel() + fake.block_nx.set() # let the executor thread finish the SET NX (it wins the lock) + + with pytest.raises(asyncio.CancelledError): + await task + + # The executor thread runs independently of the cancelled coroutine, so wait for + # its write to actually land before checking the store — otherwise the assertion + # below races the background thread instead of testing the fix. + assert fake.nx_done.wait(2.0), "executor thread never finished the SET NX" + + lock_name = backend._scoped_key("k") + ":lock" + assert lock_name not in fake._store, "lock won after cancellation must still be released" + + async def test_second_cancellation_while_draining_the_attempt_still_releases_the_lock(self): + """A cancel landing while the first one waits out the in-flight SET NX must not orphan the key. + + A plain ``asyncio.shield`` hands the *next* ``task.cancel()`` straight to the attempt + itself: its result is lost and the release skipped. + """ + fake = _FakeRedis() + fake.nx_entered = threading.Event() + fake.block_nx = threading.Event() + fake.nx_done = threading.Event() + backend = PerRequestRedisBackend(fake, tenant_id="t") + + task = asyncio.create_task(_acquire_once(backend)) + await _entered(fake.nx_entered, "executor thread never entered the SET NX call") + + task.cancel() + await asyncio.sleep(0) # first cancellation lands; acquire_lock is now waiting out the attempt + task.cancel() + fake.block_nx.set() # the executor thread finishes the SET NX and wins the lock + + with pytest.raises(asyncio.CancelledError): + await task + assert fake.nx_done.wait(2.0), "executor thread never finished the SET NX" + assert backend._scoped_key("k") + ":lock" not in fake._store, "lock won under repeated cancellation must be released" + + async def test_all_tasks_sweep_mid_attempt_still_releases_the_lock(self): + """``asyncio.run()`` teardown cancels everything in ``all_tasks()``: a round-trip run as a Task dies under the drain. + + A plain executor future is invisible to that sweep, so the win is still read and released. + """ + fake = _FakeRedis() + fake.nx_entered = threading.Event() + fake.block_nx = threading.Event() + fake.nx_done = threading.Event() + backend = PerRequestRedisBackend(fake, tenant_id="t") + + task = asyncio.create_task(_acquire_once(backend)) + await _entered(fake.nx_entered, "executor thread never entered the SET NX call") + + me = asyncio.current_task() + for t in asyncio.all_tasks(): # what asyncio.run()'s _cancel_all_tasks does + if t is not me: + t.cancel() + fake.block_nx.set() # the executor thread finishes the SET NX and wins the lock + + with pytest.raises(asyncio.CancelledError): + await task + assert fake.nx_done.wait(2.0), "executor thread never finished the SET NX" + assert backend._scoped_key("k") + ":lock" not in fake._store, "lock won during a shutdown sweep must be released" + + async def test_second_cancellation_while_the_release_is_queued_still_releases_the_lock(self): + """A cancel landing while ``lock.release`` still waits for an executor thread must not orphan the key. + + With every executor thread busy — the saturation this class exists for — the release + sits in the pool's queue, and a bare ``await to_thread(lock.release)`` lets the next + ``task.cancel()`` cancel that queued work item, so the release never runs. + """ + fake = _FakeRedis() + backend = PerRequestRedisBackend(fake, tenant_id="t") + pool = ThreadPoolExecutor(max_workers=1) + asyncio.get_running_loop().set_default_executor(pool) + holding = asyncio.Event() + + async def hold() -> None: + async with backend.acquire_lock("k", timeout=30.0, blocking_timeout=None) as acquired: + assert acquired + holding.set() + await asyncio.Event().wait() # hold the lock until cancelled + + task = asyncio.create_task(hold()) + await asyncio.wait_for(holding.wait(), 2.0) + + busy = threading.Event() + pool.submit(busy.wait) # the only executor thread is now taken; the release will queue behind it + task.cancel() + await asyncio.sleep(0) # first cancellation lands; the release is queued for the pool + task.cancel() + busy.set() + + with pytest.raises(asyncio.CancelledError): + await task + pool.shutdown(wait=True) # whatever survived in the queue has run by now + assert backend._scoped_key("k") + ":lock" not in fake._store, "lock must be released despite repeated cancellation" + + async def test_attempt_failing_during_cancellation_is_logged_not_raised(self, caplog): + """A Redis error from the in-flight attempt must not replace the ``CancelledError``; it is logged instead.""" + fake = _FakeRedis() + fake.nx_entered = threading.Event() + fake.block_nx = threading.Event() + fake.nx_error = RedisConnectionError("redis went away") + backend = PerRequestRedisBackend(fake, tenant_id="t") + + task = asyncio.create_task(_acquire_once(backend)) + await _entered(fake.nx_entered, "executor thread never entered the SET NX call") + task.cancel() + fake.block_nx.set() # the executor thread now fails the SET NX + + with pytest.raises(asyncio.CancelledError), caplog.at_level(logging.WARNING, logger="cachekit.backends.redis.provider"): + await task + assert any( + r.levelno == logging.WARNING and r.exc_info and isinstance(r.exc_info[1], RedisConnectionError) + for r in caplog.records + ), "a failed attempt swallowed by cancellation must be logged with its traceback" + + async def test_cancellation_mid_attempt_that_loses_leaves_the_holders_lock_alone(self, caplog): + """A cancelled attempt that loses to an existing holder has nothing to release: no release call, nothing logged.""" + fake = _FakeRedis() + fake.nx_entered = threading.Event() + fake.block_nx = threading.Event() + backend = PerRequestRedisBackend(fake, tenant_id="t") + lock_name = backend._scoped_key("k") + ":lock" + fake._store[lock_name] = b"someone-else" + + task = asyncio.create_task(_acquire_once(backend)) + await _entered(fake.nx_entered, "executor thread never entered the SET NX call") + task.cancel() + fake.block_nx.set() # the executor thread finishes the SET NX and loses + + with pytest.raises(asyncio.CancelledError), caplog.at_level(logging.DEBUG, logger="cachekit.backends.redis.provider"): + await task + assert not caplog.records, "a lost attempt must not try to release (a release without a token logs)" + + @pytest.mark.parametrize( + ("error", "level"), + [ + (RedisConnectionError("redis went away"), logging.WARNING), # key orphaned until its TTL: worth a warning + (LockNotOwnedError("expired"), logging.DEBUG), # already gone or taken over: nothing to orphan + ], + ) + async def test_release_failing_in_redis_is_logged_not_raised(self, caplog, monkeypatch, error, level): + """Redis failing the release is the one gap left: the caller sees no error, the key lives until its TTL.""" + fake = _FakeRedis() + monkeypatch.setattr(fake, "evalsha", Mock(side_effect=error)) + backend = PerRequestRedisBackend(fake, tenant_id="t") + + with caplog.at_level(logging.DEBUG, logger="cachekit.backends.redis.provider"): + async with backend.acquire_lock("k", timeout=30.0, blocking_timeout=None) as acquired: + assert acquired + + assert backend._scoped_key("k") + ":lock" in fake._store, "a failed release leaves the key for its TTL" + assert [r.levelno for r in caplog.records if "release" in r.getMessage()] == [level]