xiaoxia-saas/packages/shared/celery_orphan_guard.py

"""孤儿任务消息撤销与执行前状态守卫（API / Worker 共享）。

#1714 / #1710 缺陷修复：超时清理/孤儿恢复把 DB 任务标记为 failed/cancelled
后，Redis 队列里对应的 Celery 消息仍然存在；worker 重启或重新拉取时该消息
被再次执行，状态机抛「非法状态转换: failed → running」，旧实现打印 ERROR 后
继续跑，最终产出半成品。

防御两道：
1. 清理任务标 failed 时，调用 revoke_and_purge() 撤销（celery revoke 广播，
   通知在线 worker 丢弃）并直接扫描 Redis 队列移除消息体（worker 下线期间
   队列中的消息 revoke 广播收不到，必须物理移除）；
2. 任务真正开始业务逻辑前，调用 ensure_task_claimable() 校验 DB 状态，
   非 pending 的消息直接丢弃（抛 StaleTaskDiscarded，task 捕获后安全返回，
   不进入渲染/转码，不产出半成品）。
"""

from __future__ import annotations

import base64
import json
import logging
from collections.abc import Callable, Iterable
from typing import Any

logger = logging.getLogger(__name__)


class StaleTaskDiscarded(Exception):
    """任务消息已作废（DB 中任务已是终态），应安全中止、丢弃消息。"""

    def __init__(self, task_id: str, status: str):
        self.task_id = task_id
        self.status = status
        super().__init__(f"任务 {task_id} 已是终态 {status}，丢弃重复/作废消息")


# 终态状态值集合：处于这些状态的任务消息一律不执行
TERMINAL_STATUS_VALUES = frozenset({"failed", "cancelled", "completed"})


def ensure_task_claimable(
    task_id: str,
    get_status: Callable[[str], str | None],
    *,
    task_label: str = "任务",
) -> str:
    """执行前守卫：任务必须处于可领取状态（pending）。

    Args:
        task_id: 业务任务 ID
        get_status: 回调，返回 DB 中任务当前状态字符串；返回 None 表示任务不存在
        task_label: 日志用任务类型名

    Returns:
        当前状态字符串（pending）；任务不存在时返回空串（由调用方处理 not found）

    Raises:
        StaleTaskDiscarded: 任务已是终态（failed/cancelled/completed），消息必须丢弃
    """
    status = get_status(task_id)
    if status is None:
        return ""
    if status in TERMINAL_STATUS_VALUES:
        logger.warning("[%s] task_id=%s 状态已为 %s，消息作废，丢弃不执行", task_label, task_id, status)
        raise StaleTaskDiscarded(task_id, status)
    return status


def _extract_business_ids(raw: bytes) -> tuple[str | None, str | None]:
    """从 Redis 中的 Celery 消息提取 (celery 消息 ID, 业务任务 ID)。

    Redis transport 存储格式为 JSON 信封：
        {"body": base64(json), "headers": {"id": <celery id>, "task": <name>, ...}, ...}
    body 解码后 Celery task 协议为 [args, kwargs, embed]；
    generate_video / ingest_asset 均以 args=[业务任务ID] 投递。

    无法解析时返回 (None, None)（保守保留该消息，绝不误删）。
    """
    try:
        envelope = json.loads(raw)
        celery_id = None
        headers = envelope.get("headers") or {}
        if isinstance(headers, dict):
            celery_id = headers.get("id")
        body = envelope.get("body")
        if not body:
            return celery_id, None
        decoded = base64.b64decode(body)
        payload = json.loads(decoded)
        # 两种 body 形态：
        # 1. 标准 Celery task 消息：[args, kwargs, embed] 三元组 → 业务 ID 在 payload[0][0]
        # 2. 裸 producer 发布：body 即 args 数组 ["biz-id"] → 业务 ID 在 payload[0]
        args = None
        if isinstance(payload, dict):
            args = payload.get("args")
        elif isinstance(payload, (list, tuple)) and payload:
            first = payload[0]
            if isinstance(first, (list, tuple)):
                args = first  # 三元组：[args, kwargs, embed]
            else:
                args = payload  # body 本身就是 args
        if isinstance(args, (list, tuple)) and args and args[0] is not None:
            return celery_id, str(args[0])
        return celery_id, None
    except Exception:
        return None, None


def purge_stale_messages_from_queues(
    broker_url: str,
    queue_names: Iterable[str],
    business_task_ids: Iterable[str] = (),
    celery_task_ids: Iterable[str] = (),
) -> int:
    """扫描 Redis 队列，移除作废任务的待消费消息。

    同时按业务任务 ID（消息 args[0]）和 celery 消息 ID（headers.id）匹配，
    任一命中即移除。未命中或无法解析的消息原样保留（保持相对顺序）。

    Returns:
        实际移除的消息条数
    """
    biz_ids = {bid for bid in business_task_ids if bid}
    msg_ids = {mid for mid in celery_task_ids if mid}
    if not biz_ids and not msg_ids:
        return 0

    try:
        import redis
    except ImportError:
        logger.warning("redis-py 不可用，跳过队列消息清理")
        return 0

    try:
        client = redis.Redis.from_url(broker_url)
        client.ping()
    except Exception as e:
        logger.warning("连接 Redis 清理作废消息失败: %s", e)
        return 0

    removed_total = 0
    try:
        for queue in queue_names:
            removed_total += _purge_one_queue(client, queue, biz_ids, msg_ids)
    finally:
        try:
            client.close()
        except Exception:
            pass
    if removed_total:
        logger.info(
            "从 Redis 队列移除 %d 条作废消息（biz=%s, celery=%s）",
            removed_total,
            sorted(biz_ids),
            sorted(msg_ids),
        )
    return removed_total


def _purge_one_queue(client: Any, queue_name: str, biz_ids: set[str], msg_ids: set[str]) -> int:
    try:
        raw_messages = client.lrange(queue_name, 0, -1)
    except Exception as e:
        logger.warning("读取队列 %s 失败: %s", queue_name, e)
        return 0
    if not raw_messages:
        return 0

    keep: list[bytes] = []
    removed = 0
    for raw in raw_messages:
        celery_id, biz_id = _extract_business_ids(raw)
        hit = (biz_id is not None and biz_id in biz_ids) or (celery_id is not None and celery_id in msg_ids)
        if hit:
            removed += 1
            continue
        keep.append(raw)

    if removed:
        try:
            pipe = client.pipeline()
            pipe.delete(queue_name)
            if keep:
                pipe.rpush(queue_name, *keep)
            pipe.execute()
        except Exception as e:
            logger.warning("重写队列 %s 失败: %s", queue_name, e)
            return 0
    return removed


def revoke_and_purge(
    celery_app: Any,
    broker_url: str,
    business_task_ids: Iterable[str] = (),
    celery_task_ids: Iterable[str] = (),
    *,
    queue_names: Iterable[str] = ("generation", "transcode", "celery"),
) -> int:
    """撤销作废任务：revoke 广播（在线 worker）+ 物理清理 Redis 队列消息。

    Args:
        celery_app: Celery app 实例（worker 端 worker_app.celery_app.celery_app）
        broker_url: Redis broker URL
        business_task_ids: 业务任务 ID（generation_tasks.id / ingest_jobs.id）
        celery_task_ids: 入队时记录的 celery 消息 ID
        queue_names: 需要扫描清理的队列名

    Returns:
        从队列中实际移除的消息条数
    """
    for tid in celery_task_ids:
        if not tid:
            continue
        try:
            celery_app.control.revoke(tid)
        except Exception as e:
            logger.warning("revoke celery 消息 %s 失败: %s", tid, e)

    return purge_stale_messages_from_queues(
        broker_url, queue_names, business_task_ids=business_task_ids, celery_task_ids=celery_task_ids
    )