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: 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