269 lines
9.8 KiB
Python
269 lines
9.8 KiB
Python
"""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 表（不存在时建，已存在时不动，老库补精度）。
|
||
|
||
注意：last_seen/created_at 必须是 DOUBLE PRECISION，不能用 REAL。
|
||
SQLite 的 REAL 是 8 字节双精度，没问题；但 Postgres 的 REAL 是
|
||
4 字节 float4，只能存约 7 位有效数字，epoch 秒（10 位）会被
|
||
舍入到 128 秒一档，导致存活判断误差可达 ±64 秒、甚至出现负 age
|
||
（2026-10-26 主机实测：last_seen 比 NOW() 还大 20 秒）。
|
||
新表直接用 DOUBLE PRECISION；老库（已是 REAL 列）用 ALTER
|
||
TYPE 升到 DOUBLE PRECISION（float4 -> float8 无损，已存的舍入值
|
||
无法恢复，但之后的心跳不再舍入）。SQLite 下 ALTER TYPE 语法
|
||
不支持，靠 try/except 跳过——SQLite 的 REAL 本来就是双精度，
|
||
不需要升。
|
||
"""
|
||
conn = _db_connect(str(db_path))
|
||
try:
|
||
conn.execute(
|
||
"CREATE TABLE IF NOT EXISTS workers("
|
||
"worker_id TEXT PRIMARY KEY, "
|
||
"executor TEXT, "
|
||
"capacity INTEGER, "
|
||
"running INTEGER, "
|
||
"last_seen DOUBLE PRECISION, "
|
||
"created_at DOUBLE PRECISION)"
|
||
)
|
||
conn.commit()
|
||
_ensure_running_column(conn)
|
||
# 老库升精度只在确实是 Postgres REAL 列时做一次：每次心跳都
|
||
# ALTER 会反复拿 ACCESS EXCLUSIVE 锁，没必要。先查
|
||
# information_schema（SQLite 下这张表不存在，查询失败即跳过，
|
||
# SQLite 的 REAL 本来就是双精度，不需要升）。
|
||
for col in ("last_seen", "created_at"):
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT data_type FROM information_schema.columns "
|
||
"WHERE table_name = 'workers' AND column_name = ?",
|
||
(col,)).fetchone()
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
continue
|
||
if not row or row[0] != "real":
|
||
continue
|
||
try:
|
||
conn.execute(
|
||
"ALTER TABLE workers ALTER COLUMN %s "
|
||
"TYPE DOUBLE PRECISION" % col)
|
||
conn.commit()
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _ensure_running_column(conn):
|
||
"""老库补 running 列（新表建表时已含该列）。
|
||
|
||
running 是 worker 上次心跳上报的当前占用槽数（启动中占位
|
||
也算，见 cloud/caps.py）；列缺失时 _worker_to_dict 默认 0，
|
||
所以补列失败只回滚、不抛，绝不炸注册与心跳。
|
||
SQLite 用 PRAGMA 查列，Postgres 用 information_schema。
|
||
"""
|
||
try:
|
||
rows = conn.execute("PRAGMA table_info(workers)").fetchall()
|
||
names = set()
|
||
for row in rows:
|
||
try:
|
||
names.add(row[1])
|
||
except (IndexError, TypeError):
|
||
pass
|
||
if rows and "running" not in names:
|
||
conn.execute(
|
||
"ALTER TABLE workers ADD COLUMN running INTEGER DEFAULT 0")
|
||
conn.commit()
|
||
return
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT COUNT(*) FROM information_schema.columns "
|
||
"WHERE table_name = 'workers' AND column_name = 'running'"
|
||
).fetchone()
|
||
exists = bool(row and row[0])
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
return
|
||
if not exists:
|
||
try:
|
||
conn.execute(
|
||
"ALTER TABLE workers ADD COLUMN IF NOT EXISTS "
|
||
"running INTEGER DEFAULT 0")
|
||
conn.commit()
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def register_worker(db_path, worker_id, executor=None, capacity=1,
|
||
running=None):
|
||
"""注册（或刷新）一个 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 = ?,"
|
||
" running = ?, last_seen = ? WHERE worker_id = ?",
|
||
(executor, capacity, running if running is not None else 0,
|
||
now, worker_id),
|
||
)
|
||
if cur.rowcount == 0:
|
||
conn.execute(
|
||
"INSERT INTO workers(worker_id, executor, capacity,"
|
||
" running, last_seen, created_at)"
|
||
" VALUES(?, ?, ?, ?, ?, ?)",
|
||
(worker_id, executor, capacity,
|
||
running if running is not None else 0, 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,
|
||
running=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,"
|
||
" running, last_seen, created_at)"
|
||
" VALUES(?, ?, ?, ?, ?, ?)",
|
||
(worker_id, executor, capacity if capacity is not None
|
||
else 1, running if running is not None else 0, now, now),
|
||
)
|
||
else:
|
||
old = _worker_to_dict(row)
|
||
conn.execute(
|
||
"UPDATE workers SET executor = ?, capacity = ?,"
|
||
" running = ?, last_seen = ? WHERE worker_id = ?",
|
||
(executor if executor is not None else old["executor"],
|
||
capacity if capacity is not None else old["capacity"],
|
||
running if running is not None else old["running"],
|
||
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
|
||
running = _get("running")
|
||
return {
|
||
"worker_id": _get("worker_id"),
|
||
"executor": _get("executor"),
|
||
"capacity": _get("capacity"),
|
||
"running": running if running is not None else 0,
|
||
"last_seen": _get("last_seen"),
|
||
"created_at": _get("created_at"),
|
||
}
|