host
29dcae4513
Worker startup cleanup previously deleted every container whose name started with codex-, including the codex-postgres database container. On the external host the first worker restart after the prod-grade deploy deleted codex-postgres, leaving the control plane unable to connect to Postgres (connection refused on 127.0.0.1:5432, server auto-restart loop). Fix: cleanup only deletes codex-<32-hex task id> containers, which is exactly the name DockerExecutor gives task containers; infrastructure names like codex-postgres never match and are never deleted. Regression test covers postgres + non-hex names.
301 lines
12 KiB
Python
301 lines
12 KiB
Python
"""生产级 worker 超时 / 重试 / 容器清理兜底 + Postgres 备份单测。
|
||
|
||
只用标准库与本地 fake runner / 真实 sleep 进程，不跑模型、不联网；
|
||
Postgres 备份用可注入的 fake pg_dump runner，不调真实 pg_dump，
|
||
密码只经 PGPASSWORD 环境变量传给子进程（单测里不放真实密码）。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud import backup as backup_mod
|
||
from cloud.auth import AuthStore
|
||
from cloud.sandbox import DockerExecutor
|
||
from cloud.worker import Worker
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud import backup as backup_mod
|
||
from cloud.auth import AuthStore
|
||
from cloud.sandbox import DockerExecutor
|
||
from cloud.worker import Worker
|
||
|
||
|
||
class _FakeRunner:
|
||
"""同步 fake 执行器 runner：按 results 顺序返回结果。
|
||
|
||
LocalExecutor 的 runner 收的是命令列表、返回 SimpleNamespace
|
||
（returncode/stdout/stderr），与真实 subprocess 同款接口；
|
||
这里按顺序返回，不按任务 id 查表。
|
||
"""
|
||
|
||
def __init__(self, results):
|
||
self.results = list(results)
|
||
self.calls = []
|
||
|
||
def __call__(self, cmd):
|
||
self.calls.append(list(cmd))
|
||
spec = self.results[min(len(self.calls) - 1,
|
||
len(self.results) - 1)]
|
||
if isinstance(spec, Exception):
|
||
raise spec
|
||
return SimpleNamespace(
|
||
returncode=0 if spec[0] == "completed" else 1,
|
||
stdout=spec[1], stderr="")
|
||
|
||
|
||
def _make_worker(tmp, runner=None, process_factory=None,
|
||
task_timeout=1800, max_retries=1,
|
||
retry_backoff_seconds=0):
|
||
db_path = os.path.join(tmp, "control.db")
|
||
# worker 本身不查项目成员关系：init_db + 直接插任务即可，
|
||
# 不用走 AuthStore（其 create_project 只写成员表、不返项目）。
|
||
from cloud.server import init_db
|
||
init_db(db_path)
|
||
project_id = "proj"
|
||
task_id = "t1"
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||
" depends_on, priority, created_by, created_at, updated_at,"
|
||
" last_heartbeat) VALUES(?, ?, ?, 'queued', '[]', 0,"
|
||
" 'system', 0, 0, NULL)",
|
||
(task_id, project_id, "do it"))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
worker = Worker(db_path, "host.py", config_path=None, repo=None,
|
||
runner=runner,
|
||
env={"XCM_API_KEY": "test-placeholder"},
|
||
env_key=None,
|
||
process_factory=process_factory,
|
||
task_timeout=task_timeout, max_retries=max_retries,
|
||
retry_backoff_seconds=retry_backoff_seconds)
|
||
task = {"id": task_id, "project_id": project_id,
|
||
"prompt": "do it", "depends_on": "[]", "priority": 0}
|
||
return worker, db_path, task
|
||
|
||
|
||
def _task_row(db_path):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
return conn.execute(
|
||
"SELECT status, error_text FROM tasks WHERE id='t1'").fetchone()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _audit_actions(db_path):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
return [r[0] for r in conn.execute(
|
||
"SELECT action FROM audit_log").fetchall()]
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
class WorkerTimeoutTestCase(unittest.TestCase):
|
||
def test_popen_timeout_kills_process_group_and_fails(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
|
||
def factory(task):
|
||
return subprocess.Popen(
|
||
["sleep", "30"], stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE, text=True, start_new_session=True)
|
||
|
||
worker, db_path, task = _make_worker(
|
||
tmp.name, process_factory=factory, task_timeout=0.2,
|
||
max_retries=0, retry_backoff_seconds=0)
|
||
result = worker.run_one(task)
|
||
self.assertEqual(result["status"], "failed")
|
||
# 超时是失败，不是取消：库里必须是 failed，绝不能是 interrupted。
|
||
status, error_text = _task_row(db_path)
|
||
self.assertEqual(status, "failed")
|
||
self.assertIn("timeout", error_text)
|
||
self.assertIn("worker_timeout", _audit_actions(db_path))
|
||
tmp.cleanup()
|
||
|
||
def test_docker_timeout_converts_to_failed_and_audit(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
|
||
def docker_runner(cmd, timeout=None):
|
||
if cmd[1] in ("kill", "rm"):
|
||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||
raise subprocess.TimeoutExpired(cmd, timeout)
|
||
|
||
worker, db_path, task = _make_worker(
|
||
tmp.name, runner=_FakeRunner([]), max_retries=0,
|
||
retry_backoff_seconds=0)
|
||
# 测试占位 key（非真实密钥），只为通过执行器的“key 存在性”
|
||
# 前置检查，让 runner 走到 TimeoutExpired 分支。
|
||
worker.executor = DockerExecutor(
|
||
runner=docker_runner, env={"XCM_API_KEY": "test"},
|
||
timeout=5)
|
||
# VM 上无 docker 二进制：只 mock 存在性检查（_check_docker），
|
||
# 真正的 run 仍走 fake runner 的 TimeoutExpired 分支。
|
||
with mock.patch("cloud.sandbox.shutil.which",
|
||
return_value="/usr/bin/docker"):
|
||
result = worker.run_one(task)
|
||
self.assertEqual(result["status"], "failed")
|
||
status, error_text = _task_row(db_path)
|
||
self.assertEqual(status, "failed")
|
||
self.assertIn("timeout", error_text)
|
||
self.assertIn("worker_timeout", _audit_actions(db_path))
|
||
tmp.cleanup()
|
||
|
||
|
||
class WorkerRetryTestCase(unittest.TestCase):
|
||
def test_generic_failure_retries_once_then_failed(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
runner = _FakeRunner([("failed", "boom"), ("failed", "boom")])
|
||
worker, db_path, task = _make_worker(
|
||
tmp.name, runner=runner, max_retries=1,
|
||
retry_backoff_seconds=0)
|
||
result = worker.run_one(task)
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertEqual(len(runner.calls), 2)
|
||
status, error_text = _task_row(db_path)
|
||
self.assertEqual(status, "failed")
|
||
self.assertIn("[retries=1]", error_text)
|
||
self.assertIn("boom", error_text)
|
||
self.assertIn("worker_retry", _audit_actions(db_path))
|
||
tmp.cleanup()
|
||
|
||
def test_retry_success_recovers(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
runner = _FakeRunner([("failed", "boom"), ("completed", "ok")])
|
||
worker, db_path, task = _make_worker(
|
||
tmp.name, runner=runner, max_retries=1,
|
||
retry_backoff_seconds=0)
|
||
result = worker.run_one(task)
|
||
self.assertEqual(result["status"], "completed")
|
||
self.assertEqual(len(runner.calls), 2)
|
||
status, _ = _task_row(db_path)
|
||
self.assertEqual(status, "completed")
|
||
tmp.cleanup()
|
||
|
||
def test_max_retries_zero_means_no_retry(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
runner = _FakeRunner([("failed", "boom")])
|
||
worker, db_path, task = _make_worker(
|
||
tmp.name, runner=runner, max_retries=0,
|
||
retry_backoff_seconds=0)
|
||
result = worker.run_one(task)
|
||
self.assertEqual(result["status"], "failed")
|
||
self.assertEqual(len(runner.calls), 1)
|
||
_, error_text = _task_row(db_path)
|
||
self.assertNotIn("[retries=", error_text)
|
||
tmp.cleanup()
|
||
|
||
|
||
class ContainerCleanupTestCase(unittest.TestCase):
|
||
def _docker_executor(self, ps_stdout=None, fail_ps=False):
|
||
calls = []
|
||
|
||
def runner(cmd, timeout=None):
|
||
calls.append(list(cmd))
|
||
if cmd[1] == "ps":
|
||
if fail_ps:
|
||
raise RuntimeError("docker down")
|
||
return SimpleNamespace(
|
||
returncode=0, stdout=ps_stdout or "", stderr="")
|
||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||
|
||
executor = DockerExecutor(runner=runner)
|
||
return executor, calls
|
||
|
||
def test_cleanup_only_task_containers_never_infra(self):
|
||
# 回归：2026-10-26 实测，按 codex- 前缀全删会把数据库容器
|
||
# codex-postgres 当成残留任务容器杀掉，导致控制面连不上
|
||
# Postgres。清理只能删 codex-<32 位 hex 任务 id>。
|
||
executor, calls = self._docker_executor(
|
||
ps_stdout="codex-postgres\n"
|
||
"codex-3689cf40d671492992c6bfedd18b004c\n"
|
||
"codex-abc\nother\ncodex-app-server\n")
|
||
cleaned = executor.cleanup_stale_containers()
|
||
self.assertEqual(
|
||
cleaned, ["codex-3689cf40d671492992c6bfedd18b004c"])
|
||
flat = [c for call in calls for c in call]
|
||
self.assertNotIn("codex-postgres", flat)
|
||
self.assertNotIn("codex-abc", flat)
|
||
self.assertNotIn("other", flat)
|
||
# kill + rm 各调一次（仅那个任务容器）。
|
||
self.assertEqual(
|
||
sum(1 for call in calls if call[1] == "kill"), 1)
|
||
self.assertEqual(
|
||
sum(1 for call in calls if call[1] == "rm"), 1)
|
||
|
||
def test_cleanup_failure_does_not_crash_worker(self):
|
||
executor, _ = self._docker_executor(fail_ps=True)
|
||
self.assertEqual(executor.cleanup_stale_containers(), [])
|
||
tmp = tempfile.TemporaryDirectory()
|
||
worker, _, _ = _make_worker(tmp.name, runner=_FakeRunner([]))
|
||
worker.executor = executor
|
||
# worker 层兜底同样不抛异常。
|
||
self.assertEqual(worker.cleanup_stale_containers(), [])
|
||
tmp.cleanup()
|
||
|
||
def test_list_containers_unparsable_or_failed_returns_empty(self):
|
||
executor, _ = self._docker_executor(ps_stdout="")
|
||
self.assertEqual(executor.list_containers(), [])
|
||
|
||
|
||
class PostgresBackupTestCase(unittest.TestCase):
|
||
def test_backup_postgres_fake_runner_success(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
created = {}
|
||
|
||
def fake_runner(cmd, env=None):
|
||
self.assertNotIn("secret-pw", " ".join(cmd))
|
||
self.assertEqual(env.get("PGPASSWORD"), "secret-pw")
|
||
dest = cmd[cmd.index("-f") + 1]
|
||
Path(dest).write_text("-- dummy dump\n")
|
||
created["cmd"] = cmd
|
||
return SimpleNamespace(returncode=0, stdout="", stderr="")
|
||
|
||
env = {"CLOUD_DB_NAME": "codexdb", "CLOUD_DB_USER": "codexsvc",
|
||
"CLOUD_DB_PASSWORD": "secret-pw"}
|
||
with mock.patch.dict(os.environ, env):
|
||
dest = backup_mod.backup_postgres(
|
||
tmp.name, runner=fake_runner)
|
||
self.assertTrue(os.path.exists(dest))
|
||
self.assertGreater(os.path.getsize(dest), 0)
|
||
self.assertNotIn("secret-pw", dest)
|
||
self.assertIn("pg_dump", created["cmd"][0])
|
||
tmp.cleanup()
|
||
|
||
def test_backup_postgres_failure_raises_clear_error(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
|
||
def failing_runner(cmd, env=None):
|
||
return SimpleNamespace(
|
||
returncode=1, stdout="", stderr="pg_dump: boom")
|
||
|
||
env = {"CLOUD_DB_NAME": "codexdb", "CLOUD_DB_USER": "codexsvc"}
|
||
with mock.patch.dict(os.environ, env):
|
||
with self.assertRaises(RuntimeError) as ctx:
|
||
backup_mod.backup_postgres(tmp.name, runner=failing_runner)
|
||
self.assertIn("pg_dump", str(ctx.exception))
|
||
tmp.cleanup()
|
||
|
||
def test_backup_postgres_missing_binary_raises(self):
|
||
tmp = tempfile.TemporaryDirectory()
|
||
env = {"CLOUD_DB_NAME": "codexdb"}
|
||
with mock.patch.dict(os.environ, env):
|
||
with self.assertRaises(RuntimeError):
|
||
backup_mod.backup_postgres(
|
||
tmp.name, pg_dump_cmd="nonexistent-pg-dump-xyz")
|
||
tmp.cleanup()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|