1"""Postgres-backed job queue (`jobs` table, `FOR UPDATE SKIP LOCKED`): priorities, retries with backoff, dead letters, batches.2Local-first: no extra broker. Redis is used only for cross-process locks and caches."""3from __future__ import annotations45import asyncio6import logging7import socket8import traceback9from collections.abc import Awaitable, Callable10from datetime import UTC, datetime, timedelta11from typing import Any1213from sqlalchemy.ext.asyncio import AsyncConnection1415from aiatlas.config import settings16from aiatlas.db import execute, fetch_all, fetch_one, jsonb, transaction17from aiatlas.ids import new_id1819log = logging.getLogger(__name__)2021Handler = Callable[[dict[str, Any], dict[str, Any]], Awaitable[dict[str, Any] | None]]22_HANDLERS: dict[str, Handler] = {}232425def handler(kind: str) -> Callable[[Handler], Handler]:26 def deco(fn: Handler) -> Handler:27 _HANDLERS[kind] = fn28 return fn29 return deco303132async def enqueue(conn: AsyncConnection, kind: str, payload: dict[str, Any], *, priority: int = 5, run_after: datetime | None = None,33 max_attempts: int = 3, batch_id: str | None = None, dedupe_key: str | None = None) -> str | None:34 jid = new_id("queue_job")35 row = await fetch_one(conn, """insert into jobs (id, kind, payload, priority, run_after, max_attempts, batch_id, dedupe_key)36 values (:id, :k, cast(:p as jsonb), :pr, :ra, :ma, :b, :d)37 on conflict (dedupe_key) where status in ('queued','running') do nothing returning id""",38 id=jid, k=kind, p=jsonb(payload), pr=priority, ra=run_after or datetime.now(UTC), ma=max_attempts, b=batch_id, d=dedupe_key)39 return row["id"] if row else None404142async def claim_next(conn: AsyncConnection, worker: str, kinds: list[str] | None = None) -> dict[str, Any] | None:43 kind_filter = "and kind = any(cast(:kinds as text[]))" if kinds else ""44 row = await fetch_one(conn, f"""with next as (45 select id from jobs where status = 'queued' and run_after <= now() {kind_filter}46 order by priority, run_after limit 1 for update skip locked)47 update jobs j set status = 'running', locked_by = :w, locked_at = now(), started_at = now(), attempts = attempts + 148 from next where j.id = next.id returning j.*""", w=worker, kinds=kinds or [])49 return row505152async def complete(conn: AsyncConnection, job_id: str, result: dict[str, Any] | None = None) -> None:53 await execute(conn, "update jobs set status = 'done', finished_at = now(), error = null, payload = payload || cast(:r as jsonb) where id = :id",54 r=jsonb({"_result": result} if result else {}), id=job_id)555657async def fail(conn: AsyncConnection, job: dict[str, Any], error: str) -> None:58 dead = job["attempts"] >= job["max_attempts"]59 backoff = timedelta(seconds=min(3600, 30 * (3 ** max(0, job["attempts"] - 1))))60 await execute(conn, """update jobs set status = :st, finished_at = case when :dead then now() else null end, error = :e, locked_by = null, locked_at = null,61 run_after = case when :dead then run_after else now() + :backoff end where id = :id""",62 st="dead" if dead else "queued", dead=dead, e=error[:4000], backoff=backoff, id=job["id"])636465async def queue_depth(conn: AsyncConnection) -> dict[str, Any]:66 rows = await fetch_all(conn, "select kind, status, count(*) as n from jobs where status in ('queued','running','dead') or finished_at > now() - interval '1 day' group by 1, 2")67 out: dict[str, dict[str, int]] = {}68 for r in rows:69 out.setdefault(r["kind"], {})[r["status"]] = int(r["n"])70 return out717273async def run_worker(*, concurrency: int | None = None, kinds: list[str] | None = None, stop: asyncio.Event | None = None, idle_sleep: float = 3.0) -> None:74 """Long-running worker loop. Handlers are registered with `@handler(kind)` in `aiatlas.services.handlers`."""75 import aiatlas.services.handlers # noqa: F401 (registers handlers)7677 worker = f"{socket.gethostname()}:{new_id('queue_job')[-6:]}"78 stop = stop or asyncio.Event()79 sem = asyncio.Semaphore(concurrency or settings.worker_concurrency)80 log.info("worker started", extra={"worker": worker, "kinds": kinds or "all"})8182 async def one(job: dict[str, Any]) -> None:83 async with sem:84 fn = _HANDLERS.get(job["kind"])85 if fn is None:86 async with transaction() as conn:87 await fail(conn, {**job, "attempts": job["max_attempts"]}, f"no handler for {job['kind']}")88 return89 try:90 result = await fn(job["payload"], job)91 async with transaction() as conn:92 await complete(conn, job["id"], result)93 except Exception as exc: # noqa: BLE00194 log.warning("job failed", extra={"job": job["id"], "kind": job["kind"], "error": str(exc)})95 async with transaction() as conn:96 await fail(conn, job, f"{exc.__class__.__name__}: {exc}\n{traceback.format_exc()[-1500:]}")9798 tasks: set[asyncio.Task[None]] = set()99 while not stop.is_set():100 if sem.locked():101 await asyncio.sleep(0.5)102 continue103 async with transaction() as conn:104 job = await claim_next(conn, worker, kinds)105 if job is None:106 try:107 await asyncio.wait_for(stop.wait(), timeout=idle_sleep)108 except TimeoutError:109 pass110 continue111 t = asyncio.create_task(one(job))112 tasks.add(t)113 t.add_done_callback(tasks.discard)114 if tasks:115 await asyncio.gather(*tasks, return_exceptions=True)116117118async def requeue_stale(conn: AsyncConnection, *, older_than_minutes: int = 120) -> int:119 row = await fetch_one(conn, """with s as (update jobs set status = 'queued', locked_by = null, locked_at = null120 where status = 'running' and locked_at < now() - make_interval(mins => :m) returning 1) select count(*) as n from s""", m=older_than_minutes)121 return int(row["n"]) if row else 0122123124__all__ = ["enqueue", "claim_next", "complete", "fail", "queue_depth", "run_worker", "handler", "requeue_stale"]125