"""Per-Celery-task idempotency wrapper backed by ``task_dedup``.""" from __future__ import annotations import functools import inspect import logging import threading import uuid from typing import Any, Callable, Optional from celery.exceptions import Ignore, MaxRetriesExceededError from docsgpt.storage.db.repositories.idempotency import IdempotencyRepository from docsgpt.storage.db.session import db_readonly, db_session logger = logging.getLogger(__name__) # Poison-loop cap; transient-failure headroom without infinite retry. MAX_TASK_ATTEMPTS = 5 # 30s heartbeat / 60s TTL → ~2 missed ticks of slack before reclaim. LEASE_TTL_SECONDS = 60 LEASE_HEARTBEAT_INTERVAL = 30 # 10 × 60s ≈ 5 min of deferral before giving up on a held lease. LEASE_RETRY_MAX = 10 def with_idempotency( task_name: str, *, on_poison: Optional[Callable[[str, dict], None]] = None, ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: """Short-circuit on completed key; gate concurrent runs via a lease. The guard key is the caller's ``idempotency_key``, or one synthesized from ``source_id`` so a keyless dispatch is still poison-guarded. Entry short-circuits: - completed row → return cached result - live lease held → retry(countdown=LEASE_TTL_SECONDS) - attempt_count > MAX_TASK_ATTEMPTS → poison alert; ``on_poison`` fires Success writes ``completed``; exceptions leave ``pending`` for autoretry until the poison-loop guard trips. """ def decorator(fn: Callable[..., Any]) -> Callable[..., Any]: @functools.wraps(fn) def wrapper(self, *args: Any, idempotency_key: Any = None, **kwargs: Any) -> Any: explicit_key = ( idempotency_key if isinstance(idempotency_key, str) and idempotency_key else None ) # A keyless dispatch still gets the guard via a synthesized key; # None means no anchor exists — run unguarded, as before. key = explicit_key or _synthesize_guard_key(task_name, kwargs) if key is None: return fn(self, *args, idempotency_key=idempotency_key, **kwargs) cached = _lookup_completed(key) if cached is not None: logger.info( "idempotency hit for task=%s key=%s — returning cached result", task_name, key, ) return cached owner_id = str(uuid.uuid4()) attempt = _try_claim_lease( key, task_name, _safe_task_id(self), owner_id, ) if attempt is None: # Live lease held by another worker. Re-queue and bail # quickly — by the time the retry fires (LEASE_TTL # seconds), Worker 1 has either finalised (we'll hit # ``_lookup_completed`` and return cached) or its lease # has expired and we can claim. logger.info( "idempotency: live lease held; deferring task=%s key=%s", task_name, key, ) try: raise self.retry( countdown=LEASE_TTL_SECONDS, max_retries=LEASE_RETRY_MAX, ) except MaxRetriesExceededError: # The holder is simply slower than LEASE_RETRY_MAX # deferrals: a task that outruns the broker's visibility # timeout is redelivered while its first run is still # going. Letting the exhaustion propagate would report a # failure for a task that is running normally — but so # would returning a value, only less visibly. A redelivery # reuses the original task id (``Context`` carries # ``task_id`` into the retry), so a return marks the very # id the client polls SUCCESS, and ``/api/task_status`` # hands that to the UI as a finished build. ``Ignore`` # records no state at all, leaving the outcome to the run # that actually holds the lease. logger.info( "idempotency: lease still held after %s deferrals; " "leaving task=%s key=%s to its holder", LEASE_RETRY_MAX, task_name, key, ) raise Ignore() from None if attempt > MAX_TASK_ATTEMPTS: logger.error( "idempotency poison-loop guard: task=%s key=%s attempts=%s", task_name, key, attempt, extra={ "alert": "idempotency_poison_loop", "task_name": task_name, "idempotency_key": key, "attempts": attempt, }, ) poisoned = { "success": False, "error": "idempotency poison-loop guard tripped", "attempts": attempt, } _finalize(key, poisoned, status="failed") _run_poison_hook( on_poison, task_name, fn, self, args, kwargs, idempotency_key, ) return poisoned heartbeat_thread, heartbeat_stop = _start_lease_heartbeat( key, owner_id, ) try: result = fn(self, *args, idempotency_key=idempotency_key, **kwargs) _finalize(key, result, status="completed") return result except Exception: # Drop the lease so the next retry doesn't wait LEASE_TTL. _release_lease(key, owner_id) raise finally: _stop_lease_heartbeat(heartbeat_thread, heartbeat_stop) return wrapper return decorator def _synthesize_guard_key(task_name: str, kwargs: dict) -> Optional[str]: """Derive a deterministic guard key from ``source_id`` for a keyless dispatch. ``source_id`` is stable across broker redeliveries and unique per upload, so the poison-loop counter survives an OOM SIGKILL. Returns ``None`` when absent — the dispatch then runs unguarded as before. """ source_id = kwargs.get("source_id") if source_id: return f"auto:{task_name}:{source_id}" return None def _run_poison_hook( on_poison: Optional[Callable[[str, dict], None]], task_name: str, fn: Callable[..., Any], task_self: Any, args: tuple, kwargs: dict, idempotency_key: Any, ) -> None: """Invoke a task's poison-path hook with named call args; swallow failures. A hook failure must never change the poison-guard outcome. """ if on_poison is None: return try: bound = inspect.signature(fn).bind_partial( task_self, *args, idempotency_key=idempotency_key, **kwargs, ) on_poison(task_name, dict(bound.arguments)) except Exception: logger.exception( "idempotency: poison hook failed for task=%s", task_name, ) def _lookup_completed(key: str) -> Any: """Return cached ``result_json`` if a completed row exists for ``key``, else None.""" with db_readonly() as conn: row = IdempotencyRepository(conn).get_task(key) if row is None: return None if row.get("status") != "completed": return None return row.get("result_json") def _try_claim_lease( key: str, task_name: str, task_id: str, owner_id: str, ) -> Optional[int]: """Atomic CAS; returns ``attempt_count`` or ``None`` when held. DB outage → treated as ``attempt=1`` so transient failures don't block all task execution; reconciler repairs the lease columns. """ try: with db_session() as conn: return IdempotencyRepository(conn).try_claim_lease( key=key, task_name=task_name, task_id=task_id, owner_id=owner_id, ttl_seconds=LEASE_TTL_SECONDS, ) except Exception: logger.exception( "idempotency lease-claim failed for key=%s task=%s", key, task_name, ) return 1 def _finalize(key: str, result_json: Any, *, status: str) -> None: """Best-effort terminal write. Never let DB outage fail the task.""" try: with db_session() as conn: IdempotencyRepository(conn).finalize_task( key=key, result_json=result_json, status=status, ) except Exception: logger.exception( "idempotency finalize failed for key=%s status=%s", key, status, ) def _release_lease(key: str, owner_id: str) -> None: """Best-effort lease release on the wrapper's exception path.""" try: with db_session() as conn: IdempotencyRepository(conn).release_lease(key, owner_id) except Exception: logger.exception("idempotency release-lease failed for key=%s", key) def _start_lease_heartbeat( key: str, owner_id: str, ) -> tuple[threading.Thread, threading.Event]: """Spawn a daemon thread that bumps ``lease_expires_at`` every :data:`LEASE_HEARTBEAT_INTERVAL` seconds until ``stop_event`` fires. Mirrors ``docsgpt.worker._start_ingest_heartbeat`` so the two durability heartbeats share shape and cadence. """ stop_event = threading.Event() thread = threading.Thread( target=_lease_heartbeat_loop, args=(key, owner_id, stop_event, LEASE_HEARTBEAT_INTERVAL), daemon=True, name=f"idempotency-lease-heartbeat:{key[:32]}", ) thread.start() return thread, stop_event def _stop_lease_heartbeat( thread: threading.Thread, stop_event: threading.Event, ) -> None: """Signal the heartbeat thread to exit and join with a short timeout.""" stop_event.set() thread.join(timeout=10) def _lease_heartbeat_loop( key: str, owner_id: str, stop_event: threading.Event, interval: int, ) -> None: """Refresh the lease until ``stop_event`` is set or ownership is lost. A failed refresh (rowcount 0) means another worker stole the lease after expiry — at that point the damage is already possible, so we log and keep ticking. Don't escalate to thread death; the main task body needs to keep running so its outcome is at least *recorded*. """ while not stop_event.wait(interval): try: with db_session() as conn: still_owned = IdempotencyRepository(conn).refresh_lease( key=key, owner_id=owner_id, ttl_seconds=LEASE_TTL_SECONDS, ) if not still_owned: logger.warning( "idempotency lease lost mid-task for key=%s " "(another worker may have taken over)", key, ) except Exception: logger.exception( "idempotency lease-heartbeat tick failed for key=%s", key, ) def _safe_task_id(task_self: Any) -> str: """Best-effort extraction of ``self.request.id`` from a Celery task.""" try: request = getattr(task_self, "request", None) task_id: Optional[str] = ( getattr(request, "id", None) if request is not None else None ) except Exception: task_id = None return task_id or "unknown"