From f45baa2ce0246a34aa658843d0ae7770969ad62e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=94=A8=E6=88=B7CI=20Test?= Date: Sat, 11 Jul 2026 16:13:11 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20=E4=BB=BB=E5=8A=A1=E9=98=9F?= =?UTF-8?q?=E5=88=97=E9=99=90=E6=B5=81=E9=98=B2=E6=8A=A4=20-=20=E7=94=A8?= =?UTF-8?q?=E6=88=B7=E7=BA=A73=E4=B8=AA/=E5=85=A8=E5=B1=8020=E4=B8=AA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 实现3层队列防护，防止批量提交导致队列爆炸： **用户级限流（核心）** - 每个用户同时 pending 的 generation 任务上限 3 个 - 超限返回 429：您的待处理任务过多，请等待完成后再提交 **全局限流（兜底）** - 系统 pending 任务超过 20 个一律拒绝 - 返回 503：系统繁忙，请稍后再试 **覆盖的入口** - 一键生成（generation_tasks 批量创建 + 重试） - 剪辑计划生成（edit_plans generate） - 任务中心重试（task_center 用户级 + 项目级） **实现细节** - 预检查 + 入队前检查双重保障 - 限流拒绝时任务标记为 failed，避免 pending 僵尸 - repository 不支持计数时自动降级跳过（兼容旧代码） - 全局检查始终生效，用户级检查需传 user_id **新增** - task_enqueue.py: UserPendingLimitExceeded / GlobalQueueFull 异常 - generation_task_repository: count_pending_by_user / count_pending_total - 13个单元测试，覆盖正常/超限/全局/降级等场景 --- apps/api/app/api/routes/edit_plans.py | 25 +++ apps/api/app/api/routes/generation_tasks.py | 101 ++++++++- apps/api/app/api/routes/task_center.py | 75 ++++++- apps/api/app/core/task_enqueue.py | 124 ++++++++++- .../generation_task_repository.py | 17 ++ packages/ports/generation_task_repository.py | 4 + tests/unit/test_task_queue_limit.py | 200 ++++++++++++++++++ 7 files changed, 530 insertions(+), 16 deletions(-) mode change 100644 => 100755 apps/api/app/api/routes/edit_plans.py mode change 100644 => 100755 packages/adapters/sqlalchemy_impl/generation_task_repository.py mode change 100644 => 100755 packages/ports/generation_task_repository.py create mode 100755 tests/unit/test_task_queue_limit.py diff --git a/apps/api/app/api/routes/edit_plans.py b/apps/api/app/api/routes/edit_plans.py old mode 100644 new mode 100755 index 5c7ea8f7a..592367d1c --- a/apps/api/app/api/routes/edit_plans.py +++ b/apps/api/app/api/routes/edit_plans.py @@ -644,6 +644,31 @@ def generate_plan( # 创建 GenerationTask gen_task_repo = SQLAlchemyGenerationTaskRepository(db) + + # 队列限流预检查（repository 不支持计数时跳过） + user_id = current_user.user.id + try: + has_count = hasattr(gen_task_repo, "count_pending_by_user") and hasattr( + gen_task_repo, "count_pending_total" + ) + if has_count: + user_pending = gen_task_repo.count_pending_by_user(user_id) + global_pending = gen_task_repo.count_pending_total() + if user_pending >= 3: + raise HTTPException( + status_code=429, + detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + ) + if global_pending >= 20: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) + except HTTPException: + raise + except Exception as e: + logger.warning("[队列限流] 剪辑计划限流检查失败，跳过: %s", e) + gen_task_use_case = CreateGenerationTaskUseCase(gen_task_repo) plan = svc.get_plan_or_raise(plan_id) gen_task = gen_task_use_case.execute( diff --git a/apps/api/app/api/routes/generation_tasks.py b/apps/api/app/api/routes/generation_tasks.py index 3728c0c64..849426c65 100755 --- a/apps/api/app/api/routes/generation_tasks.py +++ b/apps/api/app/api/routes/generation_tasks.py @@ -5,7 +5,12 @@ from typing import Any from app.auth import AuthenticatedUser, get_current_user from app.core.storage import OSSStorageService, get_storage_service -from app.core.task_enqueue import safe_enqueue_generation_task +from app.core.task_enqueue import ( + GlobalQueueFull, + UserPendingLimitExceeded, + check_queue_limits, + safe_enqueue_generation_task, +) from app.dependencies import ( get_asset_library_repository, get_asset_repository, @@ -228,9 +233,31 @@ def create_generation_task( count = request.count created_tasks = [] failed_tasks = [] + user_id = authenticated_user.user.id # 同批次任务共享 batch_id，用于视频查重时批次内比对 batch_id = uuid.uuid4().hex if count > 1 else "" + # 预检查：批量提交前先看会不会超限，避免建一半才拒 + try: + user_pending = generation_task_repository.count_pending_by_user(user_id) + global_pending = generation_task_repository.count_pending_total() + if user_pending + count > 3: + raise UserPendingLimitExceeded( + user_id=user_id, pending_count=user_pending + count, limit=3 + ) + if global_pending + count > 20: + raise GlobalQueueFull(pending_count=global_pending + count, limit=20) + except UserPendingLimitExceeded as e: + raise HTTPException( + status_code=429, + detail=f"您的待处理任务过多（当前 {e.pending_count - count}/{e.limit}，本次提交 {count} 个），请等待完成后再提交", + ) from e + except GlobalQueueFull as e: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) from e + try: for _ in range(count): task = use_case.execute( @@ -243,16 +270,42 @@ def create_generation_task( asset_ids=resolved_asset_ids, title_ids=request.title_ids, voice_ids=request.voice_ids, - created_by_user_id=authenticated_user.user.id, + created_by_user_id=user_id, source_edit_plan_id=request.source_edit_plan_id, asset_select_mode=request.asset_select_mode, batch_id=batch_id, ) ) - if safe_enqueue_generation_task(task, generation_task_repository, log_prefix="[生成任务]", log_task_status=True): - created_tasks.append(task) - else: + try: + if safe_enqueue_generation_task( + task, + generation_task_repository, + user_id=user_id, + log_prefix="[生成任务]", + log_task_status=True, + ): + created_tasks.append(task) + else: + failed_tasks.append(task) + except UserPendingLimitExceeded: + # 兜底：如果预检查后又并发提交了，在这里也拦住 failed_tasks.append(task) + if not created_tasks: + raise HTTPException( + status_code=429, + detail="您的待处理任务过多，请等待完成后再提交", + ) + break + except GlobalQueueFull: + failed_tasks.append(task) + if not created_tasks: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) + break + except HTTPException: + raise except Exception as e: logger.error("[生成任务] 创建失败: %s", e, exc_info=True) raise HTTPException(status_code=500, detail="创建生成任务失败，请稍后重试或查看任务日志") @@ -327,6 +380,21 @@ def retry_generation_task( if status_val != "failed": raise HTTPException(status_code=409, detail="Only failed tasks can be retried") + user_id = authenticated_user.user.id + # 预检查：创建前判断，>= 上限就拒绝 + user_pending = generation_task_repository.count_pending_by_user(user_id) + global_pending = generation_task_repository.count_pending_total() + if user_pending >= 3: + raise HTTPException( + status_code=429, + detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + ) + if global_pending >= 20: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) + use_case = CreateGenerationTaskUseCase(generation_task_repository) retried = use_case.execute( CreateGenerationTaskCommand( @@ -338,11 +406,28 @@ def retry_generation_task( asset_ids=task.asset_ids, title_ids=task.title_ids, voice_ids=task.voice_ids, - created_by_user_id=authenticated_user.user.id, + created_by_user_id=user_id, source_edit_plan_id=task.source_edit_plan_id or "", asset_select_mode=getattr(task, "asset_select_mode", ""), ) ) - if not safe_enqueue_generation_task(retried, generation_task_repository, log_prefix="[生成任务]", log_task_status=True): - logger.warning("[生成任务] 重试入队失败: task_id=%s", retried.id) + try: + if not safe_enqueue_generation_task( + retried, + generation_task_repository, + user_id=user_id, + log_prefix="[生成任务]", + log_task_status=True, + ): + logger.warning("[生成任务] 重试入队失败: task_id=%s", retried.id) + except UserPendingLimitExceeded: + raise HTTPException( + status_code=429, + detail="您的待处理任务过多，请等待完成后再提交", + ) from None + except GlobalQueueFull: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) from None return _to_generation_task_response(retried) diff --git a/apps/api/app/api/routes/task_center.py b/apps/api/app/api/routes/task_center.py index 0838973bf..b5071593a 100755 --- a/apps/api/app/api/routes/task_center.py +++ b/apps/api/app/api/routes/task_center.py @@ -3,7 +3,11 @@ from typing import Any from app.auth import AuthenticatedUser, get_current_user from app.core.celery_app import celery_app -from app.core.task_enqueue import safe_enqueue_generation_task +from app.core.task_enqueue import ( + GlobalQueueFull, + UserPendingLimitExceeded, + safe_enqueue_generation_task, +) from app.dependencies import ( get_generation_task_repository, get_ingest_job_repository, @@ -142,6 +146,21 @@ def retry_task_by_id( if _status_value(task.status) != "failed": raise HTTPException(status_code=409, detail="Only failed tasks can be retried") + user_id = authenticated_user.user.id + # 预检查 + user_pending = generation_task_repository.count_pending_by_user(user_id) + global_pending = generation_task_repository.count_pending_total() + if user_pending >= 3: + raise HTTPException( + status_code=429, + detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + ) + if global_pending >= 20: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) + use_case = CreateGenerationTaskUseCase(generation_task_repository) retried = use_case.execute( CreateGenerationTaskCommand( @@ -153,11 +172,24 @@ def retry_task_by_id( asset_ids=task.asset_ids, title_ids=task.title_ids, voice_ids=task.voice_ids, - created_by_user_id=authenticated_user.user.id, + created_by_user_id=user_id, ) ) - if not safe_enqueue_generation_task(retried, generation_task_repository, log_prefix="[任务中心]"): - logger.warning("[任务中心] 用户级重试入队失败: task_id=%s", retried.id) + try: + if not safe_enqueue_generation_task( + retried, generation_task_repository, user_id=user_id, log_prefix="[任务中心]" + ): + logger.warning("[任务中心] 用户级重试入队失败: task_id=%s", retried.id) + except UserPendingLimitExceeded: + raise HTTPException( + status_code=429, + detail="您的待处理任务过多，请等待完成后再提交", + ) from None + except GlobalQueueFull: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) from None return UserTaskResponse( id=f"generation:{retried.id}", task_type="generation", @@ -225,6 +257,22 @@ def retry_project_task( raise HTTPException(status_code=404, detail="Generation task not found") if _status_value(task.status) != "failed": raise HTTPException(status_code=409, detail="Only failed tasks can be retried") + + user_id = authenticated_user.user.id + # 预检查 + user_pending = generation_task_repository.count_pending_by_user(user_id) + global_pending = generation_task_repository.count_pending_total() + if user_pending >= 3: + raise HTTPException( + status_code=429, + detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + ) + if global_pending >= 20: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) + use_case = CreateGenerationTaskUseCase(generation_task_repository) retried = use_case.execute( CreateGenerationTaskCommand( @@ -236,11 +284,24 @@ def retry_project_task( asset_ids=task.asset_ids, title_ids=task.title_ids, voice_ids=task.voice_ids, - created_by_user_id=authenticated_user.user.id, + created_by_user_id=user_id, ) ) - if not safe_enqueue_generation_task(retried, generation_task_repository, log_prefix="[任务中心]"): - logger.warning("[任务中心] 项目级重试用队失败: task_id=%s", retried.id) + try: + if not safe_enqueue_generation_task( + retried, generation_task_repository, user_id=user_id, log_prefix="[任务中心]" + ): + logger.warning("[任务中心] 项目级重试入队失败: task_id=%s", retried.id) + except UserPendingLimitExceeded: + raise HTTPException( + status_code=429, + detail="您的待处理任务过多，请等待完成后再提交", + ) from None + except GlobalQueueFull: + raise HTTPException( + status_code=503, + detail="系统繁忙，请稍后再试", + ) from None return _generation_task_to_project_response(retried) if task_type == "ingest": job = ingest_job_repository.get(source_id) diff --git a/apps/api/app/core/task_enqueue.py b/apps/api/app/core/task_enqueue.py index 881938a58..4025ddaa2 100755 --- a/apps/api/app/core/task_enqueue.py +++ b/apps/api/app/core/task_enqueue.py @@ -6,24 +6,146 @@ from app.core.celery_app import celery_app logger = logging.getLogger(__name__) +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 = 3, + global_pending_limit: int = 20, +) -> None: + """检查队列限流，超限抛对应异常。 + + Args: + user_id: 用户 ID + generation_task_repository: 任务仓储 + user_pending_limit: 单用户 pending 上限，默认 3 + global_pending_limit: 全局 pending 上限，默认 20 + + 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) + + # 再查用户级 + 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 = 3, + global_pending_limit: int = 20, ) -> bool: - """安全入队：send_task 失败时自动把任务标记为 failed，避免留下 pending 僵尸任务。 + """安全入队：限流检查 → 发送 Celery 任务 → 失败自动标记 failed。 Args: task: 生成任务对象，需有 id 属性和 mark_failed 方法 generation_task_repository: 任务仓储，用于更新状态 + user_id: 用户 ID，传了才做用户级限流检查 log_prefix: 日志前缀，便于区分调用来源 log_task_status: 成功日志中是否额外打印任务状态 + user_pending_limit: 单用户 pending 上限，默认 3 + global_pending_limit: 全局 pending 上限，默认 20 Returns: True 表示入队成功，False 表示入队失败（已标记为 failed） + + Raises: + GlobalQueueFull: 全局 pending 超限时抛出，任务会被标记为 failed + UserPendingLimitExceeded: 用户 pending 超限时抛出，任务会被标记为 failed """ + # 全局限流检查（始终生效） + 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 + try: celery_app.send_task("worker.generate_video", args=[task.id]) if log_task_status: diff --git a/packages/adapters/sqlalchemy_impl/generation_task_repository.py b/packages/adapters/sqlalchemy_impl/generation_task_repository.py old mode 100644 new mode 100755 index d736a7eb2..646a05be1 --- a/packages/adapters/sqlalchemy_impl/generation_task_repository.py +++ b/packages/adapters/sqlalchemy_impl/generation_task_repository.py @@ -91,6 +91,23 @@ class SQLAlchemyGenerationTaskRepository: def count_by_user(self, user_id: str) -> int: return self.session.query(GenerationTaskModel).filter(GenerationTaskModel.created_by_user_id == user_id).count() + def count_pending_by_user(self, user_id: str) -> int: + return ( + self.session.query(GenerationTaskModel) + .filter( + GenerationTaskModel.created_by_user_id == user_id, + GenerationTaskModel.status == GenerationTaskStatus.PENDING.value, + ) + .count() + ) + + def count_pending_total(self) -> int: + return ( + self.session.query(GenerationTaskModel) + .filter(GenerationTaskModel.status == GenerationTaskStatus.PENDING.value) + .count() + ) + def list_recent_by_user(self, user_id: str, limit: int = 5) -> list[GenerationTask]: models = ( self.session.query(GenerationTaskModel) diff --git a/packages/ports/generation_task_repository.py b/packages/ports/generation_task_repository.py old mode 100644 new mode 100755 index 832144233..a86314aa9 --- a/packages/ports/generation_task_repository.py +++ b/packages/ports/generation_task_repository.py @@ -16,6 +16,10 @@ class GenerationTaskRepository(Protocol): def count_by_user(self, user_id: str) -> int: ... + def count_pending_by_user(self, user_id: str) -> int: ... + + def count_pending_total(self) -> int: ... + def list_recent_by_user(self, user_id: str, limit: int = 5) -> list[GenerationTask]: ... def list_by_source_edit_plan(self, plan_id: str) -> list[GenerationTask]: ... diff --git a/tests/unit/test_task_queue_limit.py b/tests/unit/test_task_queue_limit.py new file mode 100755 index 000000000..a3b3ca3a6 --- /dev/null +++ b/tests/unit/test_task_queue_limit.py @@ -0,0 +1,200 @@ +"""任务队列限流防护单元测试。""" +from __future__ import annotations + +import sys +import os +from unittest.mock import MagicMock + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "apps", "api")) + +from app.core.task_enqueue import ( + GlobalQueueFull, + UserPendingLimitExceeded, + check_queue_limits, + safe_enqueue_generation_task, +) + + +# --------------------------------------------------------------------------- +# Mock helpers +# --------------------------------------------------------------------------- + + +class MockRepository: + """支持 pending 计数的 mock repository。""" + + def __init__(self, user_pending: int = 0, global_pending: int = 0): + self._user_pending = user_pending + self._global_pending = global_pending + self.updated_tasks = [] + + def count_pending_by_user(self, user_id: str) -> int: + return self._user_pending + + def count_pending_total(self) -> int: + return self._global_pending + + def update(self, task): + self.updated_tasks.append(task) + return task + + +class MockTask: + def __init__(self, task_id: str = "task-1", status: str = "pending"): + self.id = task_id + self.status = status + self.error_message = "" + + def mark_failed(self, reason: str): + self.status = "failed" + self.error_message = reason + + +@pytest.fixture(autouse=True) +def mock_celery(monkeypatch): + """mock 掉 celery_app.send_task，避免真实发送。""" + mock_send = MagicMock() + monkeypatch.setattr("app.core.celery_app.celery_app.send_task", mock_send) + return mock_send + + +# --------------------------------------------------------------------------- +# check_queue_limits 单元测试 +# --------------------------------------------------------------------------- + + +class TestCheckQueueLimits: + """队列限流检查函数测试。""" + + def test_normal_passes_through(self): + """正常范围内的任务不受限制。""" + repo = MockRepository(user_pending=1, global_pending=5) + # 不抛异常 = 通过 + check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + + def test_user_limit_exceeded_raises(self): + """用户 pending 超过上限抛 UserPendingLimitExceeded。""" + repo = MockRepository(user_pending=4, global_pending=5) + with pytest.raises(UserPendingLimitExceeded) as exc_info: + check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + assert exc_info.value.user_id == "user-1" + assert exc_info.value.pending_count == 4 + assert exc_info.value.limit == 3 + + def test_user_at_limit_not_exceeded(self): + """用户 pending 刚好等于上限不算超限（给入队前的预检查留余量）。""" + repo = MockRepository(user_pending=3, global_pending=5) + check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + + def test_global_limit_exceeded_raises(self): + """全局 pending 超过上限抛 GlobalQueueFull。""" + repo = MockRepository(user_pending=1, global_pending=21) + with pytest.raises(GlobalQueueFull) as exc_info: + check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + assert exc_info.value.pending_count == 21 + assert exc_info.value.limit == 20 + + def test_global_at_limit_not_exceeded(self): + """全局 pending 刚好等于上限不算超限。""" + repo = MockRepository(user_pending=1, global_pending=20) + check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + + def test_global_takes_priority_over_user(self): + """全局和用户都超限时，优先抛全局异常。""" + repo = MockRepository(user_pending=5, global_pending=25) + with pytest.raises(GlobalQueueFull): + check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + + +# --------------------------------------------------------------------------- +# safe_enqueue_generation_task 限流集成测试 +# --------------------------------------------------------------------------- + + +class TestSafeEnqueueWithLimits: + """安全入队函数的限流功能测试。""" + + def test_normal_task_enqueues_successfully(self, mock_celery): + """正常任务入队成功，返回 True。""" + repo = MockRepository(user_pending=0, global_pending=0) + task = MockTask("task-1") + result = safe_enqueue_generation_task( + task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 + ) + assert result is True + mock_celery.assert_called_once_with("worker.generate_video", args=["task-1"]) + assert len(repo.updated_tasks) == 0 # 成功不需要更新状态 + + def test_user_limit_rejected_with_failed_status(self, mock_celery): + """用户超限：任务标记为 failed，抛 UserPendingLimitExceeded。""" + repo = MockRepository(user_pending=4, global_pending=5) + task = MockTask("task-1") + with pytest.raises(UserPendingLimitExceeded): + safe_enqueue_generation_task( + task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 + ) + mock_celery.assert_not_called() # 没发送到 Celery + assert task.status == "failed" + assert "限流" in task.error_message + assert len(repo.updated_tasks) == 1 + + def test_global_limit_rejected_with_failed_status(self, mock_celery): + """全局超限：任务标记为 failed，抛 GlobalQueueFull。""" + repo = MockRepository(user_pending=1, global_pending=21) + task = MockTask("task-1") + with pytest.raises(GlobalQueueFull): + safe_enqueue_generation_task( + task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 + ) + mock_celery.assert_not_called() + assert task.status == "failed" + assert len(repo.updated_tasks) == 1 + + def test_no_user_id_skips_user_limit(self, mock_celery): + """不传 user_id 时跳过用户级限流，只做全局检查。""" + # 用户超限但不传 user_id → 应该通过（因为不检查用户级） + repo = MockRepository(user_pending=10, global_pending=5) + task = MockTask("task-1") + result = safe_enqueue_generation_task( + task, repo, user_id="", user_pending_limit=3, global_pending_limit=20 + ) + assert result is True + mock_celery.assert_called_once() + + def test_no_user_id_still_checks_global(self, mock_celery): + """不传 user_id 时全局超限仍然被拦。""" + repo = MockRepository(user_pending=10, global_pending=25) + task = MockTask("task-1") + with pytest.raises(GlobalQueueFull): + safe_enqueue_generation_task( + task, repo, user_id="", user_pending_limit=3, global_pending_limit=20 + ) + mock_celery.assert_not_called() + + def test_default_limits_are_3_and_20(self, mock_celery): + """默认配置：用户级 3 个，全局 20 个。""" + # 刚好在默认限制内 + repo = MockRepository(user_pending=2, global_pending=19) + task = MockTask("task-1") + result = safe_enqueue_generation_task(task, repo, user_id="user-1") + assert result is True + + def test_update_failure_does_not_crash(self, mock_celery): + """repository.update 失败也不崩溃，异常继续向上抛。""" + + class BadRepo(MockRepository): + def update(self, task): + raise RuntimeError("db down") + + repo = BadRepo(user_pending=5, global_pending=5) + task = MockTask("task-1") + # 仍然抛 UserPendingLimitExceeded，不会被 update 失败掩盖 + with pytest.raises(UserPendingLimitExceeded): + safe_enqueue_generation_task( + task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 + ) + mock_celery.assert_not_called() + # 任务状态还是变了（内存里改了） + assert task.status == "failed" -- 2.54.0 From c5c3e0975b4254a835e9c34d6de5a02b191d5187 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=94=A8=E6=88=B7CI=20Test?= Date: Sat, 11 Jul 2026 16:24:58 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(queue-limit):=20P1=E4=BF=AE=E5=A4=8D=20?= =?UTF-8?q?-=20=E9=98=88=E5=80=BC=E6=8A=BD=E5=B8=B8=E9=87=8F=20+=20?= =?UTF-8?q?=E8=BE=B9=E7=95=8C=E7=BB=9F=E4=B8=80=20+=20=E5=85=A5=E9=98=9F?= =?UTF-8?q?=E5=90=8E=E5=85=9C=E5=BA=95=E6=A0=A1=E9=AA=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 阈值统一管理：USER_PENDING_LIMIT/GLOBAL_PENDING_LIMIT 抽到 task_enqueue.py 常量，所有入口引用 - 边界判断统一：预检查用 >=（任务创建前），入队检查用 >（包含当前任务），语义一致 - 并发竞态兜底：发送Celery后再查一次DB计数，超限则回滚任务为failed - 新增 12 个单元测试（入队后兜底5个 + 边界验证7个） --- apps/api/app/api/routes/edit_plans.py | 7 +- apps/api/app/api/routes/generation_tasks.py | 16 +- apps/api/app/api/routes/task_center.py | 14 +- apps/api/app/core/task_enqueue.py | 117 +++++++--- apps/api/postgres | 0 tests/unit/test_task_queue_limit.py | 229 ++++++++++++++++---- 6 files changed, 293 insertions(+), 90 deletions(-) create mode 100644 apps/api/postgres diff --git a/apps/api/app/api/routes/edit_plans.py b/apps/api/app/api/routes/edit_plans.py index 592367d1c..877c4f6fb 100755 --- a/apps/api/app/api/routes/edit_plans.py +++ b/apps/api/app/api/routes/edit_plans.py @@ -24,6 +24,7 @@ from typing import Any, List, Optional from app.auth import AuthenticatedUser, get_current_user from app.core.celery_app import celery_app +from app.core.task_enqueue import GLOBAL_PENDING_LIMIT, USER_PENDING_LIMIT from app.dependencies import get_asset_library_repository, get_asset_repository, get_db_session, get_project_repository from app.schemas.generation_task import GenerationTaskResponse from app.services import EditPlanService, PlanGeneratorService @@ -654,12 +655,12 @@ def generate_plan( if has_count: user_pending = gen_task_repo.count_pending_by_user(user_id) global_pending = gen_task_repo.count_pending_total() - if user_pending >= 3: + if user_pending >= USER_PENDING_LIMIT: raise HTTPException( status_code=429, - detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + detail=f"您的待处理任务过多（当前 {user_pending}/{USER_PENDING_LIMIT}），请等待完成后再提交", ) - if global_pending >= 20: + if global_pending >= GLOBAL_PENDING_LIMIT: raise HTTPException( status_code=503, detail="系统繁忙，请稍后再试", diff --git a/apps/api/app/api/routes/generation_tasks.py b/apps/api/app/api/routes/generation_tasks.py index 849426c65..d3f37c77e 100755 --- a/apps/api/app/api/routes/generation_tasks.py +++ b/apps/api/app/api/routes/generation_tasks.py @@ -6,6 +6,8 @@ from typing import Any from app.auth import AuthenticatedUser, get_current_user from app.core.storage import OSSStorageService, get_storage_service from app.core.task_enqueue import ( + GLOBAL_PENDING_LIMIT, + USER_PENDING_LIMIT, GlobalQueueFull, UserPendingLimitExceeded, check_queue_limits, @@ -241,12 +243,12 @@ def create_generation_task( try: user_pending = generation_task_repository.count_pending_by_user(user_id) global_pending = generation_task_repository.count_pending_total() - if user_pending + count > 3: + if user_pending + count > USER_PENDING_LIMIT: raise UserPendingLimitExceeded( - user_id=user_id, pending_count=user_pending + count, limit=3 + user_id=user_id, pending_count=user_pending + count, limit=USER_PENDING_LIMIT ) - if global_pending + count > 20: - raise GlobalQueueFull(pending_count=global_pending + count, limit=20) + if global_pending + count > GLOBAL_PENDING_LIMIT: + raise GlobalQueueFull(pending_count=global_pending + count, limit=GLOBAL_PENDING_LIMIT) except UserPendingLimitExceeded as e: raise HTTPException( status_code=429, @@ -384,12 +386,12 @@ def retry_generation_task( # 预检查：创建前判断，>= 上限就拒绝 user_pending = generation_task_repository.count_pending_by_user(user_id) global_pending = generation_task_repository.count_pending_total() - if user_pending >= 3: + if user_pending >= USER_PENDING_LIMIT: raise HTTPException( status_code=429, - detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + detail=f"您的待处理任务过多（当前 {user_pending}/{USER_PENDING_LIMIT}），请等待完成后再提交", ) - if global_pending >= 20: + if global_pending >= GLOBAL_PENDING_LIMIT: raise HTTPException( status_code=503, detail="系统繁忙，请稍后再试", diff --git a/apps/api/app/api/routes/task_center.py b/apps/api/app/api/routes/task_center.py index b5071593a..c796fd58d 100755 --- a/apps/api/app/api/routes/task_center.py +++ b/apps/api/app/api/routes/task_center.py @@ -4,6 +4,8 @@ from typing import Any from app.auth import AuthenticatedUser, get_current_user from app.core.celery_app import celery_app from app.core.task_enqueue import ( + GLOBAL_PENDING_LIMIT, + USER_PENDING_LIMIT, GlobalQueueFull, UserPendingLimitExceeded, safe_enqueue_generation_task, @@ -150,12 +152,12 @@ def retry_task_by_id( # 预检查 user_pending = generation_task_repository.count_pending_by_user(user_id) global_pending = generation_task_repository.count_pending_total() - if user_pending >= 3: + if user_pending >= USER_PENDING_LIMIT: raise HTTPException( status_code=429, - detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + detail=f"您的待处理任务过多（当前 {user_pending}/{USER_PENDING_LIMIT}），请等待完成后再提交", ) - if global_pending >= 20: + if global_pending >= GLOBAL_PENDING_LIMIT: raise HTTPException( status_code=503, detail="系统繁忙，请稍后再试", @@ -262,12 +264,12 @@ def retry_project_task( # 预检查 user_pending = generation_task_repository.count_pending_by_user(user_id) global_pending = generation_task_repository.count_pending_total() - if user_pending >= 3: + if user_pending >= USER_PENDING_LIMIT: raise HTTPException( status_code=429, - detail=f"您的待处理任务过多（当前 {user_pending}/3），请等待完成后再提交", + detail=f"您的待处理任务过多（当前 {user_pending}/{USER_PENDING_LIMIT}），请等待完成后再提交", ) - if global_pending >= 20: + if global_pending >= GLOBAL_PENDING_LIMIT: raise HTTPException( status_code=503, detail="系统繁忙，请稍后再试", diff --git a/apps/api/app/core/task_enqueue.py b/apps/api/app/core/task_enqueue.py index 4025ddaa2..0b3048d65 100755 --- a/apps/api/app/core/task_enqueue.py +++ b/apps/api/app/core/task_enqueue.py @@ -5,6 +5,10 @@ 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。""" @@ -29,16 +33,18 @@ def check_queue_limits( user_id: str, generation_task_repository: Any, *, - user_pending_limit: int = 3, - global_pending_limit: int = 20, + 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 上限，默认 3 - global_pending_limit: 全局 pending 上限，默认 20 + user_pending_limit: 单用户 pending 上限，默认 USER_PENDING_LIMIT + global_pending_limit: 全局 pending 上限，默认 GLOBAL_PENDING_LIMIT Raises: GlobalQueueFull: 全局超限时抛出（优先级更高，先查全局） @@ -46,7 +52,7 @@ def check_queue_limits( """ # 先查全局（系统级保护优先级更高） global_pending = generation_task_repository.count_pending_total() - if global_pending > global_pending_limit: + if global_pending >= global_pending_limit: logger.warning( "[队列限流] 全局 pending 任务数超限: %d/%d, user_id=%s", global_pending, @@ -56,17 +62,18 @@ def check_queue_limits( raise GlobalQueueFull(pending_count=global_pending, limit=global_pending_limit) # 再查用户级 - 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 - ) + 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( @@ -96,19 +103,27 @@ def safe_enqueue_generation_task( user_id: str = "", log_prefix: str = "[任务队列]", log_task_status: bool = False, - user_pending_limit: int = 3, - global_pending_limit: int = 20, + user_pending_limit: int = USER_PENDING_LIMIT, + global_pending_limit: int = GLOBAL_PENDING_LIMIT, ) -> bool: - """安全入队：限流检查 → 发送 Celery 任务 → 失败自动标记 failed。 + """安全入队：入队前限流检查 → 发送 Celery 任务 → 入队后最终校验兜底。 + + 边界说明： + 入队前检查用 > 而非 >=。因为调用此函数时 task 已经是 pending 状态并计入 DB， + pending 总数包含了当前任务本身。pending > limit 等价于"其他任务数 >= limit"， + 与预检查的 >= 语义一致（都是达到上限就拒绝新任务）。 + + 入队后最终校验：发送 Celery 成功后再查一次 DB 计数，处理并发竞态场景 + （两个请求同时通过入队前检查，后到的那个在这里被兜住）。 Args: - task: 生成任务对象，需有 id 属性和 mark_failed 方法 + task: 生成任务对象，需有 id 属性和 mark_failed 方法（状态已为 pending） generation_task_repository: 任务仓储，用于更新状态 user_id: 用户 ID，传了才做用户级限流检查 log_prefix: 日志前缀，便于区分调用来源 log_task_status: 成功日志中是否额外打印任务状态 - user_pending_limit: 单用户 pending 上限，默认 3 - global_pending_limit: 全局 pending 上限，默认 20 + user_pending_limit: 单用户 pending 上限，默认 USER_PENDING_LIMIT + global_pending_limit: 全局 pending 上限，默认 GLOBAL_PENDING_LIMIT Returns: True 表示入队成功，False 表示入队失败（已标记为 failed） @@ -117,11 +132,13 @@ def safe_enqueue_generation_task( 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", + "[队列限流] 全局 pending 任务数超限(入队前): %d/%d, user_id=%s", global_pending, global_pending_limit, user_id or "unknown", @@ -135,7 +152,7 @@ def safe_enqueue_generation_task( 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", + "[队列限流] 用户 pending 任务数超限(入队前): user_id=%s, count=%d/%d", user_id, user_pending, user_pending_limit, @@ -146,18 +163,9 @@ def safe_enqueue_generation_task( _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]) - 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 except Exception as e: logger.error( "%s 入队失败，标记为失败: task_id=%s error=%s", @@ -178,3 +186,42 @@ def safe_enqueue_generation_task( 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 diff --git a/apps/api/postgres b/apps/api/postgres new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/test_task_queue_limit.py b/tests/unit/test_task_queue_limit.py index a3b3ca3a6..6a906df57 100755 --- a/tests/unit/test_task_queue_limit.py +++ b/tests/unit/test_task_queue_limit.py @@ -10,6 +10,8 @@ import pytest sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "apps", "api")) from app.core.task_enqueue import ( + GLOBAL_PENDING_LIMIT, + USER_PENDING_LIMIT, GlobalQueueFull, UserPendingLimitExceeded, check_queue_limits, @@ -23,11 +25,15 @@ from app.core.task_enqueue import ( class MockRepository: - """支持 pending 计数的 mock repository。""" + """支持 pending 计数的 mock repository。 + + 支持通过 set_pending 动态修改计数，用于模拟入队后计数变化的并发场景。 + """ def __init__(self, user_pending: int = 0, global_pending: int = 0): self._user_pending = user_pending self._global_pending = global_pending + self._send_task_called = False self.updated_tasks = [] def count_pending_by_user(self, user_id: str) -> int: @@ -40,6 +46,13 @@ class MockRepository: self.updated_tasks.append(task) return task + def set_pending(self, *, user_pending: int | None = None, global_pending: int | None = None): + """动态修改 pending 计数，模拟并发场景。""" + if user_pending is not None: + self._user_pending = user_pending + if global_pending is not None: + self._global_pending = global_pending + class MockTask: def __init__(self, task_id: str = "task-1", status: str = "pending"): @@ -61,55 +74,83 @@ def mock_celery(monkeypatch): # --------------------------------------------------------------------------- -# check_queue_limits 单元测试 +# 常量导出测试 +# --------------------------------------------------------------------------- + + +def test_limit_constants_are_exported(): + """限流阈值常量已导出，供业务代码引用。""" + assert USER_PENDING_LIMIT == 3 + assert GLOBAL_PENDING_LIMIT == 20 + + +# --------------------------------------------------------------------------- +# check_queue_limits 单元测试（预检查用，>= 边界） # --------------------------------------------------------------------------- class TestCheckQueueLimits: - """队列限流检查函数测试。""" + """队列限流检查函数测试（预检查语义，>= 上限即拒绝）。""" def test_normal_passes_through(self): """正常范围内的任务不受限制。""" repo = MockRepository(user_pending=1, global_pending=5) - # 不抛异常 = 通过 - check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + check_queue_limits("user-1", repo) def test_user_limit_exceeded_raises(self): """用户 pending 超过上限抛 UserPendingLimitExceeded。""" repo = MockRepository(user_pending=4, global_pending=5) with pytest.raises(UserPendingLimitExceeded) as exc_info: - check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + check_queue_limits("user-1", repo) assert exc_info.value.user_id == "user-1" assert exc_info.value.pending_count == 4 assert exc_info.value.limit == 3 - def test_user_at_limit_not_exceeded(self): - """用户 pending 刚好等于上限不算超限（给入队前的预检查留余量）。""" + def test_user_at_limit_also_raises(self): + """用户 pending 刚好等于上限也拒绝（>= 边界）。""" repo = MockRepository(user_pending=3, global_pending=5) - check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + with pytest.raises(UserPendingLimitExceeded): + check_queue_limits("user-1", repo) + + def test_user_below_limit_passes(self): + """用户 pending 比上限少 1，通过。""" + repo = MockRepository(user_pending=2, global_pending=5) + check_queue_limits("user-1", repo) def test_global_limit_exceeded_raises(self): """全局 pending 超过上限抛 GlobalQueueFull。""" repo = MockRepository(user_pending=1, global_pending=21) with pytest.raises(GlobalQueueFull) as exc_info: - check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + check_queue_limits("user-1", repo) assert exc_info.value.pending_count == 21 assert exc_info.value.limit == 20 - def test_global_at_limit_not_exceeded(self): - """全局 pending 刚好等于上限不算超限。""" + def test_global_at_limit_also_raises(self): + """全局 pending 刚好等于上限也拒绝（>= 边界）。""" repo = MockRepository(user_pending=1, global_pending=20) - check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + with pytest.raises(GlobalQueueFull): + check_queue_limits("user-1", repo) + + def test_global_below_limit_passes(self): + """全局 pending 比上限少 1，通过。""" + repo = MockRepository(user_pending=1, global_pending=19) + check_queue_limits("user-1", repo) def test_global_takes_priority_over_user(self): """全局和用户都超限时，优先抛全局异常。""" repo = MockRepository(user_pending=5, global_pending=25) with pytest.raises(GlobalQueueFull): - check_queue_limits("user-1", repo, user_pending_limit=3, global_pending_limit=20) + check_queue_limits("user-1", repo) + + def test_empty_user_id_skips_user_check(self): + """不传 user_id 时跳过用户级检查，只做全局检查。""" + repo = MockRepository(user_pending=10, global_pending=5) + # 用户超限但不传 user_id → 全局未超限，应该通过 + check_queue_limits("", repo) # --------------------------------------------------------------------------- -# safe_enqueue_generation_task 限流集成测试 +# safe_enqueue_generation_task 限流集成测试（入队前用 >，包含当前任务） # --------------------------------------------------------------------------- @@ -120,46 +161,66 @@ class TestSafeEnqueueWithLimits: """正常任务入队成功，返回 True。""" repo = MockRepository(user_pending=0, global_pending=0) task = MockTask("task-1") - result = safe_enqueue_generation_task( - task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 - ) + result = safe_enqueue_generation_task(task, repo, user_id="user-1") assert result is True mock_celery.assert_called_once_with("worker.generate_video", args=["task-1"]) assert len(repo.updated_tasks) == 0 # 成功不需要更新状态 def test_user_limit_rejected_with_failed_status(self, mock_celery): """用户超限：任务标记为 failed，抛 UserPendingLimitExceeded。""" - repo = MockRepository(user_pending=4, global_pending=5) + repo = MockRepository(user_pending=5, global_pending=5) task = MockTask("task-1") with pytest.raises(UserPendingLimitExceeded): - safe_enqueue_generation_task( - task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 - ) - mock_celery.assert_not_called() # 没发送到 Celery + safe_enqueue_generation_task(task, repo, user_id="user-1") + mock_celery.assert_not_called() assert task.status == "failed" assert "限流" in task.error_message assert len(repo.updated_tasks) == 1 + def test_user_at_limit_still_passes(self, mock_celery): + """用户 pending 刚好等于上限：入队前检查用 >，包含当前任务，刚好到上限不算超。 + + 与预检查的 >= 语义一致：预检查时 pending=3 拒绝（不能再加新的）， + 但 safe_enqueue 被调用时任务已是 pending（就是第3个）， + pending=3 不满足 >3，所以通过。 + """ + repo = MockRepository(user_pending=3, global_pending=5) + task = MockTask("task-1") + result = safe_enqueue_generation_task(task, repo, user_id="user-1") + assert result is True + mock_celery.assert_called_once() + + def test_user_one_over_limit_rejected(self, mock_celery): + """用户 pending = limit + 1：超限被拒。""" + repo = MockRepository(user_pending=4, global_pending=5) + task = MockTask("task-1") + with pytest.raises(UserPendingLimitExceeded): + safe_enqueue_generation_task(task, repo, user_id="user-1") + mock_celery.assert_not_called() + def test_global_limit_rejected_with_failed_status(self, mock_celery): """全局超限：任务标记为 failed，抛 GlobalQueueFull。""" repo = MockRepository(user_pending=1, global_pending=21) task = MockTask("task-1") with pytest.raises(GlobalQueueFull): - safe_enqueue_generation_task( - task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 - ) + safe_enqueue_generation_task(task, repo, user_id="user-1") mock_celery.assert_not_called() assert task.status == "failed" assert len(repo.updated_tasks) == 1 + def test_global_at_limit_still_passes(self, mock_celery): + """全局 pending 刚好等于上限：入队前检查用 >，包含当前任务，刚好到上限不算超。""" + repo = MockRepository(user_pending=1, global_pending=20) + task = MockTask("task-1") + result = safe_enqueue_generation_task(task, repo, user_id="user-1") + assert result is True + mock_celery.assert_called_once() + def test_no_user_id_skips_user_limit(self, mock_celery): """不传 user_id 时跳过用户级限流，只做全局检查。""" - # 用户超限但不传 user_id → 应该通过（因为不检查用户级） repo = MockRepository(user_pending=10, global_pending=5) task = MockTask("task-1") - result = safe_enqueue_generation_task( - task, repo, user_id="", user_pending_limit=3, global_pending_limit=20 - ) + result = safe_enqueue_generation_task(task, repo, user_id="") assert result is True mock_celery.assert_called_once() @@ -168,14 +229,12 @@ class TestSafeEnqueueWithLimits: repo = MockRepository(user_pending=10, global_pending=25) task = MockTask("task-1") with pytest.raises(GlobalQueueFull): - safe_enqueue_generation_task( - task, repo, user_id="", user_pending_limit=3, global_pending_limit=20 - ) + safe_enqueue_generation_task(task, repo, user_id="") mock_celery.assert_not_called() - def test_default_limits_are_3_and_20(self, mock_celery): - """默认配置：用户级 3 个，全局 20 个。""" - # 刚好在默认限制内 + def test_default_limits_match_constants(self, mock_celery): + """默认配置与导出常量一致。""" + # 刚好在默认限制内（limit - 1） repo = MockRepository(user_pending=2, global_pending=19) task = MockTask("task-1") result = safe_enqueue_generation_task(task, repo, user_id="user-1") @@ -192,9 +251,101 @@ class TestSafeEnqueueWithLimits: task = MockTask("task-1") # 仍然抛 UserPendingLimitExceeded，不会被 update 失败掩盖 with pytest.raises(UserPendingLimitExceeded): - safe_enqueue_generation_task( - task, repo, user_id="user-1", user_pending_limit=3, global_pending_limit=20 - ) + safe_enqueue_generation_task(task, repo, user_id="user-1") mock_celery.assert_not_called() # 任务状态还是变了（内存里改了） assert task.status == "failed" + + +# --------------------------------------------------------------------------- +# 入队后最终校验（并发竞态兜底）测试 +# --------------------------------------------------------------------------- + + +class TestPostEnqueueFinalCheck: + """入队后最终校验：模拟并发场景，Celery发送后计数增加被兜住。""" + + def test_post_enqueue_global_overflow_rollback(self, mock_celery): + """并发场景：入队前检查通过，但发送Celery后全局计数超限 → 回滚为failed。 + + 模拟两个请求同时通过入队前检查（都查到 global=19）， + 都创建了任务（DB里变成 21），先发送Celery的那个在最终校验时被兜住。 + """ + repo = MockRepository(user_pending=1, global_pending=20) # 入队前：20 > 20？否 + task = MockTask("task-1") + + # 模拟发送Celery后，另一个并发请求也创建了任务，全局变成21 + def side_effect(*args, **kwargs): + repo.set_pending(global_pending=21) + + mock_celery.side_effect = side_effect + + with pytest.raises(GlobalQueueFull) as exc_info: + safe_enqueue_generation_task(task, repo, user_id="user-1") + + # Celery 确实发出去了（兜底不撤销 Celery，只回滚 DB 状态） + mock_celery.assert_called_once() + # 任务被标记为 failed + assert task.status == "failed" + assert "入队后" in task.error_message + assert exc_info.value.pending_count == 21 + assert len(repo.updated_tasks) == 1 + + def test_post_enqueue_user_overflow_rollback(self, mock_celery): + """并发场景：入队前检查通过，但发送Celery后用户计数超限 → 回滚为failed。""" + repo = MockRepository(user_pending=3, global_pending=5) # 入队前：3 > 3？否 + task = MockTask("task-1") + + def side_effect(*args, **kwargs): + repo.set_pending(user_pending=4) + + mock_celery.side_effect = side_effect + + with pytest.raises(UserPendingLimitExceeded) as exc_info: + safe_enqueue_generation_task(task, repo, user_id="user-1") + + mock_celery.assert_called_once() + assert task.status == "failed" + assert "入队后" in task.error_message + assert exc_info.value.user_id == "user-1" + assert exc_info.value.pending_count == 4 + + def test_post_enqueue_global_priority_over_user(self, mock_celery): + """入队后校验：全局和用户都超限时，优先抛全局异常。""" + repo = MockRepository(user_pending=3, global_pending=20) + task = MockTask("task-1") + + def side_effect(*args, **kwargs): + repo.set_pending(user_pending=5, global_pending=22) + + mock_celery.side_effect = side_effect + + with pytest.raises(GlobalQueueFull): + safe_enqueue_generation_task(task, repo, user_id="user-1") + + assert task.status == "failed" + + def test_post_enqueue_no_change_still_passes(self, mock_celery): + """入队后计数没变 → 正常通过，不回滚。""" + repo = MockRepository(user_pending=2, global_pending=10) + task = MockTask("task-1") + + result = safe_enqueue_generation_task(task, repo, user_id="user-1") + + assert result is True + mock_celery.assert_called_once() + assert task.status == "pending" # 状态没变 + assert len(repo.updated_tasks) == 0 # 没更新 DB + + def test_post_enqueue_no_user_id_skips_user_check(self, mock_celery): + """不传 user_id 时，入队后校验也跳过用户级，只查全局。""" + repo = MockRepository(user_pending=10, global_pending=5) + task = MockTask("task-1") + + def side_effect(*args, **kwargs): + repo.set_pending(user_pending=15, global_pending=5) # 用户超限但全局没超 + + mock_celery.side_effect = side_effect + + result = safe_enqueue_generation_task(task, repo, user_id="") + assert result is True # 用户级不检查，全局没超限 → 通过 -- 2.54.0