Nexus/server/infrastructure/redis/server_batch_store.py

"""Server batch job live state in Redis (2h TTL)."""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from typing import Any, Optional

from server.infrastructure.redis.client import get_redis

logger = logging.getLogger("nexus.server_batch_store")

REDIS_KEY_PREFIX = "server_batch:"
REDIS_NEXT_ID = "server_batch:next_id"
REDIS_TTL_SECONDS = 7200
RECONCILE_INTERVAL_SECONDS = 60
STUCK_NO_PROGRESS_SECONDS = 3600
INTERRUPTED_MESSAGE = "服务重启导致任务中断"
STUCK_MESSAGE = "执行超时无响应（系统自动标记）"
TERMINAL_STATUSES = frozenset({"completed", "failed", "partial", "cancelled"})


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


async def next_job_id() -> int:
    redis = get_redis()
    return int(await redis.incr(REDIS_NEXT_ID))


def _key(job_id: int) -> str:
    return f"{REDIS_KEY_PREFIX}{job_id}"


def progress_stats(live: dict[str, Any]) -> dict[str, Any]:
    server_ids = live.get("server_ids") or []
    results = live.get("results") or {}
    total = len(server_ids)
    done = 0
    success = 0
    failed = 0
    for sid in server_ids:
        entry = results.get(str(sid))
        if not entry:
            continue
        done += 1
        if entry.get("success"):
            success += 1
        else:
            failed += 1
    remaining = max(0, total - done)
    return {
        "total": total,
        "done": done,
        "remaining": remaining,
        "success": success,
        "failed": failed,
        "progress": f"{success}/{total}" if total else "—",
    }


async def save_live(live: dict[str, Any]) -> None:
    redis = get_redis()
    job_id = int(live["id"])
    live["updated_at"] = _now_iso()
    await redis.set(_key(job_id), json.dumps(live, ensure_ascii=False), ex=REDIS_TTL_SECONDS)


async def list_running_live_jobs() -> list[dict[str, Any]]:
    """Scan Redis for all server batch jobs still in running status."""
    redis = get_redis()
    running: list[dict[str, Any]] = []
    cursor = 0
    while True:
        cursor, keys = await redis.scan(cursor, match=f"{REDIS_KEY_PREFIX}*", count=50)
        for key in keys:
            if key == REDIS_NEXT_ID:
                continue
            suffix = key.replace(REDIS_KEY_PREFIX, "", 1)
            try:
                job_id = int(suffix)
            except ValueError:
                continue
            live = await get_live(job_id)
            if live and live.get("status") == "running":
                running.append(live)
        if cursor == 0:
            break
    return running


async def get_live(job_id: int) -> Optional[dict[str, Any]]:
    redis = get_redis()
    raw = await redis.get(_key(job_id))
    if not raw:
        return None
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        logger.warning("Corrupt server batch live state job_id=%s", job_id)
        return None


def init_live(
    *,
    job_id: int,
    op: str,
    label: str,
    server_ids: list[int],
    operator: str,
    params: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    return {
        "id": job_id,
        "op": op,
        "label": label,
        "status": "running",
        "server_ids": list(server_ids),
        "results": {},
        "operator": operator,
        "params": params or {},
        "created_at": _now_iso(),
        "updated_at": _now_iso(),
    }


def record_result(live: dict[str, Any], item: dict[str, Any]) -> None:
    results = live.setdefault("results", {})
    results[str(item["server_id"])] = item


def finalize_status(live: dict[str, Any]) -> str:
    stats = progress_stats(live)
    total = stats["total"]
    success = stats["success"]
    failed = stats["failed"]
    if total == 0:
        status = "failed"
    elif failed == 0:
        status = "completed"
    elif success == 0:
        status = "failed"
    else:
        status = "partial"
    live["status"] = status
    return status


def parse_iso(ts: str) -> Optional[datetime]:
    try:
        return datetime.fromisoformat(ts)
    except (ValueError, TypeError):
        return None


def pending_server_ids(live: dict[str, Any]) -> list[int]:
    """Server IDs in the job with no result row yet."""
    server_ids = live.get("server_ids") or []
    results = live.get("results") or {}
    pending: list[int] = []
    for sid in server_ids:
        if str(sid) not in results:
            pending.append(int(sid))
    return pending


def apply_pending_failures(
    live: dict[str, Any],
    *,
    server_ids: list[int],
    error: str,
    auto_stuck: bool = False,
    server_name_fallback: str = "未知服务器",
) -> int:
    """Mark pending servers as failed; returns count of rows added."""
    if not server_ids:
        return 0
    results = live.setdefault("results", {})
    added = 0
    for sid in server_ids:
        key = str(sid)
        if key in results:
            continue
        item: dict[str, Any] = {
            "server_id": sid,
            "server_name": server_name_fallback,
            "success": False,
            "stdout": "",
            "error": error,
        }
        if auto_stuck:
            item["auto_stuck"] = True
        results[key] = item
        added += 1
    return added


def is_stale_running(live: dict[str, Any], now: Optional[datetime] = None) -> bool:
    """True when a running job has had no Redis update for STUCK_NO_PROGRESS_SECONDS."""
    if live.get("status") != "running":
        return False
    now = now or datetime.now(timezone.utc)
    updated = parse_iso(live.get("updated_at", ""))
    if not updated:
        updated = parse_iso(live.get("created_at", ""))
    if not updated:
        return False
    return (now - updated).total_seconds() >= STUCK_NO_PROGRESS_SECONDS


def live_from_mysql_row(row: Any) -> dict[str, Any]:
    """Rebuild minimal Redis-shaped live dict from a ServerBatchJob row."""
    import json

    try:
        server_ids = json.loads(row.server_ids or "[]")
    except json.JSONDecodeError:
        server_ids = []
    try:
        raw_results = json.loads(row.results or "[]")
    except json.JSONDecodeError:
        raw_results = []
    results: dict[str, Any] = {}
    if isinstance(raw_results, dict):
        results = dict(raw_results)
    elif isinstance(raw_results, list):
        for entry in raw_results:
            if isinstance(entry, dict) and entry.get("server_id") is not None:
                results[str(entry["server_id"])] = entry
    try:
        params = json.loads(row.params or "{}") if row.params else {}
    except json.JSONDecodeError:
        params = {}
    if not isinstance(params, dict):
        params = {}
    started = row.started_at.isoformat() if row.started_at else _now_iso()
    return {
        "id": int(row.id),
        "op": row.op or "",
        "label": row.label or "",
        "status": row.status or "running",
        "server_ids": list(server_ids),
        "results": results,
        "operator": row.operator or "",
        "params": params,
        "created_at": started,
        "updated_at": started,
    }


def live_to_api(live: dict[str, Any]) -> dict[str, Any]:
    stats = progress_stats(live)
    results = live.get("results") or {}
    server_ids = live.get("server_ids") or []
    ordered = []
    for sid in server_ids:
        entry = results.get(str(sid))
        if entry:
            ordered.append(entry)
    return {
        "job_id": int(live["id"]),
        "op": live.get("op") or "",
        "label": live.get("label") or "",
        "status": live.get("status") or "running",
        "operator": live.get("operator") or "",
        **stats,
        "results": ordered,
    }