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 上限


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

    def __init__(self, user_id: str, pending_count: int, limit: int):
        self.user_id = user_id
        self.pending_count = pending_count
        self.limit = limit
        super().__init__(f"用户 {user_id} pending 任务数 {pending_count} 超过上限 {limit}")


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

    def __init__(self, pending_count: int, limit: int):
        self.pending_count = pending_count
        self.limit = limit
        super().__init__(f"系统 pending 任务数 {pending_count} 超过上限 {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 = 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: Exception = 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