gh-christianlouis-docuelevate/app/middleware/upload_rate_limit.py

"""Per-user, health-aware upload rate limiter for DocuElevate.

This module provides a FastAPI dependency that enforces per-user upload rate
limits using a Redis-backed sliding window counter.  The effective limit is
dynamically reduced when the system is under heavy load (high Celery queue
depth or elevated CPU load average), ensuring the server remains responsive
to all users even during bulk-upload scenarios.

Usage in an endpoint::

    from app.middleware.upload_rate_limit import require_upload_rate_limit

    @router.post("/ui-upload")
    @require_login
    async def ui_upload(
        request: Request,
        _rate_ok: None = Depends(require_upload_rate_limit),
        ...
    ):
        ...

See ``docs/ConfigurationGuide.md`` for the configuration options
(``UPLOAD_RATE_LIMIT_PER_USER``, ``UPLOAD_RATE_LIMIT_WINDOW``).
"""

from __future__ import annotations

import logging
import os
import time
from typing import Any

import redis
from fastapi import HTTPException, Request, status

from app.config import settings
from app.utils.user_scope import get_current_owner_id

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Redis key prefix
# ---------------------------------------------------------------------------
_KEY_PREFIX = "docuelevate:upload_rate"

# ---------------------------------------------------------------------------
# Health-check queue names (Celery defaults used by DocuElevate)
# ---------------------------------------------------------------------------
_CELERY_QUEUES = ("document_processor", "default", "celery")

# ---------------------------------------------------------------------------
# Singleton Redis client (lazy-initialised; fail-open when unavailable)
# ---------------------------------------------------------------------------
_redis_client: redis.Redis | None = None


def _get_redis() -> redis.Redis | None:
    """Return a shared Redis client, or *None* when Redis is unavailable."""
    global _redis_client
    if _redis_client is not None:
        return _redis_client
    try:
        _redis_client = redis.Redis.from_url(
            settings.redis_url,
            decode_responses=True,
            socket_connect_timeout=2,
            socket_timeout=2,
        )
        # Quick connectivity check – raises on failure.
        _redis_client.ping()
        return _redis_client
    except Exception:  # noqa: BLE001
        logger.debug("Redis unavailable for upload rate limiter – falling back to allow-all", exc_info=True)
        _redis_client = None
        return None


# ---------------------------------------------------------------------------
# Health metrics helpers
# ---------------------------------------------------------------------------


def _get_queue_depth(r: redis.Redis) -> int:
    """Return the total number of pending tasks across all Celery queues."""
    total = 0
    for queue_name in _CELERY_QUEUES:
        try:
            total += r.llen(queue_name)
        except Exception:  # noqa: BLE001, S110
            logger.debug("Could not read queue length for %r", queue_name, exc_info=True)
    return total


def _get_cpu_load_ratio() -> float:
    """Return the 1-minute load average divided by the number of CPU cores.

    Returns ``0.0`` on platforms that do not support :func:`os.getloadavg`
    (e.g. Windows) so that the limiter never penalises on those systems.
    """
    try:
        load_1m = os.getloadavg()[0]
        cpu_count = os.cpu_count() or 1
        return load_1m / cpu_count
    except (OSError, AttributeError):
        return 0.0


def compute_effective_limit(
    base_limit: int,
    queue_depth: int = 0,
    cpu_load_ratio: float = 0.0,
) -> tuple[int, float, str]:
    """Compute the effective upload rate limit based on system health.

    The function applies a *reduction factor* (``0.0 < factor ≤ 1.0``) to the
    configured base limit.  Both queue depth and CPU load contribute
    independently; the lowest factor wins.

    Args:
        base_limit: The configured maximum uploads per window.
        queue_depth: Total pending tasks in Celery queues.
        cpu_load_ratio: 1-minute load average divided by CPU count.

    Returns:
        A 3-tuple of ``(effective_limit, factor, reason)`` where *reason*
        is a human-readable tag for logging.
    """
    factor = 1.0
    reason = "normal"

    # --- Queue-depth thresholds ---
    if queue_depth > 200:
        factor, reason = min(factor, 0.10), f"critical_queue({queue_depth})"
    elif queue_depth > 100:
        factor, reason = min(factor, 0.25), f"high_queue({queue_depth})"
    elif queue_depth > 50:
        factor, reason = min(factor, 0.50), f"moderate_queue({queue_depth})"

    # --- CPU-load thresholds ---
    if cpu_load_ratio > 3.0:
        new_factor = 0.10
        if new_factor < factor:
            factor, reason = new_factor, f"critical_cpu({cpu_load_ratio:.1f})"
    elif cpu_load_ratio > 2.0:
        new_factor = 0.25
        if new_factor < factor:
            factor, reason = new_factor, f"high_cpu({cpu_load_ratio:.1f})"
    elif cpu_load_ratio > 1.5:
        new_factor = 0.50
        if new_factor < factor:
            factor, reason = new_factor, f"moderate_cpu({cpu_load_ratio:.1f})"

    effective = max(1, int(base_limit * factor))
    return effective, factor, reason


