diff --git a/apps/api/app/api/routes/generation_preview.py b/apps/api/app/api/routes/generation_preview.py index 64e5b16ca..e8d0b8a53 100755 --- a/apps/api/app/api/routes/generation_preview.py +++ b/apps/api/app/api/routes/generation_preview.py @@ -14,6 +14,7 @@ from app.core.task_enqueue import ( USER_PENDING_LIMIT, GlobalQueueFull, UserPendingLimitExceeded, + build_rate_limit_detail, safe_enqueue_generation_task, ) from app.dependencies import ( @@ -312,12 +313,12 @@ def create_preview_generation_task( except UserPendingLimitExceeded as e: raise HTTPException( status_code=429, - detail=f"您的待处理任务过多（当前 {e.pending_count - count}/{e.limit}，本次提交 {count} 个），请等待后再提交", + detail=build_rate_limit_detail(e, generation_task_repository, scope="user"), ) from e except GlobalQueueFull as e: raise HTTPException( status_code=503, - detail="系统繁忙，请稍后再试", + detail=build_rate_limit_detail(e, generation_task_repository, scope="global"), ) from e # 确定视频比例：优先前端传入，否则从模板 mode 推断 @@ -498,6 +499,7 @@ def create_preview_generation_task( # ── 入队 ── responses: list[PreviewGenerationTaskResponse] = [] + rate_limit_exc: Exception | None = None # 记录首个限流异常，全部失败时返回结构化提示 for variant_index, task in enumerate(created_tasks): try: enqueued = safe_enqueue_generation_task( @@ -510,23 +512,29 @@ def create_preview_generation_task( if not enqueued: logger.warning("[预览生成] 任务入队失败: task_id=%s", task.id) _mark_task_failed(generation_task_repository, task, "任务入队失败") - except UserPendingLimitExceeded: + except UserPendingLimitExceeded as e: _mark_task_failed(generation_task_repository, task, "待处理任务超限") - except GlobalQueueFull: + rate_limit_exc = rate_limit_exc or e + except GlobalQueueFull as e: _mark_task_failed(generation_task_repository, task, "系统队列已满") + rate_limit_exc = rate_limit_exc or e except Exception: logger.exception("[预览生成] 入队异常: task_id=%s", task.id) _mark_task_failed(generation_task_repository, task, "任务入队异常") # enqueue 会原地更新 task 状态/进度，直接用 task 构造响应 responses.append(_to_preview_response(task)) - # 队列满/限流时若全部失败，返回明确错误码 - if all(r.status == "failed" for r in responses): - first_err = next((r.error_message for r in responses if r.error_message), "") - if "待处理任务" in first_err: - raise HTTPException(status_code=429, detail=first_err or "待处理任务超限") - if "队列" in first_err: - raise HTTPException(status_code=503, detail=first_err or "系统繁忙，请稍后再试") + # 队列满/限流时若全部失败，返回结构化错误码（前端区分"排队"与"创建失败"） + if all(r.status == "failed" for r in responses) and rate_limit_exc is not None: + if isinstance(rate_limit_exc, UserPendingLimitExceeded): + raise HTTPException( + status_code=429, + detail=build_rate_limit_detail(rate_limit_exc, generation_task_repository, scope="user"), + ) + raise HTTPException( + status_code=503, + detail=build_rate_limit_detail(rate_limit_exc, generation_task_repository, scope="global"), + ) logger.info( "[预览生成] 创建完成: %d 个变体任务, task_ids=%s", diff --git a/apps/api/app/api/routes/generation_tasks.py b/apps/api/app/api/routes/generation_tasks.py index 32391a34d..b0e01892f 100755 --- a/apps/api/app/api/routes/generation_tasks.py +++ b/apps/api/app/api/routes/generation_tasks.py @@ -10,6 +10,7 @@ from app.core.task_enqueue import ( USER_PENDING_LIMIT, GlobalQueueFull, UserPendingLimitExceeded, + build_rate_limit_detail, safe_enqueue_generation_task, ) from app.dependencies import ( @@ -417,12 +418,12 @@ def create_generation_task( except UserPendingLimitExceeded as e: raise HTTPException( status_code=429, - detail=f"您的待处理任务过多（当前 {e.pending_count - count}/{e.limit}，本次提交 {count} 个），请等待完成后再提交", + detail=build_rate_limit_detail(e, generation_task_repository, scope="user"), ) from e except GlobalQueueFull as e: raise HTTPException( status_code=503, - detail="系统繁忙，请稍后再试", + detail=build_rate_limit_detail(e, generation_task_repository, scope="global"), ) from e # 画中画已下线：strategy_id 中的 pip/voice_pip 统一映射为 one_take @@ -579,7 +580,7 @@ def create_generation_task( if not created_tasks: raise HTTPException( status_code=429, - detail="您的待处理任务过多，请等待完成后再提交", + detail=build_rate_limit_detail(_e, generation_task_repository, scope="user"), ) from _e break except GlobalQueueFull as _e: @@ -587,7 +588,7 @@ def create_generation_task( if not created_tasks: raise HTTPException( status_code=503, - detail="系统繁忙，请稍后再试", + detail=build_rate_limit_detail(_e, generation_task_repository, scope="global"), ) from _e break except HTTPException: @@ -713,15 +714,15 @@ def confirm_generation( log_task_status=True, ): logger.warning("[确认生成] 入队失败: task_id=%s", new_task.id) - except UserPendingLimitExceeded: + except UserPendingLimitExceeded as _e: raise HTTPException( status_code=429, - detail="您的待处理任务过多，请等待完成后再提交", + detail=build_rate_limit_detail(_e, generation_task_repository, scope="user"), ) from None - except GlobalQueueFull: + except GlobalQueueFull as _e: raise HTTPException( status_code=503, - detail="系统繁忙，请稍后再试", + detail=build_rate_limit_detail(_e, generation_task_repository, scope="global"), ) from None return BatchGenerationTaskResponse( @@ -803,12 +804,24 @@ def retry_generation_task( if user_pending >= USER_PENDING_LIMIT: raise HTTPException( status_code=429, - detail=f"您的待处理任务过多（当前 {user_pending}/{USER_PENDING_LIMIT}），请等待完成后再提交", + detail=build_rate_limit_detail( + UserPendingLimitExceeded( + user_id=user_id, + pending_count=user_pending, + limit=USER_PENDING_LIMIT, + ), + generation_task_repository, + scope="user", + ), ) if global_pending >= GLOBAL_PENDING_LIMIT: raise HTTPException( status_code=503, - detail="系统繁忙，请稍后再试", + detail=build_rate_limit_detail( + GlobalQueueFull(pending_count=global_pending, limit=GLOBAL_PENDING_LIMIT), + generation_task_repository, + scope="global", + ), ) use_case = CreateGenerationTaskUseCase(generation_task_repository) @@ -843,15 +856,15 @@ def retry_generation_task( log_task_status=True, ): logger.warning("[生成任务] 重试入队失败: task_id=%s", retried.id) - except UserPendingLimitExceeded: + except UserPendingLimitExceeded as _e: raise HTTPException( status_code=429, - detail="您的待处理任务过多，请等待完成后再提交", + detail=build_rate_limit_detail(_e, generation_task_repository, scope="user"), ) from None - except GlobalQueueFull: + except GlobalQueueFull as _e: raise HTTPException( status_code=503, - detail="系统繁忙，请稍后再试", + detail=build_rate_limit_detail(_e, generation_task_repository, scope="global"), ) from None return _to_generation_task_response(retried) diff --git a/apps/api/app/core/task_enqueue.py b/apps/api/app/core/task_enqueue.py index 31545f177..95e0d48cc 100755 --- a/apps/api/app/core/task_enqueue.py +++ b/apps/api/app/core/task_enqueue.py @@ -8,27 +8,145 @@ 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): + 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): + 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, diff --git a/apps/worker/worker_app/celery_app.py b/apps/worker/worker_app/celery_app.py index d953c23f4..a5051e5a4 100755 --- a/apps/worker/worker_app/celery_app.py +++ b/apps/worker/worker_app/celery_app.py @@ -22,10 +22,18 @@ celery_app.conf.imports = ( ) # Celery Beat 定时任务调度 +# 注：worker 单实例内嵌 beat（entrypoint-worker.sh -B），定时任务不会重复执行 celery_app.conf.beat_schedule = { + # pending 任务超时清理：worker 停止消费后，卡 pending 的任务 15 分钟内释放限流名额 "cleanup-stale-pending-tasks": { "task": "worker.cleanup_stale_pending_tasks", - "schedule": 600.0, # 每 10 分钟（秒） - "options": {"expires": 300}, # 5 分钟过期，避免堆积 + "schedule": 300.0, # 每 5 分钟（秒） + "options": {"expires": 240}, # 4 分钟过期，避免堆积 + }, + # running 孤儿任务巡检：容器重启/进程被杀后卡 running 的任务，20 分钟无更新则判失败 + "cleanup-stale-running-tasks": { + "task": "worker.cleanup_stale_running_tasks", + "schedule": 300.0, # 每 5 分钟（秒） + "options": {"expires": 240}, }, } diff --git a/apps/worker/worker_app/tasks/_startup.py b/apps/worker/worker_app/tasks/_startup.py index c841f4406..c3e6e8dfb 100644 --- a/apps/worker/worker_app/tasks/_startup.py +++ b/apps/worker/worker_app/tasks/_startup.py @@ -7,11 +7,34 @@ from worker_app.db import SessionLocal logger = logging.getLogger(__name__) -# 孤儿任务超时阈值：渲染任务超过此时间未更新则视为卡死 -ORPHAN_TASK_TIMEOUT_MINUTES = 10 -# Pending 任务超时阈值：pending 任务在队列中等待超过此时间则自动清理 -PENDING_TASK_TIMEOUT_MINUTES = 30 +def cleanup_stale_running_with_session(repo, timeout_minutes: int) -> int: + """清理超时未更新的 running GenerationTask（可注入 repo 的纯核心，便于单测）。 + + Returns: + 清理的任务数量 + """ + return repo.cleanup_stale_running(timeout_minutes) + + +def cleanup_stale_pending_with_session(repo, timeout_minutes: int) -> int: + """清理超时 pending GenerationTask（可注入 repo 的纯核心，便于单测）。 + + Returns: + 清理的任务数量 + """ + return repo.cleanup_stale_pending(timeout_minutes) + + +# 孤儿任务超时阈值：running 任务超过此时间无进度更新则视为卡死。 +# 依据：worker.generate_video 硬超时 time_limit=11 分钟，正常任务不可能超过； +# 20 分钟阈值覆盖硬超时 + 重试 + 余量，绝不误杀正常任务。 +ORPHAN_TASK_TIMEOUT_MINUTES = 20 + +# Pending 任务超时阈值：任务创建后超过此时间仍未被 worker 拉取， +# 说明 worker 已停止消费（容器异常/卡死），清掉释放限流名额。 +# 依据：满队列（20 pending）× 平均 2 分钟 / 并发 4 ≈ 10 分钟，15 分钟留余量。 +PENDING_TASK_TIMEOUT_MINUTES = 15 def cleanup_orphan_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> int: # pragma: no cover @@ -32,9 +55,11 @@ def cleanup_orphan_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> try: session = SessionLocal() - repo = SQLAlchemyGenerationTaskRepository(session) - count = repo.cleanup_stale_running(timeout_minutes) - session.close() + try: + repo = SQLAlchemyGenerationTaskRepository(session) + count = cleanup_stale_running_with_session(repo, timeout_minutes) + finally: + session.close() if count > 0: logger.warning("清理了 %d 个超时的孤儿 GenerationTask（超过 %d 分钟未更新）", count, timeout_minutes) else: @@ -105,7 +130,7 @@ def cleanup_stale_pending_tasks(timeout_minutes: int = PENDING_TASK_TIMEOUT_MINU session = SessionLocal() try: repo = SQLAlchemyGenerationTaskRepository(session) - count = repo.cleanup_stale_pending(timeout_minutes) + count = cleanup_stale_pending_with_session(repo, timeout_minutes) if count > 0: logger.warning("清理了 %d 个超时的 pending GenerationTask（超过 %d 分钟未处理）", count, timeout_minutes) else: diff --git a/apps/worker/worker_app/tasks/cleanup.py b/apps/worker/worker_app/tasks/cleanup.py index 28d5e7e22..5948e8a0d 100644 --- a/apps/worker/worker_app/tasks/cleanup.py +++ b/apps/worker/worker_app/tasks/cleanup.py @@ -1,14 +1,18 @@ """定期清理任务 — Celery Beat 调度。 包含： -- cleanup_stale_pending_tasks: 定期清理卡在 pending 超时的 generation_tasks +- cleanup_stale_pending_tasks: 定期清理卡在 pending 超时的 generation_tasks（worker 停止消费时占位） +- cleanup_stale_running_tasks: 定期清理卡在 running 超时的 generation_tasks（容器重启/进程被杀后的孤儿） """ import logging from celery import shared_task from worker_app.tasks._startup import ( + ORPHAN_TASK_TIMEOUT_MINUTES, PENDING_TASK_TIMEOUT_MINUTES, + cleanup_orphan_tasks, + cleanup_stale_jobs, cleanup_stale_pending_tasks, ) @@ -19,12 +23,12 @@ logger = logging.getLogger(__name__) def scheduled_cleanup_stale_pending(timeout_minutes: int = PENDING_TASK_TIMEOUT_MINUTES) -> dict: """Celery Beat 调度的定期任务：清理超时的 pending 任务。 - 每 10 分钟执行一次（由 celery_app.py 的 beat_schedule 配置）， + 每 5 分钟执行一次（由 celery_app.py 的 beat_schedule 配置）， 查找所有 status='pending' 且 created_at < NOW() - timeout_minutes - 的 generation_tasks，批量更新为 failed。 + 的 generation_tasks，批量更新为 failed，释放限流名额。 Args: - timeout_minutes: 超时时间（分钟），默认 30 分钟 + timeout_minutes: 超时时间（分钟），默认 15 分钟 Returns: {"cleaned": int} @@ -33,3 +37,33 @@ def scheduled_cleanup_stale_pending(timeout_minutes: int = PENDING_TASK_TIMEOUT_ if count > 0: logger.info("[Beat] 清理了 %d 个超时 pending 任务（超时阈值 %d 分钟）", count, timeout_minutes) return {"cleaned": count} + + +@shared_task(name="worker.cleanup_stale_running_tasks") +def scheduled_cleanup_stale_running(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> dict: + """Celery Beat 调度的定期任务：清理超时的 running 孤儿任务。 + + 每 5 分钟执行一次。worker_ready 信号只在 worker 启动时清一次， + 若 worker 没重启但任务卡死（上传挂起、进程 OOM 被内核杀掉等）， + 任务会永久卡在 running 占位。此任务做持续兜底： + 查找 status='running' 且 updated_at < NOW() - timeout_minutes 的任务， + 标记为 failed（原因：容器重启/超时中断），同时清理 Job 表孤儿。 + + Args: + timeout_minutes: 超时时间（分钟），默认 20 分钟 + （worker.generate_video 硬超时 11 分钟，正常任务不可能超过 20 分钟） + + Returns: + {"generation_tasks": int, "jobs": int} + """ + gen_count = cleanup_orphan_tasks(timeout_minutes) + job_count = cleanup_stale_jobs(timeout_minutes) + total = gen_count + job_count + if total > 0: + logger.warning( + "[Beat] 清理孤儿任务: running GenerationTask=%d, Job=%d（超时阈值 %d 分钟）", + gen_count, + job_count, + timeout_minutes, + ) + return {"generation_tasks": gen_count, "jobs": job_count} diff --git a/packages/adapters/sqlalchemy_impl/generation_task_repository.py b/packages/adapters/sqlalchemy_impl/generation_task_repository.py index 38aa44f9d..9ef851ad5 100755 --- a/packages/adapters/sqlalchemy_impl/generation_task_repository.py +++ b/packages/adapters/sqlalchemy_impl/generation_task_repository.py @@ -138,6 +138,52 @@ class SQLAlchemyGenerationTaskRepository: .count() ) + def count_running_by_user(self, user_id: str) -> int: + """统计指定用户处于 running 状态的任务数（用于限流提示展示）。""" + return ( + self.session.query(GenerationTaskModel) + .filter( + GenerationTaskModel.created_by_user_id == user_id, + GenerationTaskModel.status == GenerationTaskStatus.RUNNING.value, + ) + .count() + ) + + def count_running_total(self) -> int: + """统计全局处于 running 状态的任务数（worker 实际在执行的任务数）。""" + return ( + self.session.query(GenerationTaskModel) + .filter(GenerationTaskModel.status == GenerationTaskStatus.RUNNING.value) + .count() + ) + + def estimate_avg_duration_seconds(self, limit: int = 20, default_seconds: float = 120.0) -> float: + """估算最近完成任务的平均耗时（秒），用于 429 限流提示的等待预估。 + + 取最近 N 条 completed 任务的 (completed_at - started_at) 平均值； + 无足够历史数据时返回 default_seconds。 + 用 Python 侧计算差值，避免 SQLite/PostgreSQL 方言差异。 + """ + rows = ( + self.session.query(GenerationTaskModel.started_at, GenerationTaskModel.completed_at) + .filter( + GenerationTaskModel.status == GenerationTaskStatus.COMPLETED.value, + GenerationTaskModel.started_at.isnot(None), + GenerationTaskModel.completed_at.isnot(None), + ) + .order_by(GenerationTaskModel.completed_at.desc()) + .limit(limit) + .all() + ) + durations = [ + (completed - started).total_seconds() + for started, completed in rows + if completed and started and (completed - started).total_seconds() > 0 + ] + if not durations: + return default_seconds + return sum(durations) / len(durations) + 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 index 5c21200e5..cc79f0370 100755 --- a/packages/ports/generation_task_repository.py +++ b/packages/ports/generation_task_repository.py @@ -20,6 +20,12 @@ class GenerationTaskRepository(Protocol): def count_pending_total(self) -> int: ... + def count_running_by_user(self, user_id: str) -> int: ... + + def count_running_total(self) -> int: ... + + def estimate_avg_duration_seconds(self, limit: int = 20, default_seconds: float = 120.0) -> float: ... + 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_fault_tolerance_1709.py b/tests/unit/test_task_fault_tolerance_1709.py new file mode 100644 index 000000000..756543d2a --- /dev/null +++ b/tests/unit/test_task_fault_tolerance_1709.py @@ -0,0 +1,315 @@ +"""Issue #1709 任务容错：孤儿任务恢复 + 429 限流结构化提示。 + +覆盖： +1. 仓储层：count_running_by_user/count_running_total 计数正确（预览/正式任务都计入） +2. 仓储层：estimate_avg_duration_seconds 耗时估算（有历史/无历史） +3. 限流核心：build_rate_limit_detail 返回结构化 code/message/排队数/预计等待 +4. worker 侧：cleanup_stale_running/pending 核心函数——中断任务被重置为 failed + 且原因写明（容器重启/超时中断），正常任务不受影响 +""" + +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path +from unittest.mock import MagicMock + +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "api")) +sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "apps" / "worker")) + +# 预注入 mock worker_app.db，防止真实数据库连接初始化（与其他 worker 测试同模式） +_mock_db = MagicMock() +_mock_db.SessionLocal = MagicMock() +sys.modules.setdefault("worker_app.db", _mock_db) + +from app.core import task_enqueue # noqa: E402 +from sqlalchemy import create_engine, text # noqa: E402 +from sqlalchemy.orm import sessionmaker # noqa: E402 +from worker_app.tasks import _startup # noqa: E402 + +from packages.adapters.sqlalchemy_impl.generation_task_repository import ( # noqa: E402 + SQLAlchemyGenerationTaskRepository, +) +from packages.adapters.sqlalchemy_impl.models import Base # noqa: E402 +from packages.domain import GenerationTask, GenerationTaskStatus # noqa: E402 + + +def _repository(): + engine = create_engine("sqlite:///:memory:", connect_args={"check_same_thread": False}) + Base.metadata.create_all(engine) + session = sessionmaker(bind=engine)() + return SQLAlchemyGenerationTaskRepository(session), session, engine + + +def _make_task(**kwargs) -> GenerationTask: + defaults = dict( + project_id="proj-1", + asset_library_id="lib-1", + created_by_user_id="user-1", + ) + defaults.update(kwargs) + return GenerationTask.create(**defaults) + + +def _age_task(engine, task_id, *, updated_minutes=None, created_minutes=None): + """用 SQL 直接把 updated_at/created_at 改到过去（模拟孤儿任务）。""" + sets, params = [], {"id": task_id} + if updated_minutes is not None: + sets.append("updated_at = :uts") + params["uts"] = datetime.now(timezone.utc) - timedelta(minutes=updated_minutes) + if created_minutes is not None: + sets.append("created_at = :cts") + params["cts"] = datetime.now(timezone.utc) - timedelta(minutes=created_minutes) + with engine.connect() as conn: + conn.execute(text(f"UPDATE generation_tasks SET {', '.join(sets)} WHERE id = :id"), params) + conn.commit() + + +# --------------------------------------------------------------------------- +# 1. running 计数（限流"渲染中"数量） +# --------------------------------------------------------------------------- + + +def test_count_running_by_user_mix_statuses(): + """count_running_by_user 只统计该用户 running，不含 pending/completed/failed。""" + repo, _, _ = _repository() + t1 = _make_task(project_id="p1") + repo.create(t1) # pending + t2 = _make_task(project_id="p2") + repo.create(t2) + t2.mark_processing() + repo.update(t2) + t3 = _make_task(project_id="p3") + repo.create(t3) + t3.mark_processing() + repo.update(t3) + t4 = _make_task(project_id="p4") + repo.create(t4) + t4.mark_processing() + repo.update(t4) + t4.mark_completed() + repo.update(t4) + t5 = _make_task(project_id="p5", created_by_user_id="user-2") + repo.create(t5) + t5.mark_processing() + repo.update(t5) + + assert repo.count_running_by_user("user-1") == 2 + assert repo.count_running_by_user("user-2") == 1 + assert repo.count_running_total() == 3 + + +def test_count_running_total_empty(): + repo, _, _ = _repository() + assert repo.count_running_total() == 0 + assert repo.count_running_by_user("nobody") == 0 + + +def test_preview_tasks_counted_in_running(): + """预览任务（is_preview=True，工单实测卡 80% 的那种）同样计入 running。""" + repo, _, _ = _repository() + t = _make_task(is_preview=True) + repo.create(t) + t.mark_processing() + repo.update(t) + assert repo.count_running_by_user("user-1") == 1 + assert repo.count_running_total() == 1 + + +# --------------------------------------------------------------------------- +# 2. 平均耗时估算（429 等待预估依据） +# --------------------------------------------------------------------------- + + +def _complete_task(repo, engine, task, duration_seconds: float): + repo.create(task) + task.mark_processing() + repo.update(task) + task.mark_completed() + repo.update(task) + now = datetime.now(timezone.utc) + with engine.connect() as conn: + conn.execute( + text("UPDATE generation_tasks SET started_at = :s, completed_at = :c WHERE id = :id"), + {"s": now - timedelta(seconds=duration_seconds), "c": now, "id": task.id}, + ) + conn.commit() + + +def test_estimate_avg_duration_with_history(): + """有历史完成任务时返回平均耗时（秒）。""" + repo, _, engine = _repository() + _complete_task(repo, engine, _make_task(project_id="p1"), 60.0) + _complete_task(repo, engine, _make_task(project_id="p2"), 180.0) + + avg = repo.estimate_avg_duration_seconds(default_seconds=120.0) + assert 119.0 < avg < 121.0 # (60+180)/2 = 120 + + +def test_estimate_avg_duration_no_history_returns_default(): + """无历史数据时返回默认值。""" + repo, _, _ = _repository() + assert repo.estimate_avg_duration_seconds(default_seconds=90.0) == 90.0 + + +# --------------------------------------------------------------------------- +# 3. build_rate_limit_detail 结构化提示（前端区分"排队"与"创建失败"） +# --------------------------------------------------------------------------- + + +def test_user_rate_limit_detail_structure(): + """429 用户限流：返回 USER_QUEUE_FULL + 排队/渲染数 + 预计等待。""" + repo, _, _ = _repository() + for i in range(2): # 2 个渲染中 + t = _make_task(project_id=f"rp{i}") + repo.create(t) + t.mark_processing() + repo.update(t) + + exc = task_enqueue.UserPendingLimitExceeded(user_id="user-1", pending_count=3, limit=3) + detail = task_enqueue.build_rate_limit_detail(exc, repo, scope="user") + + assert detail["code"] == task_enqueue.ERROR_CODE_USER_QUEUE_FULL + assert detail["queued_count"] == 3 + assert detail["running_count"] == 2 + assert detail["limit"] == 3 + assert detail["estimated_wait_seconds"] > 0 + assert "排队" in detail["message"] + assert "user-1" not in detail["message"] # 不泄露内部 ID + + +def test_global_rate_limit_detail_structure(): + """503 全局繁忙：返回 SYSTEM_QUEUE_FULL。""" + repo, _, _ = _repository() + exc = task_enqueue.GlobalQueueFull(pending_count=20, limit=20) + detail = task_enqueue.build_rate_limit_detail(exc, repo, scope="global") + + assert detail["code"] == task_enqueue.ERROR_CODE_SYSTEM_QUEUE_FULL + assert detail["queued_count"] == 20 + assert detail["limit"] == 20 + assert detail["estimated_wait_seconds"] > 0 + assert "系统繁忙" in detail["message"] + + +def test_wait_estimate_uses_concurrency(): + """等待预估：排队 8 个 / 并发 4 = 2 批 × 平均耗时。""" + + class FakeRepo: + def estimate_avg_duration_seconds(self, limit=20, default_seconds=120.0): + return 100.0 + + wait = task_enqueue._estimate_wait_seconds(8, FakeRepo()) + assert wait == 200 # ceil(8/4)=2 批 × 100 秒 + + +def test_wait_estimate_repo_without_methods_uses_default(): + """仓储没有新方法（旧 mock/鸭子类型）时用默认 120 秒兜底，不抛错。""" + + class LegacyRepo: + """只实现旧接口的仓储（模拟未升级的调用方）。""" + + def count_pending_total(self): + return 0 + + wait = task_enqueue._estimate_wait_seconds(4, LegacyRepo()) + assert wait == 120 # ceil(4/4)=1 批 × 120 默认 + + +def test_rate_limit_detail_running_count_falls_back_to_zero(): + """仓储不支持 running 计数时，running_count 优雅降级为 0。""" + + class LegacyRepo: + def count_pending_total(self): + return 0 + + exc = task_enqueue.GlobalQueueFull(pending_count=20, limit=20) + detail = task_enqueue.build_rate_limit_detail(exc, LegacyRepo(), scope="global") + assert detail["running_count"] == 0 + assert detail["code"] == task_enqueue.ERROR_CODE_SYSTEM_QUEUE_FULL + + +# --------------------------------------------------------------------------- +# 4. worker 清理核心：中断任务被重置（worker 重启/超时恢复） +# --------------------------------------------------------------------------- + + +def test_worker_cleanup_resets_interrupted_running_task(): + """模拟 worker 重启：running 超 20 分钟无更新的任务被重置为 failed，原因写明。""" + repo, _, engine = _repository() + + t = _make_task(is_preview=True) # 预览任务 + repo.create(t) + t.mark_processing() # running + repo.update(t) + _age_task(engine, t.id, updated_minutes=25) # 25 分钟无进度更新 + + cleaned = _startup.cleanup_stale_running_with_session(repo, 20) + assert cleaned == 1 + + saved = repo.get(t.id) + assert saved.status == GenerationTaskStatus.FAILED + assert "中断" in saved.error_message + assert saved.error_info.get("error_type") == "WorkerInterrupted" + assert saved.completed_at is not None + + +def test_worker_cleanup_keeps_healthy_running_task(): + """正常运行中（5 分钟前有更新）的任务不被误杀。""" + repo, _, engine = _repository() + + t = _make_task() + repo.create(t) + t.mark_processing() + repo.update(t) + _age_task(engine, t.id, updated_minutes=5) + + assert _startup.cleanup_stale_running_with_session(repo, 20) == 0 + assert repo.get(t.id).status == GenerationTaskStatus.RUNNING + + +def test_worker_cleanup_resets_stale_pending_task(): + """卡 pending 超 15 分钟（worker 停止消费）的任务被重置，释放限流名额。""" + repo, _, engine = _repository() + + t = _make_task(is_preview=True) + repo.create(t) # 一直 pending + _age_task(engine, t.id, created_minutes=20) + + cleaned = _startup.cleanup_stale_pending_with_session(repo, 15) + assert cleaned == 1 + + saved = repo.get(t.id) + assert saved.status == GenerationTaskStatus.FAILED + assert saved.error_info.get("error_type") == "PendingTimeout" + # 释放名额后 pending 计数归零，新请求不再被 429 误伤 + assert repo.count_pending_total() == 0 + + +def test_worker_cleanup_pending_keeps_recent(): + """刚创建 3 分钟的 pending 任务不清理。""" + repo, _, engine = _repository() + + t = _make_task() + repo.create(t) + _age_task(engine, t.id, created_minutes=3) + + assert _startup.cleanup_stale_pending_with_session(repo, 15) == 0 + assert repo.get(t.id).status == GenerationTaskStatus.PENDING + + +def test_worker_cleanup_multiple_orphans_all_reset(): + """3 个卡死 running 任务（工单实测：3 个预览卡 80% 超 10 小时）全部恢复。""" + repo, _, engine = _repository() + + ids = [] + for i in range(3): + t = _make_task(project_id=f"p{i}", is_preview=True) + repo.create(t) + t.mark_processing() + repo.update(t) + _age_task(engine, t.id, updated_minutes=600) # 10 小时 + ids.append(t.id) + + cleaned = _startup.cleanup_stale_running_with_session(repo, 20) + assert cleaned == 3 + for tid in ids: + assert repo.get(tid).status == GenerationTaskStatus.FAILED