From 8ea54cf66b81a94cf77ef8de7e4ebe341f002359 Mon Sep 17 00:00:00 2001 From: host Date: Fri, 30 Oct 2026 19:19:48 +0000 Subject: [PATCH] workers: land host/pid/address columns with idempotent migration and pid_verified stale classification - 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. --- cloud/server.py | 30 +++- cloud/tests/test_workers_hostpid.py | 243 ++++++++++++++++++++++++++ cloud/worker.py | 14 +- cloud/workers.py | 261 ++++++++++++++++++++++++++-- 4 files changed, 524 insertions(+), 24 deletions(-) create mode 100644 cloud/tests/test_workers_hostpid.py diff --git a/cloud/server.py b/cloud/server.py index 88e1804..07c33a5 100644 --- a/cloud/server.py +++ b/cloud/server.py @@ -8641,9 +8641,12 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): def _workers_heartbeat(self, user_id): """worker 心跳：body 含 worker_id（必填，缺了 400）与可选 - executor/capacity/running。心跳本身就是注册：worker_id - 不存在时自动注册，绝不 404。running 是 worker 上报的 - 当前占用槽数（启动中占位也算）；未给时保留旧值。""" + executor/capacity/running/host/pid/address。心跳本身就是 + 注册：worker_id 不存在时自动注册，绝不 404。running 是 + worker 上报的当前占用槽数（启动中占位也算）；未给时保留 + 旧值。host/pid/address 只存调用方显式给的值：远程 worker + 的 host/pid 绝不由 server 本机推导（server 的主机名/进程 + 号不是该 worker 的），未给时保留旧值。""" body, err = self._read_json() if err: self._send_json(400, {"error": err}) @@ -8670,10 +8673,29 @@ class CloudControlPlaneHandler(BaseHTTPRequestHandler): if running < 0: self._send_json(400, {"error": "running must be >= 0"}) return + host = body.get("host") + if host is not None and not isinstance(host, str): + self._send_json(400, {"error": "host must be a string"}) + return + pid = body.get("pid") + if pid is not None: + try: + pid = int(pid) + except (TypeError, ValueError): + self._send_json(400, {"error": "pid must be an integer"}) + return + if pid <= 0: + self._send_json(400, {"error": "pid must be > 0"}) + return + address = body.get("address") + if address is not None and not isinstance(address, str): + self._send_json(400, {"error": "address must be a string"}) + return try: worker = _registry_heartbeat( self.server.db_path, worker_id, - executor=executor, capacity=capacity, running=running) + executor=executor, capacity=capacity, running=running, + host=host, pid=pid, address=address) except ValueError as exc: self._send_json(400, {"error": str(exc)}) return diff --git a/cloud/tests/test_workers_hostpid.py b/cloud/tests/test_workers_hostpid.py new file mode 100644 index 0000000..cda9859 --- /dev/null +++ b/cloud/tests/test_workers_hostpid.py @@ -0,0 +1,243 @@ +"""workers 表 host/pid/address 落地单测：幂等迁移 + 心跳 + 失联判定。 + +只用标准库 + 本地 SQLite 临时库，不跑真实模型、不联网、不杀 +任何进程或容器：失联判定只测 classify_stale_workers 的返回值， +pid 存活一律用注入的 pid_is_dead_fn，不依赖本机真实进程号； +pid_is_dead 本体只测非法 pid（False）与不存在的大进程号 +（True），不测当前存活进程（环境相关）。 +""" + +import os +import sqlite3 +import sys +import tempfile +import time +import unittest +from pathlib import Path + +try: + from cloud.workers import ( + classify_stale_workers, + ensure_workers_table, + heartbeat_worker, + list_workers, + pid_is_dead, + register_worker, + ) +except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.workers import ( + classify_stale_workers, + ensure_workers_table, + heartbeat_worker, + list_workers, + pid_is_dead, + register_worker, + ) + + +def _cols(db_path): + conn = sqlite3.connect(db_path) + try: + rows = conn.execute("PRAGMA table_info(workers)").fetchall() + return {row[1]: row[2] for row in rows} + finally: + conn.close() + + +def _set_last_seen(db_path, worker_id, last_seen): + conn = sqlite3.connect(db_path) + try: + conn.execute("UPDATE workers SET last_seen = ? WHERE worker_id = ?", + (last_seen, worker_id)) + conn.commit() + finally: + conn.close() + + +def _by_id(rows, worker_id): + for row in rows: + if row["worker_id"] == worker_id: + return row + raise AssertionError("missing worker %s" % worker_id) + + +class HostpidMigrationTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + + def tearDown(self): + self.tmp.cleanup() + + def test_new_db_table_has_host_pid_address(self): + ensure_workers_table(self.db_path) + cols = _cols(self.db_path) + self.assertIn("host", cols) + self.assertIn("pid", cols) + self.assertIn("address", cols) + + def test_migration_is_idempotent(self): + ensure_workers_table(self.db_path) + register_worker(self.db_path, "w1", executor="local", + host="h1", pid=1234, address="10.0.0.1") + ensure_workers_table(self.db_path) + ensure_workers_table(self.db_path) + rows = list_workers(self.db_path) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]["host"], "h1") + self.assertEqual(rows[0]["pid"], 1234) + self.assertEqual(rows[0]["address"], "10.0.0.1") + + def test_old_db_upgrade_keeps_rows_and_leaves_new_cols_null(self): + # 手工建旧表（无 host/pid/address 列、无 running 列外的 + # 旧列集），写一行，再跑迁移：旧行保留、新列为空。 + conn = sqlite3.connect(self.db_path) + try: + conn.execute( + "CREATE TABLE workers(" + "worker_id TEXT PRIMARY KEY, executor TEXT, " + "capacity INTEGER, last_seen DOUBLE PRECISION, " + "created_at DOUBLE PRECISION)") + conn.execute( + "INSERT INTO workers VALUES('w-old', 'docker', 1, ?, ?)", + (1000.0, 900.0)) + conn.commit() + finally: + conn.close() + ensure_workers_table(self.db_path) + ensure_workers_table(self.db_path) + rows = list_workers(self.db_path) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]["executor"], "docker") + self.assertEqual(rows[0]["last_seen"], 1000.0) + self.assertIsNone(rows[0]["host"]) + self.assertIsNone(rows[0]["pid"]) + self.assertIsNone(rows[0]["address"]) + + def test_heartbeat_refreshes_host_pid_and_preserves_when_omitted(self): + register_worker(self.db_path, "w1", executor="local", + host="h-old", pid=111, address="a-old") + row = heartbeat_worker(self.db_path, "w1", host="h-new", pid=222) + self.assertEqual(row["host"], "h-new") + self.assertEqual(row["pid"], 222) + # address 未给：保留旧值，不被洗成 None。 + self.assertEqual(row["address"], "a-old") + row = heartbeat_worker(self.db_path, "w1") + self.assertEqual(row["host"], "h-new") + self.assertEqual(row["pid"], 222) + self.assertEqual(row["address"], "a-old") + + def test_heartbeat_new_row_stores_given_values_and_nulls_when_omitted(self): + row = heartbeat_worker(self.db_path, "w-new", host="h1", pid=333) + self.assertEqual(row["host"], "h1") + self.assertEqual(row["pid"], 333) + self.assertIsNone(row["address"]) + row = heartbeat_worker(self.db_path, "w-plain") + self.assertIsNone(row["host"]) + self.assertIsNone(row["pid"]) + self.assertIsNone(row["address"]) + + def test_invalid_pid_rejected(self): + with self.assertRaises(ValueError): + heartbeat_worker(self.db_path, "w1", pid=0) + with self.assertRaises(ValueError): + heartbeat_worker(self.db_path, "w1", pid="not-a-pid") + with self.assertRaises(ValueError): + heartbeat_worker(self.db_path, "w1", pid=True) + + +class StaleClassificationTestCase(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.db_path = os.path.join(self.tmp.name, "control.db") + self.now = time.time() + # 阈值 120 秒：w-kill 超阈 + 本机 + 已死；w-eq 恰等于阈值； + # w-foreign 超阈但异机；w-nopid 超阈但无 pid；w-alive 超阈 + # 但 pid 还活着；w-fresh 未超阈。 + register_worker(self.db_path, "w-kill", host="h-local", pid=9001) + register_worker(self.db_path, "w-eq", host="h-local", pid=9002) + register_worker(self.db_path, "w-foreign", host="h-other", + pid=9003) + register_worker(self.db_path, "w-nopid", host="h-local") + register_worker(self.db_path, "w-alive", host="h-local", pid=9004) + register_worker(self.db_path, "w-fresh", host="h-local", pid=9005) + _set_last_seen(self.db_path, "w-kill", self.now - 121) + _set_last_seen(self.db_path, "w-eq", self.now - 120) + _set_last_seen(self.db_path, "w-foreign", self.now - 200) + _set_last_seen(self.db_path, "w-nopid", self.now - 200) + _set_last_seen(self.db_path, "w-alive", self.now - 200) + _set_last_seen(self.db_path, "w-fresh", self.now - 10) + + def tearDown(self): + self.tmp.cleanup() + + def _classify(self): + dead = {9001, 9002, 9003} + return classify_stale_workers( + self.db_path, stale_seconds=120, local_host="h-local", + pid_is_dead_fn=lambda pid: pid in dead, now=self.now) + + def test_only_pid_verified_stale_is_kill(self): + rows = self._classify() + self.assertEqual(_by_id(rows, "w-kill")["action"], "kill") + self.assertTrue(_by_id(rows, "w-kill")["pid_verified"]) + self.assertEqual(_by_id(rows, "w-kill")["reason"], + "pid_verified_dead") + + def test_equal_threshold_is_list_only(self): + rows = self._classify() + row = _by_id(rows, "w-eq") + self.assertEqual(row["action"], "list_only") + self.assertFalse(row["pid_verified"]) + self.assertEqual(row["reason"], "age_equals_threshold") + + def test_non_verified_stale_are_list_only(self): + rows = self._classify() + self.assertEqual(_by_id(rows, "w-foreign")["action"], "list_only") + self.assertEqual(_by_id(rows, "w-foreign")["reason"], + "host_mismatch") + self.assertEqual(_by_id(rows, "w-nopid")["action"], "list_only") + self.assertEqual(_by_id(rows, "w-nopid")["reason"], "pid_missing") + self.assertEqual(_by_id(rows, "w-alive")["action"], "list_only") + self.assertEqual(_by_id(rows, "w-alive")["reason"], + "pid_alive_or_unknown") + self.assertEqual(_by_id(rows, "w-fresh")["action"], "list_only") + self.assertEqual(_by_id(rows, "w-fresh")["reason"], "not_stale") + + def test_classification_kills_nothing(self): + # 判定函数本身绝不杀进程：跑完后本进程仍在，且返回值里 + # 除 w-kill 外无 kill 项；函数不接受任何容器名/前缀参数， + # 无法按名删容器（签名即证据，见 classify_stale_workers）。 + rows = self._classify() + kills = [r for r in rows if r["action"] == "kill"] + self.assertEqual([r["worker_id"] for r in kills], ["w-kill"]) + self.assertTrue(pid_is_dead(999999)) + + +class PgBulkSchemaBoundaryTestCase(unittest.TestCase): + def test_workers_stays_out_of_postgres_bulk_schema(self): + # 界线钉住：workers 是易失心跳数据（migrations.py 的 + # UNMIGRATED_TABLES），PG 批量 schema 刻意不含它；PG + # 生产库的 workers 表由 ensure_workers_table 懒建/补列。 + # 本轨曾把 workers 加进 POSTGRES_SCHEMA，致 P2 迁移对齐 + # 3 项变红（迁移表集、未迁移清单、覆盖率），已回退。 + try: + from cloud.db import POSTGRES_SCHEMA + except ImportError: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + from cloud.db import POSTGRES_SCHEMA + joined = "\n".join(POSTGRES_SCHEMA) + self.assertNotIn("CREATE TABLE IF NOT EXISTS workers(", + joined) + + +class PidIsDeadTestCase(unittest.TestCase): + def test_invalid_pid_is_not_dead(self): + self.assertFalse(pid_is_dead(None)) + self.assertFalse(pid_is_dead(0)) + self.assertFalse(pid_is_dead("x")) + + +if __name__ == "__main__": + unittest.main() diff --git a/cloud/worker.py b/cloud/worker.py index 993405a..c73e647 100644 --- a/cloud/worker.py +++ b/cloud/worker.py @@ -699,18 +699,26 @@ class Worker: register_worker( self.db_path, self.worker_id, executor=self._executor_name(), capacity=self.capacity, - running=self._running_count()) + running=self._running_count(), + host=self.host_name, pid=os.getpid()) except Exception as exc: log_event(self.logger, "worker_register_failed", message="worker register failed: %s" % exc) def _registry_heartbeat(self): - """向控制面写注册表心跳（与文件心跳并行，失败只记日志）。""" + """向控制面写注册表心跳（与文件心跳并行，失败只记日志）。 + + host 用 self.host_name（autoscaler 显式传的宿主名；未传 + 时退化为 socket.gethostname()，容器内即容器名，与 + worktree_host 同款 caveat），pid 用本进程号；两者每轮 + 刷新，供失联判定做本机 + pid 已死的 pid_verified 校验。 + """ try: heartbeat_worker( self.db_path, self.worker_id, executor=self._executor_name(), capacity=self.capacity, - running=self._running_count()) + running=self._running_count(), + host=self.host_name, pid=os.getpid()) except Exception as exc: log_event(self.logger, "worker_registry_heartbeat_failed", message="worker registry heartbeat failed: %s" % exc) diff --git a/cloud/workers.py b/cloud/workers.py index a3df77c..ccf209e 100644 --- a/cloud/workers.py +++ b/cloud/workers.py @@ -17,15 +17,40 @@ 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 秒。""" @@ -63,11 +88,15 @@ def ensure_workers_table(db_path): "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 下这张表不存在，查询失败即跳过， @@ -151,32 +180,217 @@ def _ensure_running_column(conn): 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): + running=None, host=None, pid=None, address=None): """注册（或刷新）一个 worker：不存在 INSERT，存在 UPDATE。 返回注册行 dict。worker_id 为空时抛 ValueError，调用方 - （server 的 heartbeat 端点）据此返回 400。 + （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: - 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: + 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, last_seen, created_at)" - " VALUES(?, ?, ?, ?, ?, ?)", + " running, host, pid, address, last_seen, created_at)" + " VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)", (worker_id, executor, capacity, - running if running is not None else 0, now, now), + running if running is not None else 0, + host, pid, address, now, now), ) conn.commit() row = conn.execute( @@ -188,14 +402,19 @@ def register_worker(db_path, worker_id, executor=None, capacity=1, def heartbeat_worker(db_path, worker_id, executor=None, capacity=None, - running=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)) @@ -206,19 +425,24 @@ def heartbeat_worker(db_path, worker_id, executor=None, capacity=None, if row is None: conn.execute( "INSERT INTO workers(worker_id, executor, capacity," - " running, last_seen, created_at)" - " VALUES(?, ?, ?, ?, ?, ?)", + " 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, now, now), + 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 = ?, last_seen = ? WHERE worker_id = ?", + " 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() @@ -263,6 +487,9 @@ def _worker_to_dict(row): "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"), }