fix(redis): release a lock won after acquire_lock cancellation (LAB-3606)

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

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: ``<scoped_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