crypto_okx_hub/lib/update_lib.py

"""中控本机快速更新:git pull + pm2 restart."""
from __future__ import annotations

import os
import shlex
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) -> tuple[bool, str]:
    try:
        proc = subprocess.run(
            ["git", *args],
            cwd=str(cwd),
            capture_output=True,
            text=True,
            timeout=15,
        )
    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 _run_pull_and_restart(root: Path, script: Path) -> None:
    log_path = _UPDATE_LOG
    try:
        log_path.parent.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    env = os.environ.copy()
    env["REPO"] = str(root)
    if not (env.get("PM2_APP_NAME") or "").strip():
        env["PM2_APP_NAME"] = "crypto_okx_hub"
    try:
        with open(log_path, "a", encoding="utf-8") as logf:
            logf.write(
                time.strftime("%Y-%m-%d %H:%M:%S")
                + f" start pull_and_restart REPO={root}\n"
            )
            logf.flush()
            proc = subprocess.run(
                ["bash", str(script)],
                cwd=str(root),
                env=env,
                stdout=logf,
                stderr=subprocess.STDOUT,
                text=True,
                timeout=180,
            )
            logf.write(
                time.strftime("%Y-%m-%d %H:%M:%S")
                + f" finish returncode={proc.returncode}\n"
            )
    except Exception as e:
        _append_log("pull_and_restart error: " + (str(e) or "unknown"))


def schedule_pull_and_restart(*, delay_seconds: float = 2.0) -> dict[str, Any]:
    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()
    script = root / "deploy" / "pull_and_restart.sh"
    if not script.is_file():
        msg = f"缺少更新脚本: {script}"
        _append_log("reject: " + msg)
        return {"ok": False, "msg": msg, "deferred": False}

    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:
            msg = "已有更新任务进行中,请稍后再试"
            return {"ok": False, "msg": msg, "deferred": False}
        _update_pending_until = now + 180.0

    version = read_git_version()
    delay = max(1.0, float(delay_seconds))

    def _worker() -> None:
        try:
            time.sleep(delay)
            _run_pull_and_restart(root, script)
        finally:
            with _update_lock:
                # 保留短暂冷却,避免连点
                pass

    try:
        t = threading.Thread(target=_worker, name="hub-self-update", daemon=True)
        t.start()
        _append_log(
            f"accepted before={version.get('short_sha') or ''} delay={delay}s root={root}"
        )
        return {
            "ok": True,
            "msg": "更新已受理,即将 git pull 并重启",
            "deferred": True,
            "before_sha": version.get("sha") or "",
            "before_short_sha": version.get("short_sha") or "",
            "branch": version.get("branch") or "",
            "repo": str(root),
            "log": str(_UPDATE_LOG),
        }
    except Exception as e:
        with _update_lock:
            _update_pending_until = 0.0
        msg = str(e) or "触发更新失败"
        _append_log("reject: " + msg)
        return {"ok": False, "msg": msg, "deferred": False}