codex-app-server/cloud/workers.py

"""Worker 注册表：多 worker 池的注册、心跳与存活判断。

零第三方依赖，只用 Python 标准库。

设计：
- 注册表存在控制面数据库的 workers 表里（worker 与 server 共用
  同一个库，所以 worker 直接写库注册，不需要为此再发 HTTP；
  HTTP 侧的 GET /workers 与 POST /workers/heartbeat 只是同一张
  表的 API 面，给有 API key 的外部 worker 用）。
- worker_id 默认由调用方给（生产 worker 应配稳定的 --worker-id）；
  未给时由 worker.py 用“主机名-进程号”生成——重启后进程号会变，
  旧行会自然变 stale（alive=false），不要用随机 uuid 让注册表
  每次重启都多一行无法对应到同一台机器。
- 心跳本身就是注册：heartbeat 时 worker_id 不存在就 INSERT，
  存在就 UPDATE，绝不因“未注册”返回 404 报错。
- alive 定义：last_seen 在阈值内（默认 120 秒，可经环境变量
  CLOUD_WORKER_ALIVE_SECONDS 调，测试可传参覆盖）。
- 本模块不存、不打任何密钥：注册信息只有 worker_id、executor、
  capacity、last_seen，没有任何凭证字段。
"""

import os
import time

from cloud.db import connect as _db_connect

DEFAULT_ALIVE_SECONDS = 120


def alive_threshold_seconds(override=None):
    """存活阈值：显式参数 > 环境变量 > 默认 120 秒。"""
    if override is not None:
        try:
            return float(override)
        except (TypeError, ValueError):
            return float(DEFAULT_ALIVE_SECONDS)
    try:
        return float(os.environ.get(
            "CLOUD_WORKER_ALIVE_SECONDS") or DEFAULT_ALIVE_SECONDS)
    except (TypeError, ValueError):
        return float(DEFAULT_ALIVE_SECONDS)


def ensure_workers_table(db_path):
    """建 workers 表（不存在时建，已存在时不动）。"""
    conn = _db_connect(str(db_path))
    try:
        conn.execute(
            "CREATE TABLE IF NOT EXISTS workers("
            "worker_id TEXT PRIMARY KEY, "
            "executor TEXT, "
            "capacity INTEGER, "
            "last_seen REAL, "
            "created_at REAL)"
        )
        conn.commit()
    finally:
        conn.close()


def register_worker(db_path, worker_id, executor=None, capacity=1):
    """注册（或刷新）一个 worker：不存在 INSERT，存在 UPDATE。

    返回注册行 dict。worker_id 为空时抛 ValueError，调用方
    （server 的 heartbeat 端点）据此返回 400。
    """
    if not worker_id or not isinstance(worker_id, str):
        raise ValueError("worker_id required")
    ensure_workers_table(db_path)
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        cur = conn.execute(
            "UPDATE workers SET executor = ?, capacity = ?,"
            " last_seen = ? WHERE worker_id = ?",
            (executor, capacity, now, worker_id),
        )
        if cur.rowcount == 0:
            conn.execute(
                "INSERT INTO workers(worker_id, executor, capacity,"
                " last_seen, created_at) VALUES(?, ?, ?, ?, ?)",
                (worker_id, executor, capacity, now, now),
            )
        conn.commit()
        row = conn.execute(
            "SELECT * FROM workers WHERE worker_id = ?",
            (worker_id,)).fetchone()
    finally:
        conn.close()
    return _worker_to_dict(row)


def heartbeat_worker(db_path, worker_id, executor=None, capacity=None):
    """心跳：就是带 last_seen 刷新的注册。

    executor/capacity 未给时保留旧值（已注册的 worker 只刷时间，
    不把 executor 洗成 None）；未注册时按给定值（或默认）注册。
    """
    if not worker_id or not isinstance(worker_id, str):
        raise ValueError("worker_id required")
    ensure_workers_table(db_path)
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        row = conn.execute(
            "SELECT * FROM workers WHERE worker_id = ?",
            (worker_id,)).fetchone()
        if row is None:
            conn.execute(
                "INSERT INTO workers(worker_id, executor, capacity,"
                " last_seen, created_at) VALUES(?, ?, ?, ?, ?)",
                (worker_id, executor, capacity if capacity is not None
                 else 1, now, now),
            )
        else:
            old = _worker_to_dict(row)
            conn.execute(
                "UPDATE workers SET executor = ?, capacity = ?,"
                " last_seen = ? WHERE worker_id = ?",
                (executor if executor is not None else old["executor"],
                 capacity if capacity is not None else old["capacity"],
                 now, worker_id),
            )
        conn.commit()
        row = conn.execute(
            "SELECT * FROM workers WHERE worker_id = ?",
            (worker_id,)).fetchone()
    finally:
        conn.close()
    return _worker_to_dict(row)


def list_workers(db_path, alive_seconds=None):
    """列出全部 worker，每行含 alive（last_seen 在阈值内）。"""
    threshold = alive_threshold_seconds(alive_seconds)
    ensure_workers_table(db_path)
    now = time.time()
    conn = _db_connect(str(db_path))
    try:
        rows = conn.execute(
            "SELECT * FROM workers ORDER BY worker_id").fetchall()
    finally:
        conn.close()
    out = []
    for row in rows:
        item = _worker_to_dict(row)
        last = item.get("last_seen")
        item["alive"] = bool(
            last is not None and (now - last) <= threshold)
        out.append(item)
    return out


def _worker_to_dict(row):
    def _get(name):
        try:
            return row[name]
        except (KeyError, IndexError, TypeError):
            return None
    return {
        "worker_id": _get("worker_id"),
        "executor": _get("executor"),
        "capacity": _get("capacity"),
        "last_seen": _get("last_seen"),
        "created_at": _get("created_at"),
    }