pdf/gateway/app/services/convert/cancellation.py

"""Per-conversion cancellation and deadlines.

Two mechanisms, both scoped to the conversion that owns them:

``CancelToken``
    A caller-supplied "should I stop?" predicate, used to honour a job DELETE.

``Deadline``
    A wall-clock budget the pipeline checks at page boundaries, so a runaway
    document stops doing work rather than merely having its result discarded.

Both travel in :mod:`contextvars` rather than module globals. A module global
is shared by every thread in the process, so with concurrent conversions one
job would overwrite another's callback — and cancelling one job would abort an
unrelated one. Context variables give each conversion its own value whether it
runs in a thread, a task, or nested inside a worker.
"""

from __future__ import annotations

import time
from collections.abc import Callable, Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass

from app.services import logs

logger = logs.get_logger(__name__)


class ConversionCancelled(RuntimeError):
    """Raised when the caller cancelled the conversion."""


class ConversionDeadlineExceeded(TimeoutError):
    """Raised when a conversion exhausted its wall-clock budget."""


CancelToken = Callable[[], bool]

_cancel_check: ContextVar[CancelToken | None] = ContextVar("convert_cancel_check", default=None)
_deadline: ContextVar[Deadline | None] = ContextVar("convert_deadline", default=None)
# Whether this conversion has already reported a broken cancel token. The check
# runs at every page boundary, so an unguarded log would produce one line per
# page for a single underlying fault.
_cancel_fault_reported: ContextVar[bool] = ContextVar(
    "convert_cancel_fault_reported", default=False
)


def _ask_cancel_token(fn: CancelToken) -> bool:
    """Ask the caller's token whether to stop, surviving a token that raises.

    Treating a raising token as "not cancelled" is the only safe answer — the
    alternative is aborting work nobody asked to abort — but doing it silently
    meant a broken token made every job **quietly uncancellable**, with no
    trace anywhere. The answer is unchanged; the silence is not.
    """
    try:
        return bool(fn())
    except Exception as exc:
        if not _cancel_fault_reported.get():
            _cancel_fault_reported.set(True)
            logs.error(
                logger,
                "cancel_token_failed",
                error=type(exc).__name__,
                detail=str(exc)[:200],
                consequence="job cannot be cancelled",
                exc_info=True,
            )
        return False


@dataclass
class Deadline:
    """A wall-clock budget for one conversion."""

    seconds: float
    started: float

    @classmethod
    def start(cls, seconds: float) -> Deadline:
        return cls(seconds=seconds, started=time.monotonic())

    @property
    def elapsed(self) -> float:
        return time.monotonic() - self.started

    @property
    def remaining(self) -> float:
        return self.seconds - self.elapsed

    @property
    def expired(self) -> bool:
        return self.seconds > 0 and self.remaining <= 0


@contextmanager
def conversion_scope(
    *, cancel_check: CancelToken | None = None, timeout: float | None = None
) -> Iterator[None]:
    """Bind a cancel token and/or deadline for the duration of one conversion."""
    cancel_reset = _cancel_check.set(cancel_check)
    deadline_reset = _deadline.set(Deadline.start(timeout) if timeout and timeout > 0 else None)
    fault_reset = _cancel_fault_reported.set(False)
    try:
        yield
    finally:
        _cancel_check.reset(cancel_reset)
        _deadline.reset(deadline_reset)
        _cancel_fault_reported.reset(fault_reset)


@contextmanager
def deadline_scope(seconds: float | None) -> Iterator[None]:
    """Bind only a wall-clock budget, leaving any cancel token in place."""
    reset = _deadline.set(Deadline.start(seconds) if seconds and seconds > 0 else None)
    try:
        yield
    finally:
        _deadline.reset(reset)


def set_cancel_check(fn: CancelToken | None) -> None:
    """Bind a cancel token for the current context.

    Prefer :func:`conversion_scope`; this exists for callers and tests that
    manage the lifetime themselves.
    """
    _cancel_check.set(fn)


def get_cancel_check() -> CancelToken | None:
    return _cancel_check.get()


def get_deadline() -> Deadline | None:
    return _deadline.get()


def remaining_seconds(default: float) -> float:
    """Budget left for a sub-step, never more than the conversion's own budget."""
    deadline = _deadline.get()
    if deadline is None or deadline.seconds <= 0:
        return default
    return max(0.0, min(default, deadline.remaining))


def check_cancelled(stage: str = "conversion") -> None:
    """Raise only if the caller cancelled — never on an expired deadline.

    For loops that handle an exhausted budget themselves by stopping early and
    returning what they have. Using :func:`check` there defeats the point: a
    single page overrunning the budget raises out of the loop and discards
    every page already recovered, which is exactly the failure the budget
    logic exists to prevent.
    """
    fn = _cancel_check.get()
    if fn is None:
        return
    if _ask_cancel_token(fn):
        raise ConversionCancelled(f"{stage} cancelled.")


def check(stage: str = "conversion") -> None:
    """Raise if the conversion was cancelled or has run out of time.

    Call at safe points — page boundaries, between OCR pages — so a runaway
    document stops working instead of merely having its result thrown away.
    """
    fn = _cancel_check.get()
    if fn is not None and _ask_cancel_token(fn):
        raise ConversionCancelled("Job cancelled.")
    deadline = _deadline.get()
    if deadline is not None and deadline.expired:
        raise ConversionDeadlineExceeded(f"{stage} exceeded its {deadline.seconds:.0f}s budget")