host
e2d6a454bf
SQLite REAL is 8-byte double, but Postgres REAL is 4-byte float4 (~7 significant digits). Epoch seconds (10 digits) were rounded to 128-second steps, so worker liveness ages were off by up to +-64s and could even go negative (observed live on the host: last_seen ahead of NOW() by 20s). New tables use DOUBLE PRECISION; existing Postgres tables are upgraded only when information_schema actually reports REAL, so the ALTER does not run on every heartbeat (it would take an ACCESS EXCLUSIVE lock each time). SQLite path is unchanged (its REAL is already double precision). Regression test asserts new tables use DOUBLE PRECISION.
207 lines
7.7 KiB
Python
207 lines
7.7 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, "
|
||
"last_seen DOUBLE PRECISION, "
|
||
"created_at DOUBLE PRECISION)"
|
||
)
|
||
conn.commit()
|
||
# 老库升精度只在确实是 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 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"),
|
||
}
|