# ---------------------------------------------------------------------------
# Core sliding-window check (Redis sorted set)
# ---------------------------------------------------------------------------


def _check_and_record(
    r: redis.Redis,
    user_id: str,
    window: int,
    effective_limit: int,
) -> dict[str, Any] | None:
    """Atomically check the user's upload count and record the new upload.

    Uses a Redis sorted set where each member is a unique timestamp-based ID
    and the score is the Unix timestamp.  Entries older than *window* seconds
    are pruned on every call so the set never grows unbounded.

    Returns:
        ``None`` if the request is allowed, or a ``dict`` with ``count``,
        ``limit``, and ``retry_after`` if the limit is exceeded.
    """
    key = f"{_KEY_PREFIX}:{user_id}"
    now = time.time()
    window_start = now - window

    pipe = r.pipeline(transaction=True)
    # 1. Remove entries outside the window
    pipe.zremrangebyscore(key, "-inf", window_start)
    # 2. Count current entries
    pipe.zcard(key)
    # 3. Retrieve the oldest entry's score (to compute retry_after)
    pipe.zrange(key, 0, 0, withscores=True)
    results = pipe.execute()

    current_count: int = results[1]
    oldest_entries: list = results[2]

    if current_count >= effective_limit:
        # Compute how long until the oldest entry expires from the window.
        if oldest_entries:
            oldest_score = oldest_entries[0][1]
            retry_after = max(1, int((oldest_score + window) - now))
        else:
            retry_after = max(1, window // 2)
        return {
            "count": current_count,
            "limit": effective_limit,
            "retry_after": retry_after,
        }

    # 4. Record this upload (unique member = timestamp with random suffix)
    member = f"{now}:{os.urandom(4).hex()}"
    pipe2 = r.pipeline(transaction=True)
    pipe2.zadd(key, {member: now})
    pipe2.expire(key, window + 60)  # TTL slightly longer than window
    pipe2.execute()

    return None


# ---------------------------------------------------------------------------
# FastAPI dependency
# ---------------------------------------------------------------------------


async def require_upload_rate_limit(request: Request) -> None:
    """FastAPI dependency that enforces per-user upload rate limits.

    The dependency is designed to **fail open**: if Redis is unavailable the
    request is allowed through so that uploads are never blocked by a
    monitoring outage.

    Raises:
        HTTPException: 429 Too Many Requests when the per-user upload limit
            is exceeded.  The ``Retry-After`` header indicates how many
            seconds the client should wait before retrying.
    """
    r = _get_redis()
    if r is None:
        # Redis unavailable – fail open.
        return

    # Identify the user (owner_id for multi-user, IP fallback).
    user_id = get_current_owner_id(request)
    if not user_id:
        user_id = f"ip:{request.client.host}" if request.client else "ip:unknown"

    base_limit: int = settings.upload_rate_limit_per_user
    window: int = settings.upload_rate_limit_window

    # Gather health metrics and compute effective limit.
    try:
        queue_depth = _get_queue_depth(r)
    except Exception:  # noqa: BLE001
        queue_depth = 0

    cpu_load_ratio = _get_cpu_load_ratio()
    effective_limit, factor, health_reason = compute_effective_limit(base_limit, queue_depth, cpu_load_ratio)

    # Sliding-window check.
    try:
        rejection = _check_and_record(r, user_id, window, effective_limit)
    except Exception as exc:  # noqa: BLE001
        logger.warning("Upload rate-limit check failed (allowing request): %s", exc)
        return

    if rejection is not None:
        retry_after = rejection["retry_after"]
        logger.warning(
            "Upload rate limit exceeded: user=%s count=%d/%d window=%ds health=%s retry_after=%ds",
            user_id,
            rejection["count"],
            rejection["limit"],
            window,
            health_reason,
            retry_after,
        )
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail=(
                f"Upload rate limit exceeded ({rejection['count']}/{rejection['limit']} "
                f"in {window}s). Retry after {retry_after}s."
            ),
            headers={"Retry-After": str(retry_after)},
        )

    if factor < 1.0:
        logger.info(
            "Upload allowed with reduced limit: user=%s effective=%d/%d health=%s",
            user_id,
            effective_limit,
            base_limit,
            health_reason,
        )