279 lines
11 KiB
Python
279 lines
11 KiB
Python
"""workers 按 pid 精确杀执行路径单测：复用判定、只杀可杀档。
|
||
|
||
只用标准库 + 本地 SQLite 临时库，不联网、不杀任何真实存活
|
||
进程、不碰任何容器：可杀档的“真杀”一律用注入的 kill_fn
|
||
测试替身（记录调用，不发真实信号）；默认杀路径只对一个
|
||
已退出并 wait 过的子进程 pid（或等价的已死 pid）跑，且该
|
||
pid 在判定与执行时均经 pid_is_dead 确认为已死。
|
||
不可杀档（异机、无 pid、pid 存活、未失联、恰等于阈值）
|
||
只列不杀：kill_fn 绝不收到它们的 pid，杀审计也绝不含它们。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud.workers import (
|
||
classify_stale_workers,
|
||
reap_stale_workers,
|
||
register_worker,
|
||
)
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.workers import (
|
||
classify_stale_workers,
|
||
reap_stale_workers,
|
||
register_worker,
|
||
)
|
||
|
||
|
||
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)
|
||
|
||
|
||
def _audit_rows(db_path):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
return conn.execute(
|
||
"SELECT action, detail FROM audit_log ORDER BY id").fetchall()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
class PidKillReapTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.now = time.time()
|
||
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)
|
||
self.dead = {9001, 9002, 9003}
|
||
self.events = []
|
||
self.kills = []
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _checker(self, pid):
|
||
return pid in self.dead
|
||
|
||
def _kill_fn(self, pid):
|
||
self.events.append(("kill", pid))
|
||
self.kills.append(pid)
|
||
|
||
def _audit_fn(self, action, detail):
|
||
self.events.append(("audit", action, detail))
|
||
|
||
def _reap(self):
|
||
return reap_stale_workers(
|
||
self.db_path, stale_seconds=120, local_host="h-local",
|
||
pid_is_dead_fn=self._checker, kill_fn=self._kill_fn,
|
||
audit_fn=self._audit_fn, now=self.now)
|
||
|
||
def test_only_kill_row_is_killed_and_audited_before_kill(self):
|
||
rows = self._reap()
|
||
# 真杀只发生在唯一可杀档，且杀前已有审计（事件顺序为证）。
|
||
self.assertEqual(self.kills, [9001])
|
||
kill_idx = self.events.index(("kill", 9001))
|
||
audit_idx = next(
|
||
i for i, e in enumerate(self.events)
|
||
if e[0] == "audit" and e[1] == "worker_pid_kill")
|
||
self.assertLess(audit_idx, kill_idx)
|
||
row = _by_id(rows, "w-kill")
|
||
self.assertTrue(row["executed"])
|
||
self.assertEqual(row["outcome"], "killed")
|
||
self.assertTrue(row["verified_dead"])
|
||
self.assertTrue(row["audit_before"])
|
||
self.assertTrue(row["audit_after"])
|
||
self.assertIsNone(row["kill_error"])
|
||
# 杀后审计在位，detail 含 outcome。
|
||
after = [e for e in self.events
|
||
if e[0] == "audit"
|
||
and e[1] == "worker_pid_kill_verified"]
|
||
self.assertEqual(len(after), 1)
|
||
self.assertIn("outcome=killed", after[0][2])
|
||
self.assertIn("worker_id=w-kill", after[0][2])
|
||
|
||
def test_non_kill_rows_are_listed_only(self):
|
||
rows = self._reap()
|
||
for worker_id in ("w-eq", "w-foreign", "w-nopid", "w-alive",
|
||
"w-fresh"):
|
||
row = _by_id(rows, worker_id)
|
||
self.assertFalse(row["executed"], worker_id)
|
||
self.assertEqual(row["outcome"], "listed", worker_id)
|
||
self.assertEqual(row["action"], "list_only", worker_id)
|
||
# 审计 detail 里绝不出现不可杀档的 worker_id。
|
||
details = " ".join(e[2] for e in self.events if e[0] == "audit")
|
||
for worker_id in ("w-eq", "w-foreign", "w-nopid", "w-alive",
|
||
"w-fresh"):
|
||
self.assertNotIn(worker_id, details)
|
||
|
||
def test_default_audit_writes_audit_log_rows(self):
|
||
# 不注 audit_fn 时默认写 audit_log 表：杀前与杀后各一条，
|
||
# 且只含 w-kill。
|
||
rows = reap_stale_workers(
|
||
self.db_path, stale_seconds=120, local_host="h-local",
|
||
pid_is_dead_fn=self._checker, kill_fn=self._kill_fn,
|
||
now=self.now)
|
||
self.assertEqual(_by_id(rows, "w-kill")["outcome"], "killed")
|
||
audit = _audit_rows(self.db_path)
|
||
actions = [a for a, _ in audit]
|
||
self.assertEqual(actions, ["worker_pid_kill",
|
||
"worker_pid_kill_verified"])
|
||
self.assertIn("worker_id=w-kill", audit[0][1])
|
||
self.assertIn("outcome=killed", audit[1][1])
|
||
|
||
def test_pid_alive_at_execution_is_skipped_not_killed(self):
|
||
# 判定后 pid 被复用（checker 后续调用变 False）：执行时
|
||
# 再验发现已不死，跳过不杀——判定快照不能当执行凭证。
|
||
# 按 pid 计次：9001 首次调用（判定内）为 True，此后
|
||
# 一律 False；其它 pid 按 dead 集判定（worker 行按
|
||
# worker_id 排序，w-alive 会先于 w-kill 调 checker，
|
||
# 全局计次会在此误伤判定）。
|
||
counts = {}
|
||
|
||
def flip_checker(pid):
|
||
counts[pid] = counts.get(pid, 0) + 1
|
||
if pid == 9001 and counts[pid] > 1:
|
||
return False
|
||
return pid in self.dead
|
||
|
||
rows = reap_stale_workers(
|
||
self.db_path, stale_seconds=120, local_host="h-local",
|
||
pid_is_dead_fn=flip_checker, kill_fn=self._kill_fn,
|
||
audit_fn=self._audit_fn, now=self.now)
|
||
self.assertEqual(self.kills, [])
|
||
row = _by_id(rows, "w-kill")
|
||
self.assertFalse(row["executed"])
|
||
self.assertEqual(row["outcome"],
|
||
"skipped_pid_alive_at_execution")
|
||
self.assertFalse(row["audit_before"])
|
||
|
||
def test_audit_failure_blocks_kill(self):
|
||
def failing_audit(action, detail):
|
||
raise RuntimeError("audit store down")
|
||
|
||
rows = reap_stale_workers(
|
||
self.db_path, stale_seconds=120, local_host="h-local",
|
||
pid_is_dead_fn=self._checker, kill_fn=self._kill_fn,
|
||
audit_fn=failing_audit, now=self.now)
|
||
self.assertEqual(self.kills, [])
|
||
row = _by_id(rows, "w-kill")
|
||
self.assertFalse(row["executed"])
|
||
self.assertEqual(row["outcome"], "audit_failed")
|
||
self.assertFalse(row["audit_before"])
|
||
|
||
def test_kill_error_with_dead_pid_still_killed(self):
|
||
# kill_fn 抛非 Lookup 的 OSError，但复核已死：成败以
|
||
# 复核为准，kill_error 如实记录，不谎报失败。
|
||
def boom_kill(pid):
|
||
self.kills.append(pid)
|
||
raise OSError("boom")
|
||
|
||
rows = reap_stale_workers(
|
||
self.db_path, stale_seconds=120, local_host="h-local",
|
||
pid_is_dead_fn=self._checker, kill_fn=boom_kill,
|
||
audit_fn=self._audit_fn, now=self.now)
|
||
row = _by_id(rows, "w-kill")
|
||
self.assertEqual(row["outcome"], "killed")
|
||
self.assertTrue(row["verified_dead"])
|
||
self.assertEqual(row["kill_error"], "boom")
|
||
|
||
def test_kill_error_with_live_pid_is_failed(self):
|
||
# 判定 True、执行再验 True、杀后复核 False：杀失败，
|
||
# outcome 如实为 kill_failed_pid_alive。
|
||
seq = {"9001": [True, True, False]}
|
||
|
||
def seq_checker(pid):
|
||
values = seq.get(str(pid))
|
||
if values:
|
||
return values.pop(0)
|
||
return pid in self.dead
|
||
|
||
rows = reap_stale_workers(
|
||
self.db_path, stale_seconds=120, local_host="h-local",
|
||
pid_is_dead_fn=seq_checker, kill_fn=self._kill_fn,
|
||
audit_fn=self._audit_fn, now=self.now)
|
||
row = _by_id(rows, "w-kill")
|
||
self.assertEqual(row["outcome"], "kill_failed_pid_alive")
|
||
self.assertFalse(row["verified_dead"])
|
||
self.assertEqual(self.kills, [9001])
|
||
|
||
def test_reap_reuses_classification(self):
|
||
# 复用证据：同一数据、同一 checker 下，reap 的判定字段
|
||
# 与直接调 classify_stale_workers 完全一致，不另造判定。
|
||
classified = classify_stale_workers(
|
||
self.db_path, stale_seconds=120, local_host="h-local",
|
||
pid_is_dead_fn=self._checker, now=self.now)
|
||
rows = self._reap()
|
||
for worker_id in ("w-kill", "w-eq", "w-foreign", "w-nopid",
|
||
"w-alive", "w-fresh"):
|
||
self.assertEqual(_by_id(rows, worker_id)["action"],
|
||
_by_id(classified, worker_id)["action"],
|
||
worker_id)
|
||
self.assertEqual(_by_id(rows, worker_id)["reason"],
|
||
_by_id(classified, worker_id)["reason"],
|
||
worker_id)
|
||
|
||
|
||
class DefaultKillDeadPidTestCase(unittest.TestCase):
|
||
def test_default_kill_on_truly_dead_pid(self):
|
||
# 默认杀路径：子进程退出并 wait 后 pid 确实已死，
|
||
# 默认 SIGKILL 遇 ProcessLookupError 当已死处理，
|
||
# 复核 pid_is_dead 为 True，outcome=killed。
|
||
proc = subprocess.Popen(
|
||
[sys.executable, "-c", "pass"],
|
||
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||
proc.wait()
|
||
dead_pid = proc.pid
|
||
tmp = tempfile.TemporaryDirectory()
|
||
try:
|
||
db_path = os.path.join(tmp.name, "control.db")
|
||
register_worker(db_path, "w-dead", host="h-local",
|
||
pid=dead_pid)
|
||
_set_last_seen(db_path, "w-dead", time.time() - 1000)
|
||
rows = reap_stale_workers(
|
||
db_path, stale_seconds=120, local_host="h-local")
|
||
row = _by_id(rows, "w-dead")
|
||
self.assertEqual(row["action"], "kill")
|
||
self.assertEqual(row["outcome"], "killed")
|
||
self.assertTrue(row["verified_dead"])
|
||
self.assertTrue(row["audit_before"])
|
||
self.assertTrue(row["audit_after"])
|
||
finally:
|
||
tmp.cleanup()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|