"""Redis-backed advisory lock for serializing operations across workers. Used for cloud-lifecycle paths that need single-flight semantics across admin double-clicks, retry storms, or cron+manual overlap — without holding a MySQL row lock across a multi-second cloud HTTP call. Pattern: async with distributed_lock(f"cloud_sync:{integration_id}", ttl=600): await sync_scopes_for_integration(integration_id) Raises :class:`LockBusy` if the key is already held — caller decides whether to surface 409, retry, or no-op. """ from __future__ import annotations import asyncio import contextlib import secrets import time from collections.abc import AsyncIterator from contextlib import asynccontextmanager from app.core.redis_client import get_redis from app.core.safe_logging import get_safe_logger logger = get_safe_logger(__name__) class LockBusy(Exception): """Raised when the requested distributed lock is already held.""" def __init__(self, key: str) -> None: super().__init__(f"distributed lock busy: {key}") self.key = key # Lua script for safe release: only delete the key if the value matches # the token we set (prevents releasing a lock that has expired and been # re-acquired by another caller). _RELEASE_SCRIPT = """ if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end """ # CAS-extend: only push the TTL out if WE still own the key (token # matches). Returns 1 on success, 0 if the lock was lost (expired + re- # acquired by another holder, or released). Uses PEXPIRE (ms) so callers # can pass any ttl without integer-second rounding surprises. _EXTEND_SCRIPT = """ if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('pexpire', KEYS[1], ARGV[2]) else return 0 end """ async def extend_lock(key: str, token: str, *, ttl: int) -> bool: """CAS-extend a held lock's TTL by ``ttl`` seconds. Returns ``True`` iff ``token`` still owns ``key`` and the expiry was pushed out; ``False`` if the lock was LOST (expired then re-acquired by someone else, or already released). Callers renewing a long-held lock must treat ``False`` as "another holder now owns this — abort". """ if not key: raise ValueError("distributed_lock key must be non-empty") redis = get_redis() full_key = f"helix:lock:{key}" res = await redis.eval( _EXTEND_SCRIPT, 1, full_key, token, str(int(ttl * 1000)), ) return bool(res) async def is_lock_held(key: str) -> bool: """Return ``True`` iff a distributed lock for ``key`` is currently held. A live ``auto_renew`` holder keeps the Redis key alive for the whole critical section (the watchdog re-extends it every ``ttl // 3``s); the key only lapses once the holder dies (hard-kill / OOM) and the TTL runs out. So "key present" is a reliable proxy for "a worker is actively inside the locked section right now" — used by reapers that must distinguish a genuinely-orphaned in-progress row from a legitimately long-running one. """ if not key: return False redis = get_redis() return bool(await redis.exists(f"helix:lock:{key}")) async def _renew_until_cancelled( key: str, token: str, ttl: int, holder: asyncio.Task, lost: dict, ) -> None: """Watchdog body: re-extend the lock every ``ttl // 3`` seconds so a long critical section never lets the lock lapse (which would let a second holder start concurrently). On a LOST lock, flag it and cancel the holder task so the protected work aborts immediately rather than running interleaved with the new holder. """ interval = max(1, ttl // 3) # Track when the real Redis key is guaranteed to have lapsed. The key was # just SET with ``ex=ttl`` right before this watchdog started, so it # expires no later than ``now + ttl``; each successful extend pushes the # deadline out again. If renews keep ERRORING (a sustained Redis outage, # not a one-tick blip), the key silently expires in Redis while we loop — # and a second worker can then acquire it and run the "single-flight" # critical section concurrently. Bounding on this deadline caps that # concurrent window to the actual TTL instead of "until Redis answers # again" (which could be minutes). See core review 2026-07-11. lock_deadline = time.monotonic() + ttl try: while True: await asyncio.sleep(interval) try: still_ours = await extend_lock(key, token, ttl=ttl) except Exception: # noqa: BLE001 # Transient Redis error — do NOT abort the holder on a single # blip; the TTL still has headroom. But if the outage has # lasted past the point where the key is certain to have # expired in Redis, treat it as LOST: another holder may now # own the key, so aborting is the only safe move. if time.monotonic() >= lock_deadline: lost["lost"] = True logger.error( "distributed_lock renew has errored past the TTL " "deadline for key=%s — the Redis key has certainly " "lapsed; aborting holder to prevent concurrent " "execution", key, exc_info=True, ) holder.cancel() return logger.warning( "distributed_lock renew errored for key=%s; will retry", key, exc_info=True, ) continue lock_deadline = time.monotonic() + ttl if not still_ours: lost["lost"] = True logger.error( "distributed_lock LOST mid-hold for key=%s — aborting " "holder to prevent concurrent execution", key, ) holder.cancel() return except asyncio.CancelledError: return # Poll interval for ``acquire_timeout`` waiters. Small enough that a # freed lock is picked up promptly, large enough that a 15-minute wait # is ~180 Redis round-trips rather than tens of thousands. _ACQUIRE_POLL_SEC = 5.0 @asynccontextmanager async def distributed_lock( key: str, *, ttl: int = 300, auto_renew: bool = False, acquire_timeout: float = 0.0, ) -> AsyncIterator[str]: """Acquire a Redis-backed advisory lock for the duration of the ``async with`` block. Auto-expires after ``ttl`` seconds even if the holder dies. ``auto_renew`` (for long critical sections — e.g. a multi-minute identity-sync sweep): spawn a background watchdog that re-extends the lock every ``ttl // 3`` seconds so it never lapses mid-work. Without it, a sweep exceeding ``ttl`` lets the lock expire and the next dispatch tick starts a SECOND concurrent sweep on the same target (interleaved reads/writes → lost updates). If the lock is ever LOST despite renewal (Redis eviction / manual delete), the watchdog cancels the holder and the ``async with`` re-raises :class:`LockBusy` so the caller's busy-handling path (skip / retry) runs instead of continuing to write under a lock someone else now holds. ``acquire_timeout`` (seconds, default 0 = the historical fail-immediately behaviour): poll for up to this long before giving up. Added 2026-08-06 for the cloud-sync triad (``cmdb.scopes_sync`` / ``roles_sync`` / ``grants_sync``), which all serialize on ONE ``cloud_sync:{integration_id}`` key while running as three independent crons scheduled minutes apart. When an earlier member overruns the gap the later one used to lose the lock and no-op — and because a lock-skip was reported as ``status="ok"``, ``cmdb.scopes.aws-primary`` silently did not run for 29 of 31 consecutive days (2026-07-07 .. 2026-08-05): ``cmdb.roles.aws-primary`` starts at 04:09 and takes ~15 min, scopes fires at 04:18 into a held lock. New AWS sub-accounts stayed invisible to Helix for over a month. Waiting turns "silently skipped forever" into "queued behind the sibling that is already doing the same tenant's work". :raises LockBusy: if another holder still has the lock when ``acquire_timeout`` elapses (immediately, when it is 0), or if the lock was lost mid-hold while ``auto_renew`` is set. """ if not key: raise ValueError("distributed_lock key must be non-empty") if acquire_timeout < 0: raise ValueError("acquire_timeout must be >= 0") redis = get_redis() token = secrets.token_hex(16) full_key = f"helix:lock:{key}" won = await redis.set(full_key, token, nx=True, ex=ttl) if not won and acquire_timeout > 0: deadline = time.monotonic() + acquire_timeout while not won: remaining = deadline - time.monotonic() if remaining <= 0: break await asyncio.sleep(min(_ACQUIRE_POLL_SEC, remaining)) won = await redis.set(full_key, token, nx=True, ex=ttl) if won: logger.info( "distributed_lock acquired for key=%s after waiting " "(acquire_timeout=%ss)", key, acquire_timeout, ) if not won: raise LockBusy(key) watchdog: asyncio.Task | None = None lost: dict = {"lost": False} if auto_renew: holder = asyncio.current_task() assert holder is not None # inside a running task by construction watchdog = asyncio.create_task( _renew_until_cancelled(key, token, ttl, holder, lost), ) try: yield token except asyncio.CancelledError: # If the watchdog cancelled us because the lock was lost, convert # to LockBusy so the caller treats it as "someone else owns this" # rather than a shutdown cancel. ``uncancel`` clears the request # so the swallowed CancelledError isn't re-raised at the next # await (Python 3.11+). if lost["lost"]: _t = asyncio.current_task() if _t is not None: with contextlib.suppress(Exception): _t.uncancel() raise LockBusy(key) from None raise finally: if watchdog is not None: watchdog.cancel() with contextlib.suppress(BaseException): await watchdog try: await redis.eval(_RELEASE_SCRIPT, 1, full_key, token) except Exception: # noqa: BLE001 # Lock will auto-expire via TTL; release-failure is not # actionable for the caller. logger.warning( "distributed_lock release failed for key=%s", key, exc_info=True, )