292 lines
10 KiB
Python
292 lines
10 KiB
Python
# app/tasks/generation_recovery_tasks.py
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any, Awaitable, Callable, Dict
|
|
|
|
from app.config import settings
|
|
from app.enums.celery_queue import CeleryQueue
|
|
from app.models.base import async_session
|
|
from app.services.redis_registry_service import get_registry_redis, redis_acquire_lock, redis_release_lock
|
|
from app.tasks.async_runner import run_async
|
|
from app.tasks.celery_app import celery_app
|
|
|
|
logger = logging.getLogger("video_gen")
|
|
|
|
RECOVERY_QUEUE = settings.CELERY_RECOVERY_QUEUE or CeleryQueue.GEN_RECOVERY.value
|
|
RecoveryRunner = Callable[[], Awaitable[Dict[str, Any]]]
|
|
|
|
|
|
async def _run_download_once() -> Dict[str, Any]:
|
|
from app.services.generation_recovery_service import recover_download_tasks_once
|
|
|
|
async with async_session() as db:
|
|
return await recover_download_tasks_once(db)
|
|
|
|
|
|
async def _run_generation_once() -> Dict[str, Any]:
|
|
from app.services.generation_recovery_service import recover_generation_tasks_once
|
|
|
|
async with async_session() as db:
|
|
return await recover_generation_tasks_once(db)
|
|
|
|
|
|
async def _run_due_poll_dispatch_once() -> Dict[str, Any]:
|
|
from app.services.generation_recovery_service import dispatch_due_poll_tasks_once
|
|
|
|
async with async_session() as db:
|
|
return await dispatch_due_poll_tasks_once(db)
|
|
|
|
|
|
async def _run_module_async_once() -> Dict[str, Any]:
|
|
from app.services.module_async_recovery_service import recover_module_async_tasks_once
|
|
|
|
async with async_session() as db:
|
|
return await recover_module_async_tasks_once(db)
|
|
|
|
|
|
async def _run_shot_split_once() -> Dict[str, Any]:
|
|
from app.services.shot_replicate_recovery_service import recover_shot_split_tasks_once
|
|
|
|
async with async_session() as db:
|
|
return await recover_shot_split_tasks_once(db)
|
|
|
|
|
|
async def _run_with_execution_lock(
|
|
*,
|
|
lock_key: str,
|
|
log_context: str,
|
|
runner: RecoveryRunner,
|
|
ttl_seconds: int | None = None,
|
|
) -> Dict[str, Any]:
|
|
"""恢复任务执行锁。
|
|
|
|
worker_ready 的启动锁只保证“只投递一次”；如果 broker 中残留旧消息，
|
|
或者人工手动触发恢复任务，仍可能并发执行。这里再加执行锁，避免多个
|
|
恢复扫描同时扫库、抢行锁、抢连接池。
|
|
"""
|
|
redis = await get_registry_redis()
|
|
token: str | None = None
|
|
if redis is not None:
|
|
token = await redis_acquire_lock(
|
|
lock_key=lock_key,
|
|
ttl_seconds=int(ttl_seconds or settings.CELERY_RECOVERY_TASK_LOCK_TTL_SECONDS or 600),
|
|
log_context=log_context,
|
|
)
|
|
if not token:
|
|
return {"skipped": "lock_held", "lock_key": lock_key}
|
|
else:
|
|
# Redis 不可用时仍允许 DB fallback 执行一次，避免恢复能力彻底失效。
|
|
logger.warning("恢复任务执行锁不可用，降级直接执行。context=%s", log_context)
|
|
|
|
try:
|
|
result = await runner()
|
|
result["execution_lock"] = "lock_acquired" if token else "redis_unavailable_run_db_fallback"
|
|
return result
|
|
finally:
|
|
if token:
|
|
await redis_release_lock(lock_key=lock_key, token=token, log_context=log_context)
|
|
|
|
|
|
async def _is_lock_held(lock_key: str) -> bool:
|
|
redis = await get_registry_redis()
|
|
if redis is None:
|
|
return False
|
|
try:
|
|
return bool(await redis.exists(lock_key))
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
async def _startup_or_generation_recovery_running() -> str | None:
|
|
# Beat 触发 dispatcher 时，如果启动容灾或完整生成容灾还在跑，直接跳过本轮。
|
|
# gen_recovery concurrency=1 已经能串行；这里是多机部署、残留消息、手动触发时的双保险。
|
|
lock_checks = [
|
|
("startup_recovery", settings.CELERY_RECOVERY_STARTUP_TASK_LOCK_KEY),
|
|
("generation_recovery", settings.GENERATION_RECOVERY_LOCK_KEY),
|
|
]
|
|
for name, lock_key in lock_checks:
|
|
if await _is_lock_held(lock_key):
|
|
return name
|
|
return None
|
|
|
|
|
|
async def _run_due_poll_dispatch_with_guard() -> Dict[str, Any]:
|
|
running = await _startup_or_generation_recovery_running()
|
|
if running:
|
|
return {"skipped": "recovery_lock_held", "lock": running}
|
|
return await _run_due_poll_dispatch_once()
|
|
|
|
|
|
async def _acquire_download_recovery_loop_lock() -> tuple[bool, str]:
|
|
"""下载恢复循环锁。
|
|
|
|
Redis 不可用时降级为直接执行 DB fallback，避免恢复能力彻底失效；
|
|
Redis 可用但锁被其他 worker 持有时，本轮跳过，不再重复投递下一轮。
|
|
"""
|
|
redis = await get_registry_redis()
|
|
if redis is None:
|
|
return True, "redis_unavailable_run_db_fallback"
|
|
token = await redis_acquire_lock(
|
|
lock_key=settings.DOWNLOAD_RECOVERY_LOOP_LOCK_KEY,
|
|
ttl_seconds=int(settings.DOWNLOAD_RECOVERY_LOOP_LOCK_TTL_SECONDS or 55),
|
|
log_context="download_recovery_loop",
|
|
)
|
|
return (bool(token), "lock_acquired" if token else "lock_held")
|
|
|
|
|
|
def _schedule_next_download_recovery_loop() -> None:
|
|
if not celery_app or not bool(getattr(settings, "DOWNLOAD_RECOVERY_LOOP_ENABLED", False)):
|
|
return
|
|
try:
|
|
recover_download_tasks_once.apply_async(
|
|
countdown=max(1, int(settings.DOWNLOAD_RECOVERY_INTERVAL_SECONDS or 60)),
|
|
queue=RECOVERY_QUEUE,
|
|
priority=settings.DOWNLOAD_TASK_PRIORITY_RECOVER,
|
|
)
|
|
except Exception:
|
|
logger.exception("下载恢复循环下一轮投递失败")
|
|
|
|
|
|
async def _run_startup_recovery_once() -> Dict[str, Any]:
|
|
"""启动容灾协调器：串行跑恢复扫描。
|
|
|
|
真实业务任务仍投递回原队列：
|
|
- 创建/提词/视频分析 -> gen_chatapi_create
|
|
- provider poll -> gen_provider_poll
|
|
- 下载/ffmpeg 切片 -> gen_result_download
|
|
恢复扫描本身只走 gen_recovery，避免堵住业务 worker。
|
|
"""
|
|
return await _run_with_execution_lock(
|
|
lock_key=settings.CELERY_RECOVERY_STARTUP_TASK_LOCK_KEY,
|
|
log_context="startup_recovery_once",
|
|
runner=_run_startup_recovery_steps,
|
|
)
|
|
|
|
|
|
async def _run_startup_recovery_steps() -> Dict[str, Any]:
|
|
results: Dict[str, Any] = {}
|
|
|
|
steps: list[tuple[str, str, str, RecoveryRunner]] = [
|
|
(
|
|
"module_async",
|
|
settings.MODULE_ASYNC_RECOVERY_LOCK_KEY,
|
|
"module_async_recovery",
|
|
_run_module_async_once,
|
|
),
|
|
(
|
|
"shot_split",
|
|
settings.SHOT_SPLIT_RECOVERY_LOCK_KEY,
|
|
"shot_split_recovery",
|
|
_run_shot_split_once,
|
|
),
|
|
(
|
|
"generation",
|
|
settings.GENERATION_RECOVERY_LOCK_KEY,
|
|
"generation_recovery",
|
|
_run_generation_once,
|
|
),
|
|
(
|
|
"download",
|
|
settings.DOWNLOAD_RECOVERY_LOCK_KEY,
|
|
"download_recovery",
|
|
_run_download_once,
|
|
),
|
|
]
|
|
|
|
for name, lock_key, log_context, runner in steps:
|
|
try:
|
|
results[name] = await _run_with_execution_lock(
|
|
lock_key=lock_key,
|
|
log_context=log_context,
|
|
runner=runner,
|
|
)
|
|
except Exception as exc:
|
|
logger.exception("启动容灾步骤执行失败。step=%s", name)
|
|
results[name] = {"error": str(exc)}
|
|
|
|
return {"steps": results}
|
|
|
|
|
|
if celery_app:
|
|
|
|
@celery_app.task(
|
|
name="recovery.startup_recovery_once",
|
|
bind=True,
|
|
soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
|
|
time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
|
|
)
|
|
def startup_recovery_once(self) -> Dict[str, Any]:
|
|
return run_async(_run_startup_recovery_once())
|
|
|
|
|
|
@celery_app.task(
|
|
name="generation.recover_download_tasks_once",
|
|
bind=True,
|
|
soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
|
|
time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
|
|
)
|
|
def recover_download_tasks_once(self) -> Dict[str, Any]:
|
|
acquired, reason = run_async(_acquire_download_recovery_loop_lock())
|
|
if not acquired:
|
|
return {"skipped": reason}
|
|
try:
|
|
result = run_async(
|
|
_run_with_execution_lock(
|
|
lock_key=settings.DOWNLOAD_RECOVERY_LOCK_KEY,
|
|
log_context="download_recovery",
|
|
runner=_run_download_once,
|
|
)
|
|
)
|
|
result["loop_lock"] = reason
|
|
return result
|
|
finally:
|
|
_schedule_next_download_recovery_loop()
|
|
|
|
|
|
@celery_app.task(
|
|
name="generation.recover_generation_tasks_once",
|
|
bind=True,
|
|
soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
|
|
time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
|
|
)
|
|
def recover_generation_tasks_once(self) -> Dict[str, Any]:
|
|
return run_async(
|
|
_run_with_execution_lock(
|
|
lock_key=settings.GENERATION_RECOVERY_LOCK_KEY,
|
|
log_context="generation_recovery",
|
|
runner=_run_generation_once,
|
|
)
|
|
)
|
|
|
|
|
|
@celery_app.task(
|
|
name="generation.dispatch_due_poll_tasks",
|
|
bind=True,
|
|
soft_time_limit=settings.CELERY_RECOVERY_SOFT_TIME_LIMIT_SECONDS,
|
|
time_limit=settings.CELERY_RECOVERY_TIME_LIMIT_SECONDS,
|
|
)
|
|
def dispatch_due_poll_tasks(self) -> Dict[str, Any]:
|
|
return run_async(
|
|
_run_with_execution_lock(
|
|
lock_key=settings.POLL_DUE_DISPATCH_LOCK_KEY,
|
|
log_context="due_poll_dispatch",
|
|
runner=_run_due_poll_dispatch_with_guard,
|
|
ttl_seconds=int(settings.POLL_DUE_DISPATCH_LOCK_TTL_SECONDS or 55),
|
|
)
|
|
)
|
|
|
|
else:
|
|
|
|
class _DisabledTask:
|
|
def delay(self, *args: Any, **kwargs: Any) -> None:
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
def apply_async(self, *args: Any, **kwargs: Any) -> None:
|
|
raise RuntimeError("Celery is disabled")
|
|
|
|
startup_recovery_once = _DisabledTask()
|
|
recover_download_tasks_once = _DisabledTask()
|
|
recover_generation_tasks_once = _DisabledTask()
|
|
dispatch_due_poll_tasks = _DisabledTask()
|