video-gen/video-gen-api/app/services/celery_runtime/recovery_service.py

from __future__ import annotations

import time
import json
from datetime import datetime, timezone
from typing import Any

from app.config import settings
from app.enums.celery_runtime import CeleryRuntimeEvent
from app.services.operation_log_service import log_operation_event
from app.services.redis_registry_service import (
    RedisExecutionLockUnavailable,
    datetime_to_epoch,
    get_registry_redis,
)


def _worker_instance_key(worker_instance_id: str) -> str:
    return f"{settings.CELERY_RUNTIME_WORKER_INSTANCE_PREFIX}:{worker_instance_id}"


def _worker_name_instances_key(worker_name: str) -> str:
    return f"{settings.CELERY_RUNTIME_WORKER_NAME_INSTANCE_ZSET_PREFIX}:{worker_name}"


def _worker_task_set_key(worker_instance_id: str) -> str:
    return f"{settings.CELERY_RUNTIME_WORKER_TASK_SET_PREFIX}:{worker_instance_id}"


async def startup_barrier_exists() -> bool:
    redis = await get_registry_redis()
    if redis is None:
        raise RedisExecutionLockUnavailable("Redis unavailable while checking startup barrier")
    return bool(await redis.exists(settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY))


async def set_startup_barrier() -> None:
    redis = await get_registry_redis()
    if redis is None:
        raise RedisExecutionLockUnavailable("Redis unavailable while setting startup barrier")
    await redis.set(
        settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY,
        str(datetime_to_epoch(datetime.now(timezone.utc))),
        ex=max(30, int(settings.CELERY_RUNTIME_STARTUP_BARRIER_TTL_SECONDS or 120)),
    )


async def clear_startup_barrier() -> None:
    redis = await get_registry_redis()
    if redis is None:
        return
    await redis.delete(settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY)


async def guard_periodic_recovery(*, check_global_lock: bool = True) -> dict[str, Any] | None:
    if await startup_barrier_exists():
        log_operation_event(
            domain="celery_runtime",
            event_type=CeleryRuntimeEvent.RECOVERY_BARRIER_SKIPPED.value,
            event_status="skipped",
            source="recovery",
            detail={"barrier_key": settings.CELERY_RUNTIME_STARTUP_BARRIER_KEY},
        )
        return {"skipped": "startup_barrier"}

    if check_global_lock:
        redis = await get_registry_redis()
        if redis is None:
            raise RedisExecutionLockUnavailable("Redis unavailable while checking global recovery lock")
        if await redis.exists(settings.CELERY_RUNTIME_GLOBAL_RECOVERY_LOCK_KEY):
            log_operation_event(
                domain="celery_runtime",
                event_type=CeleryRuntimeEvent.RECOVERY_BARRIER_SKIPPED.value,
                event_status="skipped",
                source="recovery",
                detail={"global_lock_key": settings.CELERY_RUNTIME_GLOBAL_RECOVERY_LOCK_KEY},
            )
            return {"skipped": "global_recovery_lock"}
    return None


async def garbage_collect_registry_pair(*, hash_key: str, zset_key: str, limit: int = 500) -> dict[str, int]:
    redis = await get_registry_redis()
    if redis is None:
        raise RedisExecutionLockUnavailable("Redis unavailable during registry GC")

    members = await redis.zrange(zset_key, 0, max(0, int(limit) - 1))
    if not members:
        return {"checked": 0, "removed": 0}

    values = await redis.hmget(hash_key, members)
    stale = [str(member) for member, value in zip(members, values) if not value]
    if stale:
        await redis.zrem(zset_key, *stale)
    return {"checked": len(members), "removed": len(stale)}


async def mark_stale_worker_instance_candidates(
    *,
    worker_name: str,
    current_worker_instance_id: str,
    supports_targeted_recovery: bool,
    emit_duplicate_log: bool = True,
) -> dict[str, Any]:
    """只扫描同一逻辑 Worker 名称下已经失活的旧主实例。

    同名存在其他活跃实例时保守跳过，避免滚动发布或错误重名导致互相抢占。
    旧实例中的任务只有在业务 String 锁已经不存在时才提前标记为恢复候选；
    真正接管仍由各业务恢复服务执行 DB lease/claim/attempt CAS。
    """
    if not supports_targeted_recovery:
        return {
            "skipped": "identity_fallback",
            "checked_instances": 0,
            "stale_instances": 0,
            "marked": 0,
        }

    normalized_name = str(worker_name or "").strip()
    normalized_current = str(current_worker_instance_id or "").strip()
    if not normalized_name or not normalized_current:
        return {
            "skipped": "invalid_identity",
            "checked_instances": 0,
            "stale_instances": 0,
            "marked": 0,
        }

    redis = await get_registry_redis()
    if redis is None:
        raise RedisExecutionLockUnavailable("Redis unavailable while discovering stale worker instances")

    index_key = _worker_name_instances_key(normalized_name)
    raw_instances = await redis.zrange(index_key, 0, -1, withscores=True)
    instances = [
        (str(member), int(float(score)))
        for member, score in raw_instances
        if str(member) != normalized_current
    ]
    if not instances:
        return {
            "checked_instances": 0,
            "active_instances": 0,
            "stale_instances": 0,
            "marked": 0,
            "live": 0,
            "invalid": 0,
        }

    keys = [_worker_instance_key(instance_id) for instance_id, _ in instances]
    values = await redis.mget(keys)
    active_instances = [
        instance_id
        for (instance_id, _), value in zip(instances, values)
        if value
    ]
    if active_instances:
        if emit_duplicate_log:
            log_operation_event(
                domain="celery_runtime",
                event_type=CeleryRuntimeEvent.WORKER_DUPLICATE_NAME.value,
                event_status="warning",
                source="worker_ready",
                detail={
                    "worker_name": normalized_name,
                    "current_worker_instance_id": normalized_current,
                    "other_active_worker_instance_ids": active_instances,
                },
            )
        return {
            "skipped": "duplicate_active_worker_name",
            "checked_instances": len(instances),
            "active_instances": len(active_instances),
            "stale_instances": 0,
            "marked": 0,
        }

    # 实例 key 由主进程 heartbeat 独立维护。key 已不存在即可判定旧主实例
    # 不再活跃；任务是否可以恢复仍必须继续检查业务 String 锁和 DB fencing。
    stale_instances = [
        instance_id
        for (instance_id, _heartbeat_at), value in zip(instances, values)
        if not value
    ]

    totals: dict[str, Any] = {
        "checked_instances": len(instances),
        "active_instances": 0,
        "stale_instances": len(stale_instances),
        "checked": 0,
        "marked": 0,
        "live": 0,
        "invalid": 0,
    }

    for stale_instance_id in stale_instances:
        log_operation_event(
            domain="celery_runtime",
            event_type=CeleryRuntimeEvent.WORKER_STALE_INSTANCE_FOUND.value,
            event_status="success",
            source="worker_ready",
            detail={
                "worker_name": normalized_name,
                "current_worker_instance_id": normalized_current,
                "old_worker_instance_id": stale_instance_id,
            },
        )
        result = await mark_worker_instance_runtime_candidates(
            old_worker_instance_id=stale_instance_id,
            worker_name=normalized_name,
        )
        for key in ("checked", "marked", "live", "invalid"):
            totals[key] += int(result.get(key, 0) or 0)

    return totals


async def mark_worker_instance_runtime_candidates(
    *,
    old_worker_instance_id: str,
    worker_name: str | None = None,
) -> dict[str, int]:
    """按旧 Worker 主实例 ID 精准标记失锁任务。"""
    redis = await get_registry_redis()
    if redis is None:
        raise RedisExecutionLockUnavailable("Redis unavailable while marking worker recovery candidates")

    normalized_instance = str(old_worker_instance_id or "").strip()
    if not normalized_instance:
        return {"checked": 0, "marked": 0, "live": 0, "invalid": 0}

    task_set_key = _worker_task_set_key(normalized_instance)
    runtime_ids = [str(value) for value in await redis.smembers(task_set_key)]
    if not runtime_ids:
        return {"checked": 0, "marked": 0, "live": 0, "invalid": 0}

    log_operation_event(
        domain="celery_runtime",
        event_type=CeleryRuntimeEvent.WORKER_INSTANCE_RECOVERY_START.value,
        event_status="started",
        source="worker_ready",
        detail={
            "worker_name": worker_name,
            "old_worker_instance_id": normalized_instance,
            "candidate_count": len(runtime_ids),
        },
    )

    raw_locations = await redis.hmget(settings.CELERY_RUNTIME_LOCATION_HASH_KEY, runtime_ids)
    locations: dict[str, dict[str, Any]] = {}
    lock_keys: list[str] = []
    invalid_ids: list[str] = []

    for runtime_id, raw in zip(runtime_ids, raw_locations):
        try:
            payload = json.loads(raw) if raw else None
        except Exception:
            payload = None

        if (
            not isinstance(payload, dict)
            or not payload.get("lock_key")
            or not payload.get("hash_key")
            or not payload.get("zset_key")
            or str(payload.get("worker_instance_id") or "") != normalized_instance
        ):
            invalid_ids.append(runtime_id)
            continue

        locations[runtime_id] = payload
        lock_keys.append(str(payload["lock_key"]))

    lock_values = await redis.mget(lock_keys) if lock_keys else []
    live_by_key = {key: bool(value) for key, value in zip(lock_keys, lock_values)}

    now_epoch = datetime_to_epoch(datetime.now(timezone.utc))
    marked = 0
    live = 0
    stale_by_hash: dict[str, list[str]] = {}

    for runtime_id, location in locations.items():
        lock_key = str(location["lock_key"])
        if live_by_key.get(lock_key):
            live += 1
            continue
        stale_by_hash.setdefault(str(location["hash_key"]), []).append(runtime_id)

    payload_by_runtime: dict[str, dict[str, Any]] = {}
    for hash_key, hash_runtime_ids in stale_by_hash.items():
        raw_payloads = await redis.hmget(hash_key, hash_runtime_ids)
        for runtime_id, raw_payload in zip(hash_runtime_ids, raw_payloads):
            try:
                payload = json.loads(raw_payload) if raw_payload else {}
            except Exception:
                payload = {}
            payload_by_runtime[runtime_id] = payload if isinstance(payload, dict) else {}

    pipe = redis.pipeline(transaction=False)
    for runtime_id, location in locations.items():
        lock_key = str(location["lock_key"])
        if live_by_key.get(lock_key):
            continue

        hash_key = str(location["hash_key"])
        zset_key = str(location["zset_key"])
        payload = payload_by_runtime.get(runtime_id, {})
        payload.update(
            {
                "runtime_state": "recovery_candidate",
                "check_at": now_epoch,
                "reason": "worker_instance_inactive_lock_missing",
                "old_worker_instance_id": normalized_instance,
            }
        )
        pipe.hset(hash_key, runtime_id, json.dumps(payload, ensure_ascii=False, default=str))
        pipe.zadd(zset_key, {runtime_id: now_epoch})
        marked += 1

    if invalid_ids:
        pipe.srem(task_set_key, *invalid_ids)

    await pipe.execute()

    result = {
        "checked": len(runtime_ids),
        "marked": marked,
        "live": live,
        "invalid": len(invalid_ids),
    }
    log_operation_event(
        domain="celery_runtime",
        event_type=CeleryRuntimeEvent.WORKER_INSTANCE_RECOVERY_DONE.value,
        event_status="success",
        source="worker_ready",
        detail={
            "worker_name": worker_name,
            "old_worker_instance_id": normalized_instance,
            **result,
        },
    )
    return result


async def garbage_collect_worker_registry(*, limit: int = 500) -> dict[str, int]:
    """清理 V2 Worker 名称索引和旧实例 Task Set 中的失配成员。"""
    redis = await get_registry_redis()
    if redis is None:
        raise RedisExecutionLockUnavailable("Redis unavailable during worker registry GC")

    checked_names = 0
    checked_instances = 0
    removed_instances = 0
    removed_task_members = 0
    now = int(time.time())
    grace = max(
        int(settings.CELERY_RUNTIME_WORKER_HEARTBEAT_TTL_SECONDS or 90),
        int(settings.CELERY_RUNTIME_WORKER_STALE_GRACE_SECONDS or 120),
    )
    pattern = f"{settings.CELERY_RUNTIME_WORKER_NAME_INSTANCE_ZSET_PREFIX}:*"

    async for raw_index_key in redis.scan_iter(match=pattern, count=max(10, min(limit, 500))):
        if checked_names >= limit:
            break
        checked_names += 1
        index_key = str(raw_index_key)
        raw_instances = await redis.zrange(index_key, 0, -1, withscores=True)

        for raw_instance_id, raw_score in raw_instances:
            if checked_instances >= limit:
                break
            checked_instances += 1
            instance_id = str(raw_instance_id)
            score = int(float(raw_score))
            if now - score < grace:
                continue
            if await redis.exists(_worker_instance_key(instance_id)):
                continue

            task_set_key = _worker_task_set_key(instance_id)
            runtime_ids = [str(value) for value in await redis.smembers(task_set_key)]
            if runtime_ids:
                raw_locations = await redis.hmget(settings.CELERY_RUNTIME_LOCATION_HASH_KEY, runtime_ids)
                stale_members: list[str] = []
                for runtime_id, raw_location in zip(runtime_ids, raw_locations):
                    try:
                        location = json.loads(raw_location) if raw_location else None
                    except Exception:
                        location = None
                    if (
                        not isinstance(location, dict)
                        or str(location.get("worker_instance_id") or "") != instance_id
                    ):
                        stale_members.append(runtime_id)
                if stale_members:
                    removed_task_members += int(await redis.srem(task_set_key, *stale_members) or 0)

            if int(await redis.scard(task_set_key) or 0) == 0:
                await redis.delete(task_set_key)
                removed_instances += int(await redis.zrem(index_key, instance_id) or 0)

        if int(await redis.zcard(index_key) or 0) == 0:
            await redis.delete(index_key)

    return {
        "checked_names": checked_names,
        "checked_instances": checked_instances,
        "removed_instances": removed_instances,
        "removed_task_members": removed_task_members,
    }