host
8ea54cf66b
- ensure_workers_table: new tables include host TEXT, pid BIGINT, address TEXT; old tables gain them via _ensure_hostpid_columns (PRAGMA column check, missing-only ALTER, old rows keep NULL, never fabricates host/pid). SQLite and Postgres share this path (PG via the db adapter's PRAGMA translation); POSTGRES_SCHEMA deliberately still excludes workers, which is volatile heartbeat data in UNMIGRATED_TABLES (adding it broke the P2 migration alignment tests, verified and reverted). - heartbeat/register: given host/pid/address refresh, omitted values preserve the old row; address has no derived source. Local worker passes host_name + own pid; HTTP heartbeat only stores caller-supplied values, never the server's own. - classify_stale_workers: kill only when age strictly exceeds the threshold AND host matches local host AND pid is proven dead; equal threshold, foreign host, missing pid, or live/unknown pid are list_only. Classification only, never kills, never matches by name/prefix, never touches infra containers.
496 lines
19 KiB
Python
496 lines
19 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，没有任何凭证字段。
|
||
- host/pid/address（本轨新增）：新库建表直接含三列；老库由
|
||
ensure_workers_table 幂等补列，已有行新列为空、不编造值。
|
||
两路界线：SQLite 文件库与 PG 生产库共用 ensure_workers_table
|
||
这一份迁移代码（PG 下 connect() 返回适配连接，PRAGMA 经
|
||
适配层翻译为 information_schema 查询）；PG 批量建库
|
||
（POSTGRES_SCHEMA / init_postgres_schema）刻意不含 workers
|
||
表——workers 是易失心跳数据，见 migrations.py 的
|
||
UNMIGRATED_TABLES，把它加进批量 schema 会撞 P2 迁移对齐
|
||
（MIGRATION_TABLES 必须等于 schema 表集合，本轨实测 3 项
|
||
变红后已回退），不要再加。
|
||
心跳刷新 host/pid 时未给保留旧值；address 没有新来源，只存
|
||
调用方显式给的值，未给保留旧值。本模块绝不按 worker 名或
|
||
前缀枚举/删除容器：失联清理只做判定（见
|
||
classify_stale_workers），判定为 kill 的唯一可杀对象是该行
|
||
的 pid 本人，且必须是本机 + pid 已死（pid_verified）。
|
||
"""
|
||
|
||
import os
|
||
import socket
|
||
import time
|
||
|
||
from cloud.db import connect as _db_connect
|
||
|
||
DEFAULT_ALIVE_SECONDS = 120
|
||
|
||
# workers 表新增列：(列名, 类型)。pid 用 BIGINT：SQLite 下为
|
||
# INTEGER 亲和性，Postgres 下为 8 字节整数，均可容纳进程号；
|
||
# host/address 为 TEXT，不存任何凭证。
|
||
WORKER_HOSTPID_COLUMNS = (
|
||
("host", "TEXT"),
|
||
("pid", "BIGINT"),
|
||
("address", "TEXT"),
|
||
)
|
||
|
||
|
||
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, "
|
||
"host TEXT, "
|
||
"pid BIGINT, "
|
||
"address TEXT, "
|
||
"last_seen DOUBLE PRECISION, "
|
||
"created_at DOUBLE PRECISION)"
|
||
)
|
||
conn.commit()
|
||
_ensure_running_column(conn)
|
||
_ensure_hostpid_columns(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 _ensure_hostpid_columns(conn):
|
||
"""老库补 host/pid/address 列（新表建表时已含三列）。
|
||
|
||
幂等保证：先查现有列名，只补缺失列，多次跑是 no-op；已有
|
||
行新列保持 NULL，不编造 host/pid/address。SQLite 与 PG
|
||
共用本函数：SQLite 走 PRAGMA table_info，Postgres 走适配
|
||
层翻译后的同款 PRAGMA（见 cloud/db.py 的 PRAGMA 翻译，调用
|
||
方取 row[1] 即列名，与 init_postgres_schema 补列同款）；
|
||
PRAGMA 查不到列时再走 information_schema 逐列 COUNT 兜底
|
||
（Postgres 语法 ADD COLUMN IF NOT EXISTS，SQLite 下
|
||
PRAGMA 本就可用，不会走到这条兜底）。补列失败只回滚、
|
||
不抛，绝不炸注册与心跳。PG 批量建库路径不适用本函数：
|
||
POSTGRES_SCHEMA 刻意不含 workers（易失心跳数据，见模块
|
||
头注释），PG 生产库的 workers 表同样由本函数懒建/补列。
|
||
"""
|
||
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:
|
||
for col, col_type in WORKER_HOSTPID_COLUMNS:
|
||
if col not in names:
|
||
conn.execute(
|
||
"ALTER TABLE workers ADD COLUMN %s %s"
|
||
% (col, col_type))
|
||
conn.commit()
|
||
return
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
for col, col_type in WORKER_HOSTPID_COLUMNS:
|
||
try:
|
||
row = conn.execute(
|
||
"SELECT COUNT(*) FROM information_schema.columns "
|
||
"WHERE table_name = 'workers' AND column_name = '%s'"
|
||
% col).fetchone()
|
||
exists = bool(row and row[0])
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
continue
|
||
if not exists:
|
||
try:
|
||
conn.execute(
|
||
"ALTER TABLE workers ADD COLUMN IF NOT EXISTS "
|
||
"%s %s" % (col, col_type))
|
||
conn.commit()
|
||
except Exception:
|
||
try:
|
||
conn.rollback()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _normalize_pid(pid):
|
||
"""归一进程号：None -> None；bool 或不可转 int -> ValueError。
|
||
|
||
pid 列存整数，不存字符串形式的进程号；调用方传错类型时
|
||
明确报错，不静默写成 0 或把旧值洗掉。
|
||
"""
|
||
if pid is None:
|
||
return None
|
||
if isinstance(pid, bool):
|
||
raise ValueError("pid must be an integer")
|
||
try:
|
||
value = int(pid)
|
||
except (TypeError, ValueError):
|
||
raise ValueError("pid must be an integer")
|
||
if value <= 0:
|
||
raise ValueError("pid must be a positive integer")
|
||
return value
|
||
|
||
|
||
def pid_is_dead(pid):
|
||
"""本机进程是否明确已死：os.kill(pid, 0) 的语义。
|
||
|
||
ProcessLookupError -> 已死（True）；PermissionError -> 进程
|
||
存在但无权发信号，不能证明已死（False）；其它 OSError /
|
||
非法 pid -> 不能证明已死（False）。绝不把“查不到/无权”
|
||
当成已死，这是 pid_verified 不误杀的底线。
|
||
"""
|
||
try:
|
||
value = _normalize_pid(pid)
|
||
except ValueError:
|
||
return False
|
||
if value is None:
|
||
return False
|
||
try:
|
||
os.kill(value, 0)
|
||
except ProcessLookupError:
|
||
return True
|
||
except PermissionError:
|
||
return False
|
||
except OSError:
|
||
return False
|
||
return False
|
||
|
||
|
||
def classify_stale_workers(db_path, stale_seconds=None, local_host=None,
|
||
pid_is_dead_fn=None, now=None):
|
||
"""失联 worker 判定（只判定、不杀任何进程或容器）。
|
||
|
||
返回全部 worker 行，每行含 action 与 reason：
|
||
- action == "kill" 当且仅当同时满足：age 严格大于阈值
|
||
（恰等于阈值不杀）且 pid_verified（行 host 等于本机 host、
|
||
pid 为已知正整数、且 pid_is_dead_fn(pid) 为 True）。
|
||
- 其余一律 action == "list_only"：未失联、恰等于阈值、
|
||
host 不是本机、pid 缺失、pid 还活着或无法证明已死，
|
||
都只列不杀。
|
||
- 本函数不枚举容器、不按 worker_id 前缀匹配、不碰任何
|
||
基础设施容器：kill 项的唯一可杀对象就是该行的 pid 本人，
|
||
调用方若要清理也必须按 pid 精确杀，绝不按名/前缀删。
|
||
- stale_seconds 未给时复用 alive 阈值语义（显式参数 >
|
||
环境变量 > 120 秒）；local_host 未给时为本机主机名；
|
||
pid_is_dead_fn 未给时为本机 pid_is_dead（测试可注入）。
|
||
"""
|
||
threshold = alive_threshold_seconds(stale_seconds)
|
||
host = local_host if local_host is not None else socket.gethostname()
|
||
checker = pid_is_dead_fn if pid_is_dead_fn is not None else pid_is_dead
|
||
moment = time.time() if now is None else float(now)
|
||
ensure_workers_table(db_path)
|
||
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")
|
||
age = (moment - last) if last is not None else None
|
||
stale = bool(age is not None and age > threshold)
|
||
verified = False
|
||
if stale:
|
||
row_pid = item.get("pid")
|
||
if (item.get("host") == host and row_pid is not None
|
||
and checker(row_pid)):
|
||
verified = True
|
||
if verified:
|
||
action, reason = "kill", "pid_verified_dead"
|
||
elif age is None:
|
||
action, reason = "list_only", "no_last_seen"
|
||
elif age == threshold:
|
||
action, reason = "list_only", "age_equals_threshold"
|
||
elif age <= threshold:
|
||
action, reason = "list_only", "not_stale"
|
||
elif item.get("host") != host:
|
||
action, reason = "list_only", "host_mismatch"
|
||
elif item.get("pid") is None:
|
||
action, reason = "list_only", "pid_missing"
|
||
else:
|
||
action, reason = "list_only", "pid_alive_or_unknown"
|
||
item["age_seconds"] = age
|
||
item["stale"] = stale
|
||
item["pid_verified"] = verified
|
||
item["action"] = action
|
||
item["reason"] = reason
|
||
out.append(item)
|
||
return out
|
||
|
||
|
||
def register_worker(db_path, worker_id, executor=None, capacity=1,
|
||
running=None, host=None, pid=None, address=None):
|
||
"""注册（或刷新）一个 worker：不存在 INSERT，存在 UPDATE。
|
||
|
||
返回注册行 dict。worker_id 为空时抛 ValueError，调用方
|
||
（server 的 heartbeat 端点）据此返回 400。新列语义：新行
|
||
按给定值写（未给为 NULL）；已存在行上给定值覆盖、未给
|
||
保留旧值——绝不因一次不带 host/pid 的注册把已有的身份
|
||
洗成 NULL，也不编造 host/pid/address。
|
||
"""
|
||
if not worker_id or not isinstance(worker_id, str):
|
||
raise ValueError("worker_id required")
|
||
pid = _normalize_pid(pid)
|
||
ensure_workers_table(db_path)
|
||
now = time.time()
|
||
conn = _db_connect(str(db_path))
|
||
try:
|
||
old_row = conn.execute(
|
||
"SELECT * FROM workers WHERE worker_id = ?",
|
||
(worker_id,)).fetchone()
|
||
if old_row is not None:
|
||
old = _worker_to_dict(old_row)
|
||
conn.execute(
|
||
"UPDATE workers SET executor = ?, capacity = ?,"
|
||
" running = ?, host = ?, pid = ?, address = ?,"
|
||
" last_seen = ? WHERE worker_id = ?",
|
||
(executor, capacity,
|
||
running if running is not None else 0,
|
||
host if host is not None else old["host"],
|
||
pid if pid is not None else old["pid"],
|
||
address if address is not None else old["address"],
|
||
now, worker_id),
|
||
)
|
||
else:
|
||
conn.execute(
|
||
"INSERT INTO workers(worker_id, executor, capacity,"
|
||
" running, host, pid, address, last_seen, created_at)"
|
||
" VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||
(worker_id, executor, capacity,
|
||
running if running is not None else 0,
|
||
host, pid, address, 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, host=None, pid=None, address=None):
|
||
"""心跳：就是带 last_seen 刷新的注册。
|
||
|
||
executor/capacity 未给时保留旧值（已注册的 worker 只刷时间，
|
||
不把 executor 洗成 None）；未注册时按给定值（或默认）注册。
|
||
host/pid/address 同款保留语义：给定则刷新（本机 worker 心跳
|
||
应传本机 host 与本进程 pid，见 worker._registry_heartbeat），
|
||
未给则保留旧值；address 没有新来源，绝不由本函数推导
|
||
（例如不从 worker_id 或 socket 猜地址），只存显式给的值。
|
||
"""
|
||
if not worker_id or not isinstance(worker_id, str):
|
||
raise ValueError("worker_id required")
|
||
pid = _normalize_pid(pid)
|
||
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, host, pid, address, last_seen, created_at)"
|
||
" VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
||
(worker_id, executor, capacity if capacity is not None
|
||
else 1, running if running is not None else 0,
|
||
host, pid, address, now, now),
|
||
)
|
||
else:
|
||
old = _worker_to_dict(row)
|
||
conn.execute(
|
||
"UPDATE workers SET executor = ?, capacity = ?,"
|
||
" running = ?, host = ?, pid = ?, address = ?,"
|
||
" 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"],
|
||
host if host is not None else old["host"],
|
||
pid if pid is not None else old["pid"],
|
||
address if address is not None else old["address"],
|
||
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,
|
||
"host": _get("host"),
|
||
"pid": _get("pid"),
|
||
"address": _get("address"),
|
||
"last_seen": _get("last_seen"),
|
||
"created_at": _get("created_at"),
|
||
}
|