o2o-site-AEO/backend/worker/runner.py

"""워커 루프 + reaper. (LPS `worker/runner.py` 이식)

워커는 큐에서 잡을 원자적으로 claim → 핸들러 실행 → complete/fail 한다.
- 처리 중 heartbeat 로 lease 를 갱신(긴 잡이 reaper 에 회수되지 않게).
  수집·비전분석은 몇 분이 정상이라 heartbeat 없이는 lease 가 먼저 만료된다.
- 핸들러엔 데드라인(job_deadline_sec)을 건다 — heartbeat 가 lease 를 계속 갱신하므로
  핸들러가 행하면 reaper 로는 영원히 회수 불가.
  초과 시 취소 후 fail 처리 → 백오프 재큐(소진 시 DEAD), 워커 슬롯은 즉시 다음 잡으로.
- claim 이 비면 LISTEN 알림 또는 짧은 타임아웃까지 대기(폴링 최소화).
- 핸들러는 주입식(async def(job)->dict) — 프로덕션은 수집 파이프라인, 테스트는 fake.
"""

import asyncio

from common.enums import JobStatus
from common.logger import LOG
from crud.job_crud import JobQueue, compute_backoff


class Worker:
    """큐에서 잡을 원자적으로 claim → 핸들러 실행 → complete/fail 하는 워커 루프."""

    def __init__(
        self,
        worker_id: str,
        queue: JobQueue,
        handler,
        lease_sec: int = 120,
        backoff_fn=compute_backoff,
        job_deadline_sec: float = 900.0,
    ):
        self.worker_id = worker_id
        self.queue = queue
        self.handler = handler
        self.lease_sec = lease_sec
        self.backoff_fn = backoff_fn
        # 잡 1건 처리 시간 상한(0 이면 무제한 — 테스트용).
        # 수집 파이프라인은 Perplexity(10~30s) + 크롤링 + Vision(사진 20~50장) 이라 15분을 기본으로 둔다.
        self.job_deadline_sec = job_deadline_sec

    async def process_one(self) -> bool:
        """대기 잡 1건을 claim·처리. 처리했으면 True, 없으면 False."""
        job = await self.queue.claim(self.worker_id, self.lease_sec)
        if not job:
            return False
        await self._process(job)
        return True

    async def drain(self) -> int:
        """큐가 빌 때까지 처리(테스트/일회성 배치용). 처리한 잡 수 반환."""
        n = 0
        while await self.process_one():
            n += 1
        return n

    async def run(self, listener=None, stop: asyncio.Event | None = None, idle_timeout: float = 5.0):
        """상시 루프. stop 이 설정될 때까지 처리하고, 유휴 시 알림/타임아웃까지 대기."""
        stop = stop or asyncio.Event()
        while not stop.is_set():
            try:
                worked = await self.process_one()
            except Exception as ex:
                # claim 자체가 실패(DB 순단 등) — 루프를 죽이지 않는다. 잠깐 쉬고 다시 시도.
                LOG.e_no_callstack(f"[{self.worker_id}] claim 실패(계속): {type(ex).__name__}: {ex}")
                worked = False
            if not worked:
                if listener is not None:
                    await listener.wait(idle_timeout)
                else:
                    await asyncio.sleep(idle_timeout)

    async def _process(self, job: dict):
        jid = job["job_id"]
        hb = asyncio.create_task(self._heartbeat(jid))
        try:
            if self.job_deadline_sec > 0:
                result = await asyncio.wait_for(self.handler(job), timeout=self.job_deadline_sec)
            else:
                result = await self.handler(job)
            await self.queue.complete(jid, self.worker_id, result)
            LOG.d(f"[{self.worker_id}] done {jid}")
        except (asyncio.TimeoutError, TimeoutError):
            # 데드라인 초과 — wait_for 가 핸들러 태스크를 취소한 뒤 여기로 온다.
            # 행이 워커 슬롯을 영구 점유하는 것보다 낫다.
            backoff = self.backoff_fn(job["attempts"])
            st = await self.queue.fail(jid, self.worker_id, f"JobDeadlineExceeded: {self.job_deadline_sec:.0f}s", backoff)
            LOG.w(f"[{self.worker_id}] deadline {jid} → {JobStatus(st).name if st else '?'} "
                  f"({self.job_deadline_sec:.0f}s 초과, 핸들러 취소)")
        except Exception as ex:
            backoff = self.backoff_fn(job["attempts"])
            st = await self.queue.fail(jid, self.worker_id, f"{type(ex).__name__}: {ex}", backoff)
            LOG.w(f"[{self.worker_id}] fail {jid} → {JobStatus(st).name if st else '?'} ({type(ex).__name__}: {ex})")
        finally:
            hb.cancel()
            try:
                await hb
            except asyncio.CancelledError:
                pass

    async def _heartbeat(self, jid: str):
        interval = max(1, self.lease_sec // 3)
        while True:
            await asyncio.sleep(interval)
            try:
                await self.queue.renew_lease(jid, self.worker_id, self.lease_sec)
            except Exception as ex:
                # 갱신 실패는 치명적이지 않다(다음 틱 재시도). 계속 실패하면 lease 만료 → reaper 회수.
                LOG.w(f"[{self.worker_id}] lease 갱신 실패(계속): {type(ex).__name__}")


async def run_reaper(queue: JobQueue, stop: asyncio.Event, interval: float = 30.0):
    """만료 lease(워커 사망) 잡을 주기적으로 회수. 재시도 남으면 재큐, 소진되면 DEAD.

    도커에서 워커 컨테이너가 재시작되면 진행 중이던 잡이 여기서 되살아난다."""
    while not stop.is_set():
        try:
            reclaimed = await queue.reap()
            if reclaimed:
                LOG.w(f"[reaper] reclaimed {len(reclaimed)} stale job(s)")
        except Exception as ex:
            LOG.e_no_callstack(f"[reaper] {type(ex).__name__}: {ex}")
        try:
            await asyncio.wait_for(stop.wait(), interval)
        except asyncio.TimeoutError:
            pass