RAG/backend/app/worker.py

"""Async, lease-fenced background worker runtime.

Business handlers are registered by dependency injection. This module owns only
queue lifecycle, heartbeats, retries, maintenance, and graceful process stop.
"""

from __future__ import annotations

import asyncio
import logging
import os
import signal
import socket
import time
import uuid
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import Protocol

from app.core.config import Settings
from app.persistence.job_queue import (
    BackgroundJob,
    JobLease,
    JobState,
    LeaseHeartbeat,
    LeaseLostError,
    PsycopgJobQueue,
)

LOGGER = logging.getLogger("geological_rag.worker")
DEFAULT_REAPER_LOCK_KEY = 7_221_016_471_511_937

type JobHandler = Callable[[BackgroundJob], Awaitable[None]]


class JobQueue(Protocol):
    def claim(
        self,
        *,
        worker_id: str,
        worker_capabilities: Sequence[str],
        lease_seconds: int,
    ) -> BackgroundJob | None: ...

    def heartbeat(self, lease: JobLease, *, lease_seconds: int) -> LeaseHeartbeat: ...

    def complete(self, lease: JobLease) -> JobState: ...

    def fail_or_retry(
        self,
        lease: JobLease,
        *,
        error_code: str,
        error_message: str,
        retry_delay_seconds: int,
    ) -> JobState: ...

    def reap_expired(
        self,
        *,
        lock_key: int,
        batch_size: int = 100,
    ) -> tuple[JobState, ...]: ...


@dataclass(frozen=True, slots=True)
class WorkerConfig:
    worker_id: str
    capabilities: tuple[str, ...]
    lease_seconds: int = 60
    heartbeat_seconds: float = 20.0
    poll_seconds: float = 1.0
    retry_delay_seconds: int = 30
    reaper_interval_seconds: float = 30.0
    reaper_batch_size: int = 100
    reaper_lock_key: int = DEFAULT_REAPER_LOCK_KEY

    def __post_init__(self) -> None:
        if not self.worker_id.strip() or len(self.worker_id) > 200:
            raise ValueError("worker_id must contain 1 to 200 characters")
        if not self.capabilities or any(not item.strip() for item in self.capabilities):
            raise ValueError("capabilities must not be empty")
        if len(self.capabilities) != len(set(self.capabilities)):
            raise ValueError("capabilities must not contain duplicates")
        if isinstance(self.lease_seconds, bool) or not 1 <= self.lease_seconds <= 86_400:
            raise ValueError("lease_seconds must be between 1 and 86400")
        if not 0 < self.heartbeat_seconds <= self.lease_seconds / 3:
            raise ValueError(
                "heartbeat_seconds must be positive and no longer than one third of the lease"
            )
        if not 0 < self.poll_seconds <= 60:
            raise ValueError("poll_seconds must be between 0 and 60")
        if (
            isinstance(self.retry_delay_seconds, bool)
            or not 0 <= self.retry_delay_seconds <= 86_400
        ):
            raise ValueError("retry_delay_seconds must be between 0 and 86400")
        if not 0 < self.reaper_interval_seconds <= 3600:
            raise ValueError("reaper_interval_seconds must be between 0 and 3600")
        if isinstance(self.reaper_batch_size, bool) or not 1 <= self.reaper_batch_size <= 1000:
            raise ValueError("reaper_batch_size must be between 1 and 1000")
        if isinstance(self.reaper_lock_key, bool) or not -(2**63) <= self.reaper_lock_key < 2**63:
            raise ValueError("reaper_lock_key must be a signed 64-bit integer")


