xiaoxia-saas/apps/worker/worker_app/tasks/_startup.py

"""Worker 启动时的初始化任务 — 孤儿任务清理等。"""

import logging

from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

logger = logging.getLogger(__name__)

ORPHAN_TASK_TIMEOUT_MINUTES = 10


def cleanup_orphan_tasks(timeout_minutes: int = ORPHAN_TASK_TIMEOUT_MINUTES) -> int:
    """清理数据库中超时未更新的 running 任务。

    worker 重启或崩溃后，之前处于 running 状态的任务会变成孤儿任务，
    一直卡在 running 不动。通过 updated_at 超时判断并标记为 failed。

    Args:
        timeout_minutes: 超时时间（分钟），默认 10 分钟

    Returns:
        清理的任务数量
    """
    from packages.adapters.sqlalchemy_impl.generation_task_repository import (
        SQLAlchemyGenerationTaskRepository,
    )

    try:
        session = SessionLocal()
        repo = SQLAlchemyGenerationTaskRepository(session)
        count = repo.cleanup_stale_running(timeout_minutes)
        session.close()
        if count > 0:
            logger.warning("清理了 %d 个超时的孤儿 running 任务", count)
        else:
            logger.info("无孤儿 running 任务需要清理")
        return count
    except Exception as e:
        logger.error("清理孤儿任务失败: %s", e, exc_info=True)
        return 0


@celery_app.on_after_configure.connect
def _setup_periodic_tasks(sender, **kwargs):
    """Celery 配置完成后，注册 worker 启动钩子。"""
    pass


@celery_app.on_worker_ready.connect
def _on_worker_ready(sender, **kwargs):
    """Worker 启动完成后执行 — 清理孤儿任务。"""
    logger.info("Worker 启动完成，开始清理孤儿 running 任务...")
    count = cleanup_orphan_tasks()
    logger.info("Worker 启动清理完成，共清理 %d 个孤儿任务", count)