o2o-negosium-original/lps/crud/job_crud.py

"""작업 큐 CRUD — PostgreSQL 을 '제대로' 큐로 쓴다.

레퍼런스의 반면교사를 전부 뒤집는다:
 - 할당은 **단일 문장 원자 claim**: FOR UPDATE SKIP LOCKED 서브쿼리 + 같은 UPDATE + RETURNING.
   → 워커/디스패처가 몇이든 같은 잡 이중 할당이 원천 불가. fetch 와 claim 을 분리하지 않는다.
 - 모든 전이는 **조건부 CAS**(WHERE 에 status/worker_id 가드) + RETURNING.
 - 복구는 timeout 추측이 아니라 **lease 만료 소유권**(reaper 가 회수).
 - 재시도/백오프/dead-letter 를 큐에 내장(스크립트 난립 제거). 외부전송도 OUTBOX 잡으로.
"""

import json

from sqlalchemy import text

from common.database.db_session_manager import DB_SESSION_MNG
from common.enums import DBType, DBWRType, JobStatus


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):
        """쓰기 트랜잭션(commit/rollback 은 여기서 책임)."""
        s = await DB_SESSION_MNG.start_session(self.DB, DBWRType.DB_WRITE.value)
        try:
            res = await fn(s)
            await s.commit()
            return res
        except Exception:
            await s.rollback()
            raise
        finally:
            await DB_SESSION_MNG.end_session(self.DB, DBWRType.DB_WRITE.value)

    # ---- 적재 -----------------------------------------------------------
    async def enqueue(self, job_type: int, payload: dict, priority: int = 100, dedupe_key: str | None = None, max_attempts: int = 3):
        """잡 적재. dedupe_key 가 활성(PENDING/RUNNING) 중복이면 삽입 없이 None 반환."""
        sql = text("""
            INSERT INTO job (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),
                "dk": dedupe_key, "ma": max_attempts,
            })).first()
            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):
        """대기 잡 1건을 원자적으로 점유. 없으면 None.
        FOR UPDATE SKIP LOCKED 로 잠근 행을 같은 UPDATE 에서 RUNNING 으로 전이 → 이중 할당 불가."""
        sql = text("""
            UPDATE job 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 job
                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 job SET status = 3, result = CAST(:result AS jsonb),
                           lease_until = NULL, worker_id = NULL, updated_at = now()
            WHERE job_id = :id 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) 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 job 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 = :id 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 job SET lease_until = now() + make_interval(secs => :lease), updated_at = now()
            WHERE job_id = :id 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 job 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,
                   result, last_error, run_after, created_at, updated_at
            FROM job WHERE job_id = :id
        """)
        s = await DB_SESSION_MNG.start_session(self.DB, DBWRType.DB_READ.value)
        try:
            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"])
            if isinstance(d.get("result"), str):
                d["result"] = json.loads(d["result"])
            return d
        finally:
            await DB_SESSION_MNG.end_session(self.DB, DBWRType.DB_READ.value)

    # ---- 관측(관리 API/메트릭용) ---------------------------------------
    async def counts(self) -> dict[str, int]:
        """상태별 잡 개수(관리 API·알림용). 수동 psql 스크립트를 대체한다."""
        async def run(s):
            rows = (await s.execute(text("SELECT status, count(*) FROM job 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 self._tx(run)