"""작업 큐 CRUD — PostgreSQL 을 '제대로' 큐로 쓴다. (LPS `crud/job_crud.py` 이식) - 할당은 **단일 문장 원자 claim**: FOR UPDATE SKIP LOCKED 서브쿼리 + 같은 UPDATE + RETURNING. → 워커 컨테이너가 몇 개든 같은 잡 이중 할당이 원천 불가. fetch 와 claim 을 분리하지 않는다. - 모든 전이는 **조건부 CAS**(WHERE 에 status/worker_id 가드) + RETURNING. - 복구는 timeout 추측이 아니라 **lease 만료 소유권**(reaper 가 회수). - 재시도/백오프/dead-letter 를 큐에 내장. 전이가 조회/변경으로 나뉘지 않으므로(RETURNING) execute_lambda_write 로 실행한다. 큐 전이만 raw SQL 이다 — 다른 crud 는 전부 SQLAlchemy 표현식을 쓴다. """ import json from sqlalchemy import text from common.database.db_session_manager import DB_SESSION_MNG from common.enums import DBType, DBWRType, JobStatus # 잡 적재 시 워커를 즉시 깨우는 LISTEN/NOTIFY 채널(폴링 제거). JOB_NOTIFY_CHANNEL = "web4ai_job" def compute_backoff(attempts: int, base: float = 5.0, cap: float = 600.0) -> float: """지수 백오프(초). attempts 회 시도 후 다음 재시도까지 대기 = base * 2^(attempts-1), cap 상한.""" return min(cap, base * (2 ** max(0, attempts - 1))) class JobQueue: DB = DBType.MAIN.value async def _tx(self, fn): """쓰기 트랜잭션 — 값 반환이 필요한 큐 전이 전용 진입점.""" return await DB_SESSION_MNG.execute_lambda_write(self.DB, fn) # ---- 적재 ---- async def enqueue( self, job_type: int, payload: dict, priority: int = 100, dedupe_key: str | None = None, max_attempts: int = 3, ) -> str | None: """잡 적재. dedupe_key 가 활성(PENDING/RUNNING) 중복이면 삽입 없이 None 반환.""" sql = text(""" INSERT INTO jobs (job_type, priority, payload, dedupe_key, max_attempts) VALUES (:t, :p, CAST(:payload AS jsonb), :dk, :ma) ON CONFLICT (dedupe_key) WHERE status IN (1, 2) AND dedupe_key IS NOT NULL DO NOTHING RETURNING job_id """) async def run(s): row = (await s.execute(sql, { "t": job_type, "p": priority, "payload": json.dumps(payload, ensure_ascii=False), "dk": dedupe_key, "ma": max_attempts, })).first() if row: # 커밋 시 전달됨 → LISTEN 중인 유휴 워커를 즉시 깨운다(중복 스킵 시엔 알림 안 함). await s.execute(text("SELECT pg_notify(:ch, '')"), {"ch": JOB_NOTIFY_CHANNEL}) return str(row[0]) if row else None return await self._tx(run) # ---- 원자적 claim ---- async def claim(self, worker_id: str, lease_sec: int = 120) -> dict | None: """대기 잡 1건을 원자적으로 점유. 없으면 None. FOR UPDATE SKIP LOCKED 로 잠근 행을 같은 UPDATE 에서 RUNNING 으로 전이 → 이중 할당 불가.""" sql = text(""" UPDATE jobs SET status = 2, worker_id = :wid, lease_until = now() + make_interval(secs => :lease), run_started_at = now(), attempts = attempts + 1, updated_at = now() WHERE job_id = ( SELECT job_id FROM jobs WHERE status = 1 AND run_after <= now() ORDER BY priority ASC, created_at ASC FOR UPDATE SKIP LOCKED LIMIT 1 ) RETURNING job_id, job_type, payload, attempts, max_attempts """) async def run(s): row = (await s.execute(sql, {"wid": worker_id, "lease": lease_sec})).mappings().first() if not row: return None d = dict(row) d["job_id"] = str(d["job_id"]) if isinstance(d.get("payload"), str): d["payload"] = json.loads(d["payload"]) return d return await self._tx(run) # ---- 완료/실패 (소유권 가드) ---- async def complete(self, job_id: str, worker_id: str, result: dict | None = None) -> bool: sql = text(""" UPDATE jobs SET status = 3, result = CAST(:result AS jsonb), lease_until = NULL, worker_id = NULL, updated_at = now() WHERE job_id = CAST(:id AS uuid) AND status = 2 AND worker_id = :wid RETURNING job_id """) async def run(s): row = (await s.execute(sql, { "id": job_id, "wid": worker_id, "result": json.dumps(result, ensure_ascii=False) if result is not None else None, })).first() return row is not None return await self._tx(run) async def fail(self, job_id: str, worker_id: str, error: str, backoff_sec: float = 5.0) -> int | None: """실패 처리. 시도 남으면 PENDING(run_after=백오프)으로 재큐, 소진되면 DEAD(dead-letter). 전이 후 status(JobStatus 값)를 반환. 소유 불일치면 None.""" sql = text(""" UPDATE jobs SET status = CASE WHEN attempts >= max_attempts THEN 4 ELSE 1 END, run_after = CASE WHEN attempts >= max_attempts THEN run_after ELSE now() + make_interval(secs => :backoff) END, last_error = :err, lease_until = NULL, worker_id = NULL, updated_at = now() WHERE job_id = CAST(:id AS uuid) AND status = 2 AND worker_id = :wid RETURNING status """) async def run(s): row = (await s.execute(sql, { "id": job_id, "wid": worker_id, "err": error[:2000], "backoff": backoff_sec, })).first() return int(row[0]) if row else None return await self._tx(run) # ---- lease 갱신(heartbeat) / 회수(reaper) ---- async def renew_lease(self, job_id: str, worker_id: str, lease_sec: int = 120) -> bool: sql = text(""" UPDATE jobs SET lease_until = now() + make_interval(secs => :lease), updated_at = now() WHERE job_id = CAST(:id AS uuid) AND worker_id = :wid AND status = 2 RETURNING job_id """) async def run(s): row = (await s.execute(sql, {"id": job_id, "wid": worker_id, "lease": lease_sec})).first() return row is not None return await self._tx(run) async def reap(self) -> list[str]: """만료된 lease(워커 사망 등)의 RUNNING 잡을 회수. 시도 남으면 즉시 재큐, 소진되면 DEAD. 회수된 job_id 목록 반환.""" sql = text(""" UPDATE jobs SET status = CASE WHEN attempts >= max_attempts THEN 4 ELSE 1 END, run_after = now(), last_error = COALESCE(last_error, '') || ' [lease-expired reclaim]', lease_until = NULL, worker_id = NULL, updated_at = now() WHERE status = 2 AND lease_until IS NOT NULL AND lease_until < now() RETURNING job_id """) async def run(s): rows = (await s.execute(sql)).all() return [str(r[0]) for r in rows] return await self._tx(run) # ---- 단건 조회 (상태 폴링) ---- async def get(self, job_id: str) -> dict | None: """잡 단건 조회(읽기). 없으면 None. status 는 정수(JobStatus 값).""" sql = text(""" SELECT job_id, job_type, status, priority, attempts, max_attempts, payload, result, last_error, run_after, run_started_at, created_at, updated_at FROM jobs WHERE job_id = CAST(:id AS uuid) """) async def run(s): row = (await s.execute(sql, {"id": job_id})).mappings().first() if not row: return None d = dict(row) d["job_id"] = str(d["job_id"]) for key in ("payload", "result"): if isinstance(d.get(key), str): d[key] = json.loads(d[key]) return d return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run) async def find_active(self, dedupe_key: str) -> dict | None: """dedupe_key 로 활성(PENDING/RUNNING) 잡을 찾는다. enqueue 가 중복으로 None 을 돌려줬을 때, 이미 돌고 있는 잡의 id 를 알려주기 위함.""" sql = text(""" SELECT job_id, job_type, status FROM jobs WHERE dedupe_key = :dk AND status IN (1, 2) LIMIT 1 """) async def run(s): row = (await s.execute(sql, {"dk": dedupe_key})).mappings().first() if not row: return None d = dict(row) d["job_id"] = str(d["job_id"]) return d return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run) # ---- 관측(관리 API/알림용) ---- async def counts(self) -> dict[str, int]: """상태별 잡 개수.""" async def run(s): rows = (await s.execute(text("SELECT status, count(*) FROM jobs GROUP BY status"))).all() by_val = {int(st): int(c) for st, c in rows} return {js.name: by_val.get(js.value, 0) for js in JobStatus} return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run) async def ops(self) -> dict: """운영 스냅샷(모니터링·알림용): 상태별 카운트 + 큐 지연(가장 오래된 PENDING 나이) + 최근 1시간 DEAD + stuck(좀비 신호). stuck 은 두 축 — lease 만료(워커 사망인데 reaper 미회수) OR 실행 10분 초과(핸들러 행 — heartbeat 가 lease 를 계속 갱신해 lease 축엔 안 잡히므로 run_started_at 으로 따로 본다).""" sql = text(""" SELECT count(*) FILTER (WHERE status = 1) AS pending, count(*) FILTER (WHERE status = 2) AS running, count(*) FILTER (WHERE status = 3) AS done, count(*) FILTER (WHERE status = 4) AS dead, count(*) FILTER (WHERE status = 4 AND updated_at > now() - interval '1 hour') AS dead_1h, count(*) FILTER (WHERE status = 2 AND ( (lease_until IS NOT NULL AND lease_until < now()) OR run_started_at < now() - interval '10 minutes' )) AS stuck_running, COALESCE(EXTRACT(EPOCH FROM (now() - min(created_at) FILTER (WHERE status = 1)))::int, 0) AS oldest_pending_sec, count(*) FILTER (WHERE last_error LIKE 'JobDeadlineExceeded%' AND updated_at > now() - interval '1 hour') AS deadline_1h FROM jobs """) async def run(s): return {k: int(v) for k, v in dict((await s.execute(sql)).mappings().first()).items()} return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run) async def requeue(self, job_id: str) -> str | None: """DEAD 잡 재큐(관리자 액션): attempts 리셋 + PENDING 전이 + 워커 깨움. DEAD 가 아니거나 없으면 None. 같은 dedupe_key 의 활성 잡이 있으면 부분 유니크 위반.""" sql = text(""" UPDATE jobs SET status = 1, attempts = 0, run_after = now(), lease_until = NULL, worker_id = NULL, run_started_at = NULL, last_error = NULL, updated_at = now() WHERE job_id = CAST(:jid AS uuid) AND status = 4 RETURNING job_id """) async def run(s): row = (await s.execute(sql, {"jid": job_id})).first() if row: await s.execute(text("SELECT pg_notify(:ch, '')"), {"ch": JOB_NOTIFY_CHANNEL}) return str(row[0]) if row else None return await self._tx(run)