class Worker:
    """Single-concurrency worker with fenced heartbeats and terminal updates."""

    def __init__(
        self,
        queue: JobQueue,
        config: WorkerConfig,
        *,
        handlers: Mapping[str, JobHandler],
        monotonic: Callable[[], float] | None = None,
    ) -> None:
        invalid_job_types = [name for name in handlers if not name.strip()]
        if invalid_job_types:
            raise ValueError("handler job types must not be empty")
        self._queue = queue
        self._config = config
        self._handlers = dict(handlers)
        self._stop = asyncio.Event()
        self._monotonic = monotonic or time.monotonic
        self._next_reap_at = 0.0

    @property
    def stopping(self) -> bool:
        return self._stop.is_set()

    def request_stop(self) -> None:
        """Stop claiming new work; an in-flight fenced job drains gracefully."""

        self._stop.set()

    async def run(self) -> None:
        LOGGER.info("worker_started", extra={"worker_id": self._config.worker_id})
        try:
            while not self._stop.is_set():
                try:
                    worked = await self.run_once()
                except Exception as exc:  # noqa: BLE001 - iteration boundary stays alive
                    LOGGER.error(
                        "worker_iteration_failed",
                        extra={"error_type": type(exc).__name__},
                    )
                    worked = False
                if not worked:
                    await self._wait_for_stop(self._config.poll_seconds)
        finally:
            LOGGER.info("worker_stopped", extra={"worker_id": self._config.worker_id})

    async def run_once(self) -> bool:
        """Run maintenance and at most one job; useful for deterministic supervision."""

        await self._reap_if_due()
        if self._stop.is_set():
            return False
        job = await asyncio.to_thread(
            self._queue.claim,
            worker_id=self._config.worker_id,
            worker_capabilities=self._config.capabilities,
            lease_seconds=self._config.lease_seconds,
        )
        if job is None:
            return False
        await self._process(job)
        return True

    async def _reap_if_due(self) -> None:
        now = self._monotonic()
        if now < self._next_reap_at:
            return
        reaped = await asyncio.to_thread(
            self._queue.reap_expired,
            lock_key=self._config.reaper_lock_key,
            batch_size=self._config.reaper_batch_size,
        )
        self._next_reap_at = now + self._config.reaper_interval_seconds
        if reaped:
            LOGGER.warning("expired_job_leases_reaped", extra={"count": len(reaped)})

    async def _process(self, job: BackgroundJob) -> None:
        handler = self._handlers.get(job.job_type)
        if handler is None:
            await self._safe_fail(
                job.lease,
                error_code="UNKNOWN_JOB_TYPE",
                error_message="No registered handler exists for this job type.",
            )
            return

        try:
            await self._invoke_with_heartbeat(job, handler)
        except LeaseLostError:
            LOGGER.warning("job_lease_lost", extra={"job_id": str(job.id)})
            return
        except asyncio.CancelledError:
            raise
        except Exception as exc:  # noqa: BLE001 - handler failures become safe queue metadata
            LOGGER.error(
                "job_handler_failed error_type=%s",
                type(exc).__name__,
                extra={
                    "job_id": str(job.id),
                    "job_type": job.job_type,
                    "error_type": type(exc).__name__,
                },
            )
            await self._safe_fail(
                job.lease,
                error_code="JOB_HANDLER_FAILED",
                error_message="Registered job handler failed.",
            )
            return

        try:
            await asyncio.to_thread(self._queue.complete, job.lease)
        except LeaseLostError:
            LOGGER.warning("job_completion_fence_rejected", extra={"job_id": str(job.id)})

    async def _safe_fail(
        self,
        lease: JobLease,
        *,
        error_code: str,
        error_message: str,
    ) -> None:
        try:
            await asyncio.to_thread(
                self._queue.fail_or_retry,
                lease,
                error_code=error_code,
                error_message=error_message,
                retry_delay_seconds=self._config.retry_delay_seconds,
            )
        except LeaseLostError:
            LOGGER.warning(
                "job_failure_fence_rejected",
                extra={"job_id": str(lease.job_id)},
            )

    async def _invoke_with_heartbeat(
        self,
        job: BackgroundJob,
        handler: JobHandler,
    ) -> None:
        heartbeat_stop = asyncio.Event()
        handler_task: asyncio.Future[None] = asyncio.ensure_future(handler(job))
        heartbeat_task = asyncio.create_task(
            self._heartbeat_loop(job.lease, heartbeat_stop),
            name=f"heartbeat-{job.id}",
        )

        try:
            done, _ = await asyncio.wait(
                (handler_task, heartbeat_task),
                return_when=asyncio.FIRST_COMPLETED,
            )
            if heartbeat_task in done:
                heartbeat_error = heartbeat_task.exception()
                if heartbeat_error is not None:
                    raise heartbeat_error
                if not handler_task.done():
                    raise RuntimeError("heartbeat stopped before the handler completed")

            heartbeat_stop.set()
            await heartbeat_task
            await handler_task
        except BaseException:
            heartbeat_stop.set()
            for task in (handler_task, heartbeat_task):
                if not task.done():
                    task.cancel()
            await asyncio.gather(handler_task, heartbeat_task, return_exceptions=True)
            raise

    async def _heartbeat_loop(self, lease: JobLease, stop: asyncio.Event) -> None:
        while not await self._wait_for_event(stop, self._config.heartbeat_seconds):
            await asyncio.to_thread(
                self._queue.heartbeat,
                lease,
                lease_seconds=self._config.lease_seconds,
            )

    async def _wait_for_stop(self, wait_seconds: float) -> None:
        await self._wait_for_event(self._stop, wait_seconds)

    @staticmethod
    async def _wait_for_event(event: asyncio.Event, wait_seconds: float) -> bool:
        try:
            await asyncio.wait_for(event.wait(), timeout=wait_seconds)
        except TimeoutError:
            return False
        return True


