FabledCurator/backend/app/services/platform_lock.py

"""Per-platform download concurrency cap.

Some platforms (Patreon) are API-rate-sensitive enough that two *simultaneous*
walks can trip the server's rate limit even with each source pacing its own
requests. The platform-cooldown handles the AFTERMATH of a 429; this is the
preventive half — it serializes downloads PER PLATFORM to one at a time.

Different platforms still run concurrently up to the worker's concurrency; only
a second walk on the SAME serialized platform waits. The lock lives in Redis
(the Celery broker) with a TTL, so a SIGKILL'd worker can't wedge a platform —
the lock auto-expires shortly after the download hard time limit.
"""
from __future__ import annotations

import logging

import redis

from ..config import get_config
from .download_backends import NATIVE_INGESTER_PLATFORMS

log = logging.getLogger(__name__)

# Platforms walked one-at-a-time: every native-ingester platform, one paced
# scrape/API walk at a time. gallery-dl platforms are not: each runs as a
# self-pacing subprocess and they are lower-volume. Discord most of all needs
# it: every source walks on the operator's ONE user token, and parallel walks
# on a user account are both how its rate limit trips and what gets it flagged.
SERIALIZED_PLATFORMS = NATIVE_INGESTER_PLATFORMS

_LOCK_PREFIX = "fc:download_lock:"

_client: redis.Redis | None = None


def _redis() -> redis.Redis:
    # One client per worker process (Celery prefork forks before tasks run, so
    # each process lazily builds its own). redis-py pools connections.
    global _client
    if _client is None:
        _client = redis.from_url(get_config().celery_broker_url)
    return _client


def platform_lock(platform: str, *, ttl_seconds: int):
    """A non-blocking Redis lock for `platform`, or None when the platform is
    not serialized. Caller does `.acquire(blocking=False)` / `.release()`.

    Returns None (rather than raising) on any Redis error so a broker hiccup
    degrades to the prior behaviour (uncapped) instead of stalling downloads.
    """
    if platform not in SERIALIZED_PLATFORMS:
        return None
    try:
        return _redis().lock(
            f"{_LOCK_PREFIX}{platform}",
            timeout=ttl_seconds,
            blocking=False,
        )
    except redis.RedisError as exc:  # pragma: no cover - broker outage
        log.warning("platform_lock unavailable for %s: %s", platform, exc)
        return None


def release_all_platform_locks() -> int:
    """Drop every serialized platform's lock. Returns how many were held.

    Only for the download lane's boot (#4433). A worker that restarts mid-walk
    is SIGKILLed past its stop grace, so its `finally` never releases the lock,
    and the TTL keeps every other source on that platform bouncing for up to
    27 minutes after the new worker is ready. At boot no walk of ours can be
    running, so a held lock names a dead one. Assumes one download consumer —
    the only shape FC deploys; a second replica booting would free a live
    walk's lock (not corrupt it: the walk runs on, a second walk may overlap it).
    """
    try:
        client = _redis()
        return int(client.delete(*(f"{_LOCK_PREFIX}{p}" for p in SERIALIZED_PLATFORMS)))
    except redis.RedisError as exc:  # pragma: no cover - broker outage
        log.warning("could not release platform locks at boot: %s", exc)
        return 0