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.
244 lines
9.7 KiB
Python
244 lines
9.7 KiB
Python
"""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()
|