def _environment_integer(name: str, default: int, minimum: int, maximum: int) -> int:
    raw_value = os.getenv(name)
    if raw_value is None:
        return default
    try:
        value = int(raw_value)
    except ValueError as exc:
        raise RuntimeError(f"{name} must be an integer") from exc
    if not minimum <= value <= maximum:
        raise RuntimeError(f"{name} is outside its allowed range")
    return value


def _environment_float(name: str, default: float, minimum: float, maximum: float) -> float:
    raw_value = os.getenv(name)
    if raw_value is None:
        return default
    try:
        value = float(raw_value)
    except ValueError as exc:
        raise RuntimeError(f"{name} must be numeric") from exc
    if not minimum <= value <= maximum:
        raise RuntimeError(f"{name} is outside its allowed range")
    return value


def config_from_environment(capability_csv: str) -> WorkerConfig:
    capabilities = tuple(item.strip() for item in capability_csv.split(",") if item.strip())
    lease_seconds = _environment_integer("WORKER_LEASE_SECONDS", 60, 1, 86_400)
    default_heartbeat = min(20.0, lease_seconds / 3)
    worker_id = os.getenv("WORKER_ID") or (
        f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4().hex[:8]}"
    )
    return WorkerConfig(
        worker_id=worker_id,
        capabilities=capabilities,
        lease_seconds=lease_seconds,
        heartbeat_seconds=_environment_float(
            "WORKER_HEARTBEAT_SECONDS",
            default_heartbeat,
            0.1,
            86_399,
        ),
        poll_seconds=_environment_float("WORKER_POLL_SECONDS", 1.0, 0.01, 60),
        retry_delay_seconds=_environment_integer("WORKER_RETRY_DELAY_SECONDS", 30, 0, 86_400),
        reaper_interval_seconds=_environment_float(
            "WORKER_REAPER_INTERVAL_SECONDS", 30.0, 0.1, 3600
        ),
        reaper_batch_size=_environment_integer("WORKER_REAPER_BATCH_SIZE", 100, 1, 1000),
        reaper_lock_key=_environment_integer(
            "WORKER_REAPER_LOCK_KEY",
            DEFAULT_REAPER_LOCK_KEY,
            -(2**63),
            2**63 - 1,
        ),
    )


def build_worker(
    *,
    settings: Settings | None = None,
    handlers: Mapping[str, JobHandler] | None = None,
) -> Worker:
    runtime_settings = settings or Settings()
    dsn = runtime_settings.database_url().set(drivername="postgresql")
    queue = PsycopgJobQueue(dsn.render_as_string(hide_password=False))
    config = config_from_environment(runtime_settings.worker_capabilities)
    registered_handlers = (
        dict(handlers)
        if handlers is not None
        else build_default_handlers(runtime_settings, config.capabilities)
    )
    return Worker(
        queue,
        config,
        handlers=registered_handlers,
    )


def build_default_handlers(
    settings: Settings,
    capabilities: Sequence[str],
) -> dict[str, JobHandler]:
    """Build only handlers whose trust boundary is granted to this process."""

    requested = set(capabilities)
    supported = {"document_parse", "embedding"}
    unsupported = requested - supported
    if unsupported:
        names = ", ".join(sorted(unsupported))
        raise RuntimeError(f"unsupported worker capabilities: {names}")

    registered: dict[str, JobHandler] = {}
    if "document_parse" in requested:
        from app.workers.document_jobs import build_document_handlers

        registered.update(build_document_handlers(settings))
    if "embedding" in requested:
        from app.workers.indexing_jobs import build_indexing_handlers

        registered.update(build_indexing_handlers(settings))
    return registered


def install_signal_handlers(
    worker: Worker,
    loop: asyncio.AbstractEventLoop,
) -> tuple[signal.Signals, ...]:
    installed: list[signal.Signals] = []
    for process_signal in (signal.SIGTERM, signal.SIGINT):
        try:
            loop.add_signal_handler(process_signal, worker.request_stop)
        except NotImplementedError:
            continue
        installed.append(process_signal)
    return tuple(installed)


async def run_worker(worker: Worker) -> None:
    loop = asyncio.get_running_loop()
    installed = install_signal_handlers(worker, loop)
    try:
        await worker.run()
    finally:
        for process_signal in installed:
            loop.remove_signal_handler(process_signal)


def main() -> None:
    logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO").upper())
    asyncio.run(run_worker(build_worker()))


if __name__ == "__main__":
    main()