xiaoxia-saas/apps/api/app/core/task_enqueue.py

import logging
from typing import Any

from app.core.celery_app import celery_app

logger = logging.getLogger(__name__)

# ── 限流阈值常量（全系统统一管理，不要在业务代码里硬编码） ──
USER_PENDING_LIMIT = 3  # 单用户 pending 上限
GLOBAL_PENDING_LIMIT = 20  # 全局 pending 上限
WORKER_CONCURRENCY = 4  # worker 渲染并发数（infra/docker/compose.yml WORKER_CONCURRENCY 默认值）

# 限流错误码：前端据此区分"排队等待"与"创建失败"
ERROR_CODE_USER_QUEUE_FULL = "USER_QUEUE_FULL"  # 429：用户自己的任务排队中
ERROR_CODE_SYSTEM_QUEUE_FULL = "SYSTEM_QUEUE_FULL"  # 503：系统整体繁忙


class UserPendingLimitExceeded(Exception):
    """用户 pending 任务数超限，返回 429。"""

    def __init__(
        self,
        user_id: str,
        pending_count: int,
        limit: int,
        *,
        running_count: int = 0,
        requested_count: int = 1,
        queue_ahead: int = 0,
        estimated_wait_seconds: int = 0,
    ):
        self.user_id = user_id
        self.pending_count = pending_count
        self.limit = limit
        # 排队上下文（用于 429 结构化提示，前端展示"排队中"而非"创建失败"）
        self.running_count = running_count
        self.requested_count = requested_count
        self.queue_ahead = queue_ahead
        self.estimated_wait_seconds = estimated_wait_seconds
        super().__init__(f"用户 {user_id} pending 任务数 {pending_count} 超过上限 {limit}")


class GlobalQueueFull(Exception):
    """全局限流，返回 503。"""

    def __init__(
        self,
        pending_count: int,
        limit: int,
        *,
        running_count: int = 0,
        queue_ahead: int = 0,
        estimated_wait_seconds: int = 0,
    ):
        self.pending_count = pending_count
        self.limit = limit
        self.running_count = running_count
        self.queue_ahead = queue_ahead
        self.estimated_wait_seconds = estimated_wait_seconds
        super().__init__(f"系统 pending 任务数 {pending_count} 超过上限 {limit}")


def _estimate_wait_seconds(queue_ahead: int, generation_task_repository: Any) -> int:
    """根据排队任务数 + worker 并发数 + 历史平均任务耗时估算等待秒数。

    估算公式：ceil(排队任务数 / 并发数) × 平均单任务耗时。
    拿不到历史数据时仓储层返回默认 120 秒。
    """
    import math

    if queue_ahead <= 0:
        return 0
    try:
        estimator = getattr(generation_task_repository, "estimate_avg_duration_seconds", None)
        avg_seconds = estimator() if estimator is not None else 120.0
    except Exception:
        avg_seconds = 120.0
    return int(math.ceil(queue_ahead / WORKER_CONCURRENCY) * avg_seconds)


def build_rate_limit_detail(
    exc: Exception,
    generation_task_repository: Any,
    *,
    scope: str = "user",
) -> dict:
    """构造结构化限流响应体（HTTPException 的 detail）。

    前端按 detail.code 判断场景：
    - USER_QUEUE_FULL (429)：用户自己的任务在排队，应提示"等待/继续排队"，不是创建失败
    - SYSTEM_QUEUE_FULL (503)：系统繁忙，稍后重试

    detail 字段：
    - code: 错误码
    - message: 可读中文提示（可直接展示）
    - queued_count: 当前排队(pending)任务数
    - running_count: 当前渲染中(running)任务数
    - queue_ahead: 前方排队任务数（预计等待批次依据）
    - estimated_wait_seconds: 预计等待秒数
    - limit: 对应限流上限
    """
    if scope == "user" and isinstance(exc, UserPendingLimitExceeded):
        running = exc.running_count
        if not running:
            try:
                counter = getattr(generation_task_repository, "count_running_by_user", None)
                running = counter(exc.user_id) if counter is not None else 0
            except Exception:
                running = 0
        queue_ahead = exc.queue_ahead or max(exc.pending_count, 0)
        wait = exc.estimated_wait_seconds or _estimate_wait_seconds(queue_ahead, generation_task_repository)
        wait_minutes = max(1, round(wait / 60))
        message = (
            f"您有 {exc.pending_count} 个任务正在排队、{running} 个正在渲染，"
            f"同一时间最多提交 {exc.limit} 个任务。请等待约 {wait_minutes} 分钟后再提交"
        )
        return {
            "code": ERROR_CODE_USER_QUEUE_FULL,
            "message": message,
            "queued_count": exc.pending_count,
            "running_count": running,
            "queue_ahead": queue_ahead,
            "estimated_wait_seconds": wait,
            "limit": exc.limit,
        }

    # 全局繁忙
    pending = getattr(exc, "pending_count", 0)
    running = getattr(exc, "running_count", 0)
    if not running:
        try:
            counter = getattr(generation_task_repository, "count_running_total", None)
            running = counter() if counter is not None else 0
        except Exception:
            running = 0
    queue_ahead = getattr(exc, "queue_ahead", 0) or pending
    wait = getattr(exc, "estimated_wait_seconds", 0) or _estimate_wait_seconds(queue_ahead, generation_task_repository)
    wait_minutes = max(1, round(wait / 60))
    return {
        "code": ERROR_CODE_SYSTEM_QUEUE_FULL,
        "message": f"系统繁忙：当前 {pending} 个任务排队中、{running} 个渲染中，预计等待约 {wait_minutes} 分钟，请稍后再试",
        "queued_count": pending,
        "running_count": running,
        "queue_ahead": queue_ahead,
        "estimated_wait_seconds": wait,
        "limit": getattr(exc, "limit", GLOBAL_PENDING_LIMIT),
    }


def check_queue_limits(
    user_id: str,
    generation_task_repository: Any,
    *,
    user_pending_limit: int = USER_PENDING_LIMIT,
    global_pending_limit: int = GLOBAL_PENDING_LIMIT,
) -> None:
    """检查队列限流（预检查用，任务创建前调用），超限抛对应异常。

    边界语义：>= 上限即拒绝（达到上限就不能再加新任务）。

    Args:
        user_id: 用户 ID
        generation_task_repository: 任务仓储
        user_pending_limit: 单用户 pending 上限，默认 USER_PENDING_LIMIT
        global_pending_limit: 全局 pending 上限，默认 GLOBAL_PENDING_LIMIT

    Raises:
        GlobalQueueFull: 全局超限时抛出（优先级更高，先查全局）
        UserPendingLimitExceeded: 用户超限时抛出
    """
    # 先查全局（系统级保护优先级更高）
    global_pending = generation_task_repository.count_pending_total()
    if global_pending >= global_pending_limit:
        logger.warning(
            "[队列限流] 全局 pending 任务数超限: %d/%d, user_id=%s",
            global_pending,
            global_pending_limit,
            user_id,
        )
        raise GlobalQueueFull(pending_count=global_pending, limit=global_pending_limit)

    # 再查用户级
    if user_id:
        user_pending = generation_task_repository.count_pending_by_user(user_id)
        if user_pending >= user_pending_limit:
            logger.warning(
                "[队列限流] 用户 pending 任务数超限: user_id=%s, count=%d/%d",
                user_id,
                user_pending,
                user_pending_limit,
            )
            raise UserPendingLimitExceeded(user_id=user_id, pending_count=user_pending, limit=user_pending_limit)


def _mark_task_failed_safely(
    task: Any,
    generation_task_repository: Any,
    log_prefix: str,
    reason: str,
) -> None:
    """安全地把任务标记为 failed，更新失败只打日志不崩溃。"""
    try:
        task.mark_failed(f"任务被限流拒绝: {reason}")
        generation_task_repository.update(task)
    except Exception as update_err:
        logger.error(
            "%s 限流后更新状态也失败: task_id=%s error=%s",
            log_prefix,
            task.id,
            update_err,
            exc_info=True,
        )


def safe_enqueue_generation_task(
    task: Any,
    generation_task_repository: Any,
    *,
    user_id: str = "",
    log_prefix: str = "[任务队列]",
    log_task_status: bool = False,
    user_pending_limit: int = USER_PENDING_LIMIT,
    global_pending_limit: int = GLOBAL_PENDING_LIMIT,
) -> bool:
    """安全入队：入队前限流检查 → 发送 Celery 任务 → 入队后最终校验兜底。

    边界说明：
    入队前检查用 > 而非 >=。因为调用此函数时 task 已经是 pending 状态并计入 DB，
    pending 总数包含了当前任务本身。pending > limit 等价于"其他任务数 >= limit"，
    与预检查的 >= 语义一致（都是达到上限就拒绝新任务）。

    入队后最终校验：发送 Celery 成功后再查一次 DB 计数，处理并发竞态场景
    （两个请求同时通过入队前检查，后到的那个在这里被兜住）。

    Args:
        task: 生成任务对象，需有 id 属性和 mark_failed 方法（状态已为 pending）
        generation_task_repository: 任务仓储，用于更新状态
        user_id: 用户 ID，传了才做用户级限流检查
        log_prefix: 日志前缀，便于区分调用来源
        log_task_status: 成功日志中是否额外打印任务状态
        user_pending_limit: 单用户 pending 上限，默认 USER_PENDING_LIMIT
        global_pending_limit: 全局 pending 上限，默认 GLOBAL_PENDING_LIMIT

    Returns:
        True 表示入队成功，False 表示入队失败（已标记为 failed）

    Raises:
        GlobalQueueFull: 全局 pending 超限时抛出，任务会被标记为 failed
        UserPendingLimitExceeded: 用户 pending 超限时抛出，任务会被标记为 failed
    """
    # ── 入队前检查：任务已是 pending，用 > 判断（包含当前任务） ──

    # 全局限流检查（始终生效）
    global_pending = generation_task_repository.count_pending_total()
    if global_pending > global_pending_limit:
        logger.warning(
            "[队列限流] 全局 pending 任务数超限(入队前): %d/%d, user_id=%s",
            global_pending,
            global_pending_limit,
            user_id or "unknown",
        )
        exc: Exception = GlobalQueueFull(pending_count=global_pending, limit=global_pending_limit)
        _mark_task_failed_safely(task, generation_task_repository, log_prefix, str(exc))
        raise exc

    # 用户级限流检查（传了 user_id 才做）
    if user_id:
        user_pending = generation_task_repository.count_pending_by_user(user_id)
        if user_pending > user_pending_limit:
            logger.warning(
                "[队列限流] 用户 pending 任务数超限(入队前): user_id=%s, count=%d/%d",
                user_id,
                user_pending,
                user_pending_limit,
            )
            exc = UserPendingLimitExceeded(user_id=user_id, pending_count=user_pending, limit=user_pending_limit)
            _mark_task_failed_safely(task, generation_task_repository, log_prefix, str(exc))
            raise exc

    # ── 发送 Celery 任务 ──
    try:
        celery_app.send_task("worker.generate_video", args=[task.id])
    except Exception as e:
        logger.error(
            "%s 入队失败，标记为失败: task_id=%s error=%s",
            log_prefix,
            task.id,
            e,
            exc_info=True,
        )
        try:
            task.mark_failed(f"任务入队失败: {e}")
            generation_task_repository.update(task)
        except Exception as update_err:
            logger.error(
                "%s 入队失败后更新状态也失败: task_id=%s error=%s",
                log_prefix,
                task.id,
                update_err,
                exc_info=True,
            )
        return False

    # ── 入队后最终校验：并发竞态兜底 ──
    # 发送成功后再查一次，防止两个请求同时通过入队前检查导致超限
    global_after = generation_task_repository.count_pending_total()
    user_after = generation_task_repository.count_pending_by_user(user_id) if user_id else 0

    global_over = global_after > global_pending_limit
    user_over = bool(user_id and user_after > user_pending_limit)

    if global_over or user_over:
        if global_over:
            reason = f"全局 pending 超限(入队后): {global_after}/{global_pending_limit}"
            exc = GlobalQueueFull(pending_count=global_after, limit=global_pending_limit)
        else:
            reason = f"用户 pending 超限(入队后): {user_after}/{user_pending_limit}"
            exc = UserPendingLimitExceeded(user_id=user_id, pending_count=user_after, limit=user_pending_limit)

        logger.warning(
            "[队列限流] %s, task_id=%s, user_id=%s — 回滚状态为 failed",
            reason,
            task.id,
            user_id or "unknown",
        )
        _mark_task_failed_safely(task, generation_task_repository, log_prefix, reason)
        raise exc

    # 入队成功日志
    if log_task_status:
        logger.info(
            "%s 入队成功: task_id=%s, status=%s",
            log_prefix,
            task.id,
            task.status,
        )
    else:
        logger.info("%s 入队成功: task_id=%s", log_prefix, task.id)
    return True