FabledCurator/backend/app/celery_signals.py

"""FC-3i: task_run lifecycle via Celery signals.

Subscribes to task_prerun / task_postrun / task_failure / task_retry
and persists one task_run row per task attempt. Drop-in for every
existing and future Celery task — no per-task instrumentation.

Signal handlers run inside the worker process (sync context); DB
writes go through the existing shared sync engine
(backend.app.tasks._sync_engine.sync_session_factory) — one engine
per worker process, not per-task, so we don't blow Postgres
max_connections under load (the reason FC-3g shared-engine fix
existed).

Failure-mode discipline (operator-pressed point): every handler is
wrapped in try/except that swallows + logs. If the DB is down or the
handler has a bug, the real task still runs — the dashboard goes
dark for that interval. Monitoring NEVER breaks the thing it's
monitoring.
"""

import functools
import logging
from datetime import UTC, datetime

from celery.exceptions import SoftTimeLimitExceeded
from celery.signals import (
    task_failure,
    task_postrun,
    task_prerun,
    task_retry,
    worker_ready,
)

from .models import TaskRun
from .tasks._sync_engine import sync_session_factory

log = logging.getLogger(__name__)

# Celery-internal tasks that would generate dashboard noise without
# operational value. Conservative list; extend only when a specific
# task proves noisy.
_UNTRACKED_TASK_NAMES = frozenset({
    "celery.chord_unlock",
    "celery.backend_cleanup",
    "celery.chunks",
})

_MAX_ERROR_MESSAGE_LEN = 2000
_MAX_ARGS_SUMMARY_LEN = 255
_MAX_WORKER_HOSTNAME_LEN = 128

# PostgreSQL Integer is signed 32-bit. Tasks called with a first-arg
# int outside this range (e.g. an absurdly large mock value, or a
# string-of-digits coercible to int but bigger than 2^31-1) would crash
# the INSERT with NumericValueOutOfRange. Bound the recorded value to
# the column's range; values outside become None (target_id is
# nullable, so this is safe).
_INT32_MAX = 2_147_483_647
_INT32_MIN = -2_147_483_648


def _queue_for(task) -> str:
    """The queue Celery routes this task to — asked of the router itself.

    This was a hand-kept copy of `celery_app.task_routes`, and it drifted
    twice (the 2026-06-02 audit, then #4432). Long-lane jobs were recorded as
    `maintenance`, and translation and gpu_queue runs as `default`, where the
    5-minute stall sweep failed healthy 35-minute translation runs. The router
    answers from the same table the broker uses, so the two cannot disagree.
    """
    name = getattr(task, "name", "") or ""
    app = getattr(task, "app", None)
    if app is None:
        from .celery_app import celery as app
    return _routed_queue(app, name)


@functools.lru_cache(maxsize=1024)
def _routed_queue(app, name: str) -> str:
    try:
        queue = app.amqp.router.route({}, name).get("queue")
    except Exception:  # noqa: BLE001 — monitoring never breaks the task
        log.warning("task_run: could not resolve the queue for %s", name)
        return "default"
    return getattr(queue, "name", None) or (queue if isinstance(queue, str) else "default")


def _target_id_from_args(args) -> int | None:
    """Best-effort: if the first positional arg parses as int AND fits
    in the column's signed-32-bit range, record it as target_id
    (image_id, source_id, etc.). Never raises."""
    if not args:
        return None
    try:
        value = int(args[0])
    except (TypeError, ValueError):
        return None
    if value < _INT32_MIN or value > _INT32_MAX:
        return None
    return value


def _truncate(s, limit: int) -> str | None:
    if s is None:
        return None
    text = str(s)
    return text if len(text) <= limit else text[:limit]


def _is_tracked(task_name: str | None) -> bool:
    return bool(task_name) and task_name not in _UNTRACKED_TASK_NAMES


@task_prerun.connect
def _on_prerun(sender=None, task_id=None, task=None, args=None,
               kwargs=None, **_):
    if not _is_tracked(getattr(task, "name", None)):
        return
    try:
        Session = sync_session_factory()
        with Session() as session:
            session.add(TaskRun(
                celery_task_id=task_id or "",
                queue=_queue_for(task),
                task_name=task.name,
                target_id=_target_id_from_args(args),
                started_at=datetime.now(UTC),
                status="running",
                args_summary=_truncate(repr(args), _MAX_ARGS_SUMMARY_LEN),
                worker_hostname=_truncate(
                    getattr(sender, "hostname", None),
                    _MAX_WORKER_HOSTNAME_LEN,
                ),
            ))
            session.commit()
    except Exception:  # noqa: BLE001 — never break the worker
        log.exception("task_run prerun insert failed (task=%s)",
                      getattr(task, "name", "?"))


def _finalize(task_id: str, *, status: str,
              error_type: str | None = None,
              error_message: str | None = None,
              retry_count: int | None = None) -> None:
    """Shared write path for postrun/failure/retry. Picks the most-
    recent task_run row for this celery_task_id that's still 'running'
    (retries reuse the same celery_task_id; each new attempt's prerun
    inserts a fresh row, so finalize targets the latest running row)."""
    try:
        from sqlalchemy import select

        Session = sync_session_factory()
        now = datetime.now(UTC)
        with Session() as session:
            row = session.execute(
                select(TaskRun)
                .where(TaskRun.celery_task_id == task_id)
                .where(TaskRun.status == "running")
                .order_by(TaskRun.id.desc())
                .limit(1)
            ).scalar_one_or_none()
            if row is None:
                return  # no prerun row (untracked, insert failed, or already finalized)
            row.finished_at = now
            row.duration_ms = int(
                (now - row.started_at).total_seconds() * 1000
            )
            row.status = status
            if error_type is not None:
                row.error_type = _truncate(error_type, 128)
            if error_message is not None:
                row.error_message = _truncate(error_message, _MAX_ERROR_MESSAGE_LEN)
            if retry_count is not None:
                row.retry_count = retry_count
            session.commit()
    except Exception:  # noqa: BLE001
        log.exception("task_run finalize failed (task_id=%s)", task_id)


@task_postrun.connect
def _on_postrun(sender=None, task_id=None, task=None, args=None,
                kwargs=None, retval=None, state=None, **_):
    if not _is_tracked(getattr(task, "name", None)):
        return
    # state is one of SUCCESS/FAILURE/RETRY/etc. Only handle SUCCESS;
    # task_failure handles FAILURE explicitly (with the exception).
    if state != "SUCCESS":
        return
    _finalize(task_id, status="ok")


@task_failure.connect
def _on_failure(sender=None, task_id=None, exception=None,
                args=None, kwargs=None, einfo=None, **_):
    if not _is_tracked(getattr(sender, "name", None)):
        return
    status = ("timeout"
              if isinstance(exception, SoftTimeLimitExceeded)
              else "error")
    _finalize(
        task_id, status=status,
        error_type=type(exception).__name__ if exception else "Unknown",
        error_message=str(exception) if exception else None,
    )


@task_retry.connect
def _on_retry(sender=None, request=None, reason=None, einfo=None, **_):
    if not _is_tracked(getattr(sender, "name", None)):
        return
    task_id = getattr(request, "id", None)
    if not task_id:
        return
    # Mark current attempt's row as 'retry' (terminal for this row).
    # The next attempt's task_prerun inserts a fresh row.
    _finalize(
        task_id, status="retry",
        error_type=type(reason).__name__ if reason else "Retry",
        error_message=str(reason) if reason else None,
        retry_count=getattr(request, "retries", 0),
    )


def _consumed_queues(consumer) -> set[str]:
    """The queue names this worker process consumes (its `-Q`), or empty when
    they can't be read — which makes the boot hook below a no-op, not a guess."""
    try:
        return {q.name for q in consumer.task_consumer.queues}
    except Exception:  # noqa: BLE001 — shape varies across celery versions
        return set()


@worker_ready.connect
def _on_worker_ready(sender=None, **_):
    """The download lane clears what its previous process left behind (#4433).

    A restart SIGKILLs any walk that outlives the stop grace, so its event is
    never finalized and its platform lock is never released. Both used to wait
    out timers — a 30-min stall sweep that then blamed the source, and a 27-min
    lock TTL that stalled every other source on the platform. Only the process
    consuming `download` does this; the other lanes booting beside it must not.
    Best-effort: a failure here is logged and the worker still starts.
    """
    if "download" not in _consumed_queues(sender):
        return
    booted_at = datetime.now(UTC)
    try:
        from .services.download_recovery import interrupt_orphaned_download_events
        from .services.platform_lock import release_all_platform_locks

        with sync_session_factory()() as session:
            closed = interrupt_orphaned_download_events(session, booted_at=booted_at)
            session.commit()
        released = release_all_platform_locks()
        if closed or released:
            log.info(
                "download lane boot: closed %d orphaned download event(s) as "
                "interrupted, released %d platform lock(s)", closed, released,
            )
    except Exception:  # noqa: BLE001 — never block the worker from starting
        log.exception("download lane boot recovery failed")