FabledCurator/backend/app/tasks/download.py

"""download_source Celery task — runs DownloadService for one source."""

import asyncio
import logging
from datetime import UTC, datetime
from pathlib import Path

from celery.exceptions import SoftTimeLimitExceeded
from sqlalchemy import select
from sqlalchemy.exc import DBAPIError, OperationalError
from sqlalchemy.orm import Session as SyncSession

from ..celery_app import celery
from ..models import DownloadEvent, ImportSettings, Source
from ..services.credential_crypto import CredentialCrypto
from ..services.credential_service import CredentialService
from ..services.download_service import DownloadService
from ..services.gallery_dl import GalleryDLService
from ..services.importer import Importer
from ..services.thumbnailer import Thumbnailer
from ._async_session import async_session_factory
from .import_file import _sync_session_factory

log = logging.getLogger(__name__)

IMAGES_ROOT = Path("/images")
_KEY_PATH = IMAGES_ROOT / "secrets" / "credential_key.b64"

# Celery time budget for one download_source run. The ceiling that
# governs *clean* teardown is the SOFT limit: it raises a catchable
# SoftTimeLimitExceeded in-process, whereas the HARD limit SIGKILLs the
# worker (no chance to finalize). Both gallery-dl subprocess budgets
# (gallery_dl.py: _DEFAULT_GDL_TIMEOUT_SECONDS=870 tick,
# BACKFILL_CHUNK_SECONDS=600 per backfill chunk, plan #693) MUST sit below the soft limit
# so subprocess.run raises its own TimeoutExpired first — that path
# captures partial stdout/stderr and finalizes the DownloadEvent. soft is
# max-subprocess (1170) + ~180s phase-3 persist headroom; hard is soft +
# 150s SIGKILL backstop. Audit 2026-06-03 (Anduo #39912): the old
# soft=900 sat BELOW the 1170 backfill budget, so SoftTimeLimitExceeded
# preempted TimeoutExpired and the event stranded empty. The recovery
# sweep's DOWNLOAD_STALL_THRESHOLD_MINUTES (30 min) still trails the new
# 25-min hard kill by 5 min, so it stays a true backstop. Invariant
# guarded by test_timeout_ladder_keeps_subprocess_budgets_under_soft_limit.
DOWNLOAD_SOFT_TIME_LIMIT = 1350
DOWNLOAD_HARD_TIME_LIMIT = 1500


def _finalize_soft_limited(session: SyncSession, source_id: int) -> None:
    """Defense in depth for the soft-time-limit kill path.

    A SoftTimeLimitExceeded unwinds download_source before phase 3 can
    finalize the DownloadEvent, leaving it 'running' until the recovery
    sweep stamps a context-free "stranded" error 30 min later — AND
    leaving backfill_runs_remaining undecremented so the source re-runs
    and re-strands every tick (Anduo #39912, 2026-06-03). Flip the
    in-flight event to error with a real reason, mirror phase 3's
    source-health write, and decrement any backfill budget so a
    chronically-slow source self-heals back to tick mode.

    The caller owns the commit. All mutations are gated on actually
    finding a running event, so a benign late soft-limit (phase 3 already
    committed) is a no-op.
    """
    now = datetime.now(UTC)
    ev = session.execute(
        select(DownloadEvent)
        .where(DownloadEvent.source_id == source_id)
        .where(DownloadEvent.status == "running")
        .order_by(DownloadEvent.id.desc())
        .limit(1)
    ).scalar_one_or_none()
    if ev is None:
        return
    ev.status = "error"
    ev.finished_at = now
    ev.error = (
        f"killed by Celery soft time limit ({DOWNLOAD_SOFT_TIME_LIMIT}s) "
        "before the gallery-dl subprocess returned — the run exceeded its "
        "budget and its stdout/stderr were lost with the worker thread. "
        "If this recurs, the source is too large for one run; the backfill "
        "budget was decremented so the next tick walks less."
    )
    ev.metadata_ = {
        **(ev.metadata_ or {}),
        "error_type": "timeout",
        "soft_time_limited": True,
    }
    src = session.get(Source, source_id)
    if src is not None:
        src.consecutive_failures = (src.consecutive_failures or 0) + 1
        src.last_error = "soft time limit exceeded"
        src.error_type = "timeout"
        src.last_checked_at = now
        if (src.backfill_runs_remaining or 0) > 0:
            src.backfill_runs_remaining = max(0, src.backfill_runs_remaining - 1)


@celery.task(
    name="backend.app.tasks.download.download_source",
    bind=True,
    acks_late=True,
    autoretry_for=(OperationalError, DBAPIError, OSError),
    retry_backoff=10,
    retry_backoff_max=120,
    retry_jitter=True,
    max_retries=3,
    soft_time_limit=DOWNLOAD_SOFT_TIME_LIMIT,
    time_limit=DOWNLOAD_HARD_TIME_LIMIT,
)
def download_source(self, source_id: int) -> int:
    """Returns the DownloadEvent.id."""

    async def _run():
        async_factory, async_engine = async_session_factory()
        SyncFactory = _sync_session_factory()
        try:
            with SyncFactory() as sync_session:
                settings = ImportSettings.load_sync(sync_session)
                rate_limit = settings.download_rate_limit_seconds
                validate_files = settings.download_validate_files

            gdl = GalleryDLService(
                images_root=IMAGES_ROOT,
                rate_limit=rate_limit,
                validate_files=validate_files,
            )
            crypto = CredentialCrypto(_KEY_PATH)

            async with async_factory() as async_session:
                cred_service = CredentialService(async_session, crypto)
                with SyncFactory() as sync_session:
                    sync_settings = ImportSettings.load_sync(sync_session)
                    importer = Importer(
                        session=sync_session,
                        images_root=IMAGES_ROOT,
                        import_root=IMAGES_ROOT,
                        thumbnailer=Thumbnailer(images_root=IMAGES_ROOT),
                        settings=sync_settings,
                    )
                    svc = DownloadService(
                        async_session=async_session,
                        sync_session=sync_session,
                        gdl=gdl,
                        importer=importer,
                        cred_service=cred_service,
                    )
                    return await svc.download_source(source_id)
        finally:
            await async_engine.dispose()

    try:
        return asyncio.run(_run())
    except SoftTimeLimitExceeded:
        # phase 3 never ran — salvage the in-flight event so the operator
        # sees a real reason instead of the recovery sweep's generic
        # "stranded" 30 min later (Anduo #39912). Best-effort: a failure
        # here must not mask the timeout. Re-raise so Celery + the
        # task_run signal handler still record the kill.
        try:
            SyncFactory = _sync_session_factory()
            with SyncFactory() as session:
                _finalize_soft_limited(session, source_id)
                session.commit()
        except Exception:  # noqa: BLE001 — cleanup must not swallow the kill
            log.exception("soft-limit finalize failed for source %s", source_id)
        raise