feat(generation): worker孤儿任务自动恢复 + 429限流结构化提示 (#1677) #1710

Merged
auto-approve-bot merged 1 commits from feature/task-fault-tolerance-orphan-recovery into develop 2026-09-05 11:37:51 +08:00
9 changed files with 614 additions and 41 deletions
+19 -11
View File
@@ -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",
+27 -14
View File
@@ -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)
+120 -2
View File
@@ -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,
+10 -2
View File
@@ -22,10 +22,18 @@ celery_app.conf.imports = (
)
# Celery Beat 定时任务调度
# 注:worker 单实例内嵌 beatentrypoint-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},
},
}
+33 -8
View File
@@ -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:
+38 -4
View File
@@ -1,14 +1,18 @@
"""定期清理任务 — Celery Beat 调度。
包含:
- cleanup_stale_pending_tasks: 定期清理卡在 pending 超时的 generation_tasks
- cleanup_stale_pending_tasks: 定期清理卡在 pending 超时的 generation_tasksworker 停止消费时占位)
- 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}
@@ -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)
@@ -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]: ...
@@ -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