agento3/scripts/daemon.py

#!/usr/bin/env python3
"""
AO3 Mirror - 被动守护进程 (Daemon v2)
不再主动 e2e 检查或重启 healthy worker。
职责：
  1. 每 30s 检查 worker 是否活着，只有连续 3 轮都挂掉才重启
  2. 不做 e2e，不做 caddy 检查（让 systemd 管）
  3. 每 5min 写状态文件
"""
import json
import os
import subprocess
import sys
import time
import logging
import signal

BASE_DIR = "/home/ubuntu/ao3-mirror"
SHM_DIR = "/dev/shm/ao3"
WORKER_PORTS = [8081, 8082]
CHECK_INTERVAL = 30
STATUS_FILE = f"{SHM_DIR}/status.json"

os.makedirs(SHM_DIR, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [DAEMON] %(message)s",
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger("ao3-daemon")

state = {
    "started_at": time.time(),
    "consecutive_dead": {p: 0 for p in WORKER_PORTS},
    "restarts_today": 0,
    "alerts": [],
}


def _http_get(url: str, timeout: int = 5) -> tuple[bool, str]:
    import urllib.request
    try:
        req = urllib.request.Request(url)
        resp = urllib.request.urlopen(req, timeout=timeout)
        return (resp.status == 200, resp.read().decode("utf-8", errors="replace")[:200])
    except Exception as e:
        return (False, str(e)[:200])


def check_worker(port: int) -> bool:
    ok, body = _http_get(f"http://127.0.0.1:{port}/health", timeout=4)
    if ok:
        try:
            return json.loads(body).get("status") == "ok"
        except Exception:
            pass
    return False


def restart_worker(port: int):
    logger.warning(f"Restarting worker on port {port}")
    state["restarts_today"] += 1
    subprocess.run(["fuser", "-k", f"{port}/tcp"], capture_output=True, timeout=10)
    time.sleep(1)
    python_bin = "/home/ubuntu/.hermes/hermes-agent/venv/bin/python3"
    i = WORKER_PORTS.index(port)
    log_file = f"{BASE_DIR}/worker-{i}.log"
    proc = subprocess.Popen(
        ["taskset", "-c", str(i), python_bin, "-m", "uvicorn", "app:app",
         "--host", "127.0.0.1", "--port", str(port), "--workers", "1",
         "--loop", "uvloop", "--http", "httptools",
         "--log-level", "warning", "--timeout-keep-alive", "30"],
        cwd=BASE_DIR, stdout=open(log_file, "a"), stderr=open(log_file, "a"),
    )
    pid_file = f"{SHM_DIR}/worker-{i}.pid"
    with open(pid_file, "w") as f:
        f.write(str(proc.pid))
    logger.info(f"Worker {port} restarted (PID: {proc.pid})")


def check_workers():
    """Only restart if worker has been dead for 3 consecutive checks (90s)."""
    for port in WORKER_PORTS:
        if check_worker(port):
            state["consecutive_dead"][port] = 0
        else:
            state["consecutive_dead"][port] += 1
            if state["consecutive_dead"][port] >= 3:
                restart_worker(port)
                state["consecutive_dead"][port] = 0
    dead_count = sum(1 for p in WORKER_PORTS if state["consecutive_dead"][p] > 0)
    alive = len(WORKER_PORTS) - dead_count
    if dead_count > 0:
        logger.info(f"{alive}/{len(WORKER_PORTS)} workers alive ({dead_count} degraded)")


def save_status():
    worker_status = {}
    for i, port in enumerate(WORKER_PORTS):
        worker_status[f"worker_{i}"] = {
            "port": port, "healthy": check_worker(port), "uptime_s": int(time.time() - state["started_at"]),
        }
    data = {
        "timestamp": time.time(), "uptime": int(time.time() - state["started_at"]),
        "workers": worker_status, "restarts_today": state["restarts_today"],
    }
    with open(STATUS_FILE + ".tmp", "w") as f:
        json.dump(data, f)
    os.rename(STATUS_FILE + ".tmp", STATUS_FILE)


def handle_signal(signum, frame):
    logger.info(f"Signal {signum}, exiting")
    sys.exit(0)


def main():
    signal.signal(signal.SIGTERM, handle_signal)
    signal.signal(signal.SIGINT, handle_signal)
    logger.info("Daemon v2 started (passive mode)")
    logger.info(f"Watching {WORKER_PORTS}, check every {CHECK_INTERVAL}s, restart after 3 missed checks")
    
    # Initial startup: start all workers immediately
    for port in WORKER_PORTS:
        if not check_worker(port):
            restart_worker(port)
            time.sleep(2)
        state["consecutive_dead"][port] = 0
    
    time.sleep(5)  # Let them settle
    last_status = 0
    while True:
        try:
            check_workers()
            now = time.time()
            if now - last_status >= 300:
                save_status()
                last_status = now
        except Exception as e:
            logger.error(f"Loop error: {e}")
        time.sleep(CHECK_INTERVAL)


if __name__ == "__main__":
    main()