crypto_okx_hub/lib/update_lib.py

"""中控本机快速更新:先同步 git pull,再异步 pm2 restart."""
from __future__ import annotations

import os
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any

_ROOT = Path(__file__).resolve().parent.parent
_update_lock = threading.Lock()
_update_pending_until = 0.0
_UPDATE_LOG = _ROOT / "logs" / "hub_update.log"


def _repo_root() -> Path:
    env_root = (os.getenv("REPO") or "").strip()
    if env_root:
        return Path(env_root)
    return _ROOT


def _run_git(args: list[str], *, cwd: Path, timeout: float = 15) -> tuple[bool, str]:
    try:
        proc = subprocess.run(
            ["git", *args],
            cwd=str(cwd),
            capture_output=True,
            text=True,
            timeout=timeout,
        )
    except FileNotFoundError:
        return False, "未找到 git"
    except subprocess.TimeoutExpired:
        return False, "git 超时"
    except Exception as e:
        return False, str(e) or "git 失败"
    out = (proc.stdout or "").strip() or (proc.stderr or "").strip()
    return proc.returncode == 0, out


def _append_log(line: str) -> None:
    try:
        _UPDATE_LOG.parent.mkdir(parents=True, exist_ok=True)
        with open(_UPDATE_LOG, "a", encoding="utf-8") as f:
            f.write(time.strftime("%Y-%m-%d %H:%M:%S") + " " + line + "\n")
    except Exception:
        pass


def read_git_version() -> dict[str, Any]:
    root = _repo_root()
    out: dict[str, Any] = {
        "ok": True,
        "repo": str(root),
        "sha": "",
        "short_sha": "",
        "branch": "",
        "dirty": False,
        "msg": "",
        "online": True,
        "platform": sys.platform,
    }
    if not (root / ".git").exists() and not (root / ".git").is_file():
        out["ok"] = False
        out["msg"] = "非 git 仓库"
        return out
    ok, sha = _run_git(["rev-parse", "HEAD"], cwd=root)
    if not ok or not sha:
        out["ok"] = False
        out["msg"] = sha or "无法读取 HEAD"
        return out
    out["sha"] = sha
    out["short_sha"] = sha[:7]
    ok_b, branch = _run_git(["rev-parse", "--abbrev-ref", "HEAD"], cwd=root)
    if ok_b and branch:
        out["branch"] = branch
    ok_d, dirty = _run_git(["status", "--porcelain"], cwd=root)
    if ok_d:
        out["dirty"] = bool(dirty.strip())
    return out


def _pm2_app_name() -> str:
    return (os.getenv("PM2_APP_NAME") or "").strip() or "crypto_okx_hub"


def _schedule_pm2_restart(*, delay_seconds: float = 1.5) -> None:
    app_name = _pm2_app_name()
    delay = max(0.8, float(delay_seconds))

    def _worker() -> None:
        time.sleep(delay)
        _append_log(f"pm2 restart {app_name} --update-env")
        try:
            proc = subprocess.run(
                ["pm2", "restart", app_name, "--update-env"],
                capture_output=True,
                text=True,
                timeout=120,
            )
            _append_log(
                f"pm2 restart returncode={proc.returncode} "
                f"out={(proc.stdout or proc.stderr or '')[:300]}"
            )
        except Exception as e:
            _append_log("pm2 restart error: " + (str(e) or "unknown"))

    threading.Thread(target=_worker, name="hub-pm2-restart", daemon=True).start()


def schedule_pull_and_restart(*, delay_seconds: float = 1.5) -> dict[str, Any]:
    """同步 git pull,再异步 pm2 restart(便于前端立刻知道是否有新版本)."""
    if not sys.platform.startswith("linux"):
        msg = f"仅 Linux 服务器支持更新(当前 platform={sys.platform})"
        _append_log("reject: " + msg)
        return {"ok": False, "msg": msg, "deferred": False, "platform": sys.platform}

    root = _repo_root()
    if not (root / "ecosystem.config.cjs").is_file():
        msg = f"缺少 ecosystem.config.cjs: {root}"
        _append_log("reject: " + msg)
        return {"ok": False, "msg": msg, "deferred": False}

    global _update_pending_until
    now = time.time()
    with _update_lock:
        if now < _update_pending_until:
            return {"ok": False, "msg": "已有更新任务进行中,请稍后再试", "deferred": False}
        _update_pending_until = now + 120.0

    before = read_git_version()
    before_sha = str(before.get("sha") or "")
    _append_log(f"git pull start before={before.get('short_sha') or ''} root={root}")
    ok_pull, pull_out = _run_git(["pull", "--ff-only"], cwd=root, timeout=90)
    if not ok_pull:
        # 回退再试普通 pull
        ok_pull, pull_out = _run_git(["pull"], cwd=root, timeout=90)
    after = read_git_version()
    after_sha = str(after.get("sha") or "")
    changed = bool(before_sha and after_sha and before_sha != after_sha)
    _append_log(
        f"git pull done ok={ok_pull} changed={changed} "
        f"after={after.get('short_sha') or ''} msg={pull_out[:200]}"
    )

    if not ok_pull and not after.get("ok"):
        with _update_lock:
            _update_pending_until = 0.0
        return {
            "ok": False,
            "msg": pull_out or "git pull 失败",
            "deferred": False,
            "before_sha": before_sha,
            "before_short_sha": before.get("short_sha") or "",
        }

    # pull 失败但 HEAD 仍可读:仍尝试重启(可能网络问题但本地可跑)
    try:
        _schedule_pm2_restart(delay_seconds=delay_seconds)
    except Exception as e:
        with _update_lock:
            _update_pending_until = 0.0
        return {"ok": False, "msg": str(e) or "触发重启失败", "deferred": False}

    if not ok_pull:
        msg = "git pull 未成功,已触发重启: " + (pull_out or "")
    elif changed:
        msg = f"已拉取新版本 {(before.get('short_sha') or '?')} → {(after.get('short_sha') or '?')},即将重启"
    else:
        msg = "代码已是最新,即将重启进程"

    return {
        "ok": True,
        "msg": msg,
        "deferred": True,
        "changed": changed,
        "pull_ok": ok_pull,
        "pull_msg": pull_out[:500],
        "before_sha": before_sha,
        "before_short_sha": before.get("short_sha") or "",
        "after_sha": after_sha,
        "after_short_sha": after.get("short_sha") or "",
        "branch": after.get("branch") or before.get("branch") or "",
        "repo": str(root),
        "log": str(_UPDATE_LOG),
    }