405 lines
14 KiB
Python
405 lines
14 KiB
Python
|
"""重启恢复（``reconcile``）测试。
|
|||
|
|
|
|||
|
|
这一块最容易出错的地方是**误判**：把一个活得好好的任务判成丢失。所以用例的重点不是
|
|||
|
|
"能标记僵尸"，而是"**不**标记活着的任务"：
|
|||
|
|
|
|||
|
|
* 有监督者（租约新鲜）→ 一动不动；
|
|||
|
|
* 长时间不调 ``checkpoint()`` 的健康任务 → 同样一动不动（心跳停滞 ≠ 进程死了）；
|
|||
|
|
* 还在排队的任务 → 保住状态，只报告；
|
|||
|
|
* 原来的监督者又活过来并抢先写入 → 放手，绝不覆盖。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import tempfile
|
|||
|
|
import time
|
|||
|
|
from datetime import UTC, datetime, timedelta
|
|||
|
|
from pathlib import Path
|
|||
|
|
from typing import Any
|
|||
|
|
|
|||
|
|
from chongming_task.controller import (
|
|||
|
|
Controller,
|
|||
|
|
ReconcileOutcome,
|
|||
|
|
ReconcileReport,
|
|||
|
|
)
|
|||
|
|
from chongming_task.errors import ConfigurationError
|
|||
|
|
from chongming_task.models import FailureKind, TaskSnapshot
|
|||
|
|
from chongming_task.states import TaskStatus
|
|||
|
|
from chongming_task.stores.memory import InMemoryTaskStore
|
|||
|
|
from chongming_task.stores.sqlite import SqliteTaskStore
|
|||
|
|
|
|||
|
|
_WINDOW = 10.0
|
|||
|
|
_OLD = datetime(2026, 1, 1, tzinfo=UTC)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _wait_until(predicate: Any, timeout: float = _WINDOW) -> bool:
|
|||
|
|
deadline = time.monotonic() + timeout
|
|||
|
|
while time.monotonic() < deadline:
|
|||
|
|
if predicate():
|
|||
|
|
return True
|
|||
|
|
time.sleep(0.01)
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _orphan(
|
|||
|
|
task_id: str,
|
|||
|
|
status: TaskStatus,
|
|||
|
|
*,
|
|||
|
|
lease_until: datetime | None = None,
|
|||
|
|
lease_owner: str | None = None,
|
|||
|
|
started_at: datetime | None = _OLD,
|
|||
|
|
) -> TaskSnapshot:
|
|||
|
|
"""构造一个"上次进程留下的"任务行。"""
|
|||
|
|
return TaskSnapshot(
|
|||
|
|
task_id=task_id,
|
|||
|
|
name="遗留任务",
|
|||
|
|
status=status,
|
|||
|
|
status_version=3,
|
|||
|
|
attempt=1,
|
|||
|
|
max_attempts=1,
|
|||
|
|
created_at=_OLD,
|
|||
|
|
started_at=started_at,
|
|||
|
|
last_heartbeat_at=_OLD,
|
|||
|
|
lease_owner=lease_owner,
|
|||
|
|
lease_until=lease_until,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 基本语义
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_reconcile_requires_a_store() -> None:
|
|||
|
|
"""没有共享存储就无从判断谁还在管——如实报错，不要猜。"""
|
|||
|
|
controller = Controller()
|
|||
|
|
try:
|
|||
|
|
controller.reconcile()
|
|||
|
|
except ConfigurationError as exc:
|
|||
|
|
assert "共享存储" in str(exc)
|
|||
|
|
return
|
|||
|
|
raise AssertionError("没有存储时 reconcile 应当报错")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_stale_running_task_is_marked_lost() -> None:
|
|||
|
|
"""占着执行资源却没有监督者：载体已随进程消失。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("zombie", TaskStatus.RUNNING))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert [item.task_id for item in report.lost] == ["zombie"]
|
|||
|
|
|
|||
|
|
restored = store.get("zombie")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.LOST
|
|||
|
|
assert restored.failure is not None
|
|||
|
|
assert restored.failure.kind is FailureKind.LOST
|
|||
|
|
assert "失联" in restored.failure.message
|
|||
|
|
assert restored.ended_at is not None, "终态必须带 ended_at"
|
|||
|
|
assert restored.validate() == ()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_reconcile_records_an_audit_trail() -> None:
|
|||
|
|
"""判丢失必须留下可追溯的迁移记录，而不是悄悄改一个字段。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("zombie", TaskStatus.RUNNING))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
controller.reconcile(reason="机房断电")
|
|||
|
|
|
|||
|
|
history = store.history("zombie")
|
|||
|
|
assert [item.to_status for item in history] == [TaskStatus.LOST]
|
|||
|
|
assert history[0].from_status is TaskStatus.RUNNING
|
|||
|
|
assert history[0].reason == "机房断电"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_all_resource_holding_states_are_reconciled() -> None:
|
|||
|
|
"""``starting`` / ``running`` / ``pausing`` / ``stopping`` 都占着载体。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
statuses = (
|
|||
|
|
TaskStatus.STARTING,
|
|||
|
|
TaskStatus.RUNNING,
|
|||
|
|
TaskStatus.PAUSING,
|
|||
|
|
TaskStatus.STOPPING,
|
|||
|
|
)
|
|||
|
|
for index, status in enumerate(statuses):
|
|||
|
|
store.create(_orphan(f"z{index}", status))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert len(report.lost) == len(statuses)
|
|||
|
|
for index in range(len(statuses)):
|
|||
|
|
restored = store.get(f"z{index}")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.LOST
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_queued_task_is_reported_but_keeps_its_status() -> None:
|
|||
|
|
"""不占资源的任务只是"没人接手"，不是"丢了"——改状态会把它的排队事实抹掉。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("waiting", TaskStatus.QUEUED))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert report.lost == ()
|
|||
|
|
assert [item.task_id for item in report.adoptable] == ["waiting"]
|
|||
|
|
assert report.adoptable[0].previous_status is TaskStatus.QUEUED
|
|||
|
|
|
|||
|
|
restored = store.get("waiting")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.QUEUED, "不该被改动"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_paused_task_is_adoptable_not_lost() -> None:
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("paused-task", TaskStatus.PAUSED))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert report.lost == ()
|
|||
|
|
assert len(report.adoptable) == 1
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 不该动的东西（这部分的用例更重要）
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_a_task_with_a_fresh_lease_is_never_touched() -> None:
|
|||
|
|
"""有活着的监督者就必须放手。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(
|
|||
|
|
_orphan(
|
|||
|
|
"alive",
|
|||
|
|
TaskStatus.RUNNING,
|
|||
|
|
lease_owner="worker-1",
|
|||
|
|
lease_until=datetime.now(UTC) + timedelta(seconds=60),
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert report.entries == ()
|
|||
|
|
restored = store.get("alive")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.RUNNING
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_a_healthy_task_that_never_checkpoints_is_not_killed() -> None:
|
|||
|
|
"""**这是本模块最重要的一条。**
|
|||
|
|
|
|||
|
|
心跳由任务自己的 ``checkpoint()`` 驱动。一个长时间不调 checkpoint 的健康任务
|
|||
|
|
（纯 CPU 计算）心跳同样会停滞；若拿心跳判断"进程是否还活着"，它会被误判失联并
|
|||
|
|
判 ``LOST``——把一个正常任务杀掉。
|
|||
|
|
|
|||
|
|
判据必须是租约。这里心跳很旧、但租约新鲜，所以一动不动。
|
|||
|
|
"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
snapshot = _orphan(
|
|||
|
|
"cpu-bound",
|
|||
|
|
TaskStatus.RUNNING,
|
|||
|
|
lease_owner="worker-1",
|
|||
|
|
lease_until=datetime.now(UTC) + timedelta(seconds=60),
|
|||
|
|
)
|
|||
|
|
assert snapshot.last_heartbeat_at == _OLD, "前提：心跳确实很旧"
|
|||
|
|
store.create(snapshot)
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
assert controller.reconcile().entries == ()
|
|||
|
|
restored = store.get("cpu-bound")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.RUNNING
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_a_just_started_task_gets_a_grace_period() -> None:
|
|||
|
|
"""刚提交、工作线程还没轮到第一次续租：不能立刻判它失联。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("fresh", TaskStatus.RUNNING, started_at=datetime.now(UTC)))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert report.entries == (), "宽限期内不该动它"
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_a_just_started_task_becomes_reconcilable_after_the_grace_period() -> None:
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("fresh", TaskStatus.RUNNING, started_at=datetime.now(UTC)))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile(stale_after=0.0)
|
|||
|
|
assert [item.task_id for item in report.lost] == ["fresh"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_tasks_owned_by_this_controller_are_skipped() -> None:
|
|||
|
|
"""自己正在跑的任务当然活着——不靠租约也该放过它。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
controller = Controller(store=store, command_poll_interval=0.05)
|
|||
|
|
controller.register("block", _block)
|
|||
|
|
run = controller.submit("block", task_id="mine")
|
|||
|
|
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
|
|||
|
|
|
|||
|
|
report = controller.reconcile(stale_after=0.0)
|
|||
|
|
assert report.lost == ()
|
|||
|
|
assert run.status is TaskStatus.RUNNING
|
|||
|
|
|
|||
|
|
controller.stop("mine")
|
|||
|
|
controller.wait_all(timeout=_WINDOW)
|
|||
|
|
controller.close(timeout=_WINDOW)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_a_terminal_task_is_not_examined() -> None:
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("done", TaskStatus.SUCCEEDED))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
assert controller.reconcile().entries == ()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_reconcile_lets_go_when_someone_writes_first() -> None:
|
|||
|
|
"""原本的监督者又活过来并抢先写入 → 必须放手，强行覆盖会抹掉对方的状态。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("racy", TaskStatus.RUNNING, lease_owner="worker-1"))
|
|||
|
|
|
|||
|
|
# 模拟"对方在我们读取之后、写入之前提交了一次迁移"。
|
|||
|
|
original_query = store.query
|
|||
|
|
|
|||
|
|
def query_with_interference(query: Any = None) -> list[TaskSnapshot]:
|
|||
|
|
found = original_query(query)
|
|||
|
|
current = store.get("racy")
|
|||
|
|
if current is not None and current.status is TaskStatus.RUNNING:
|
|||
|
|
store.save(
|
|||
|
|
TaskSnapshot(
|
|||
|
|
task_id="racy",
|
|||
|
|
name="遗留任务",
|
|||
|
|
status=TaskStatus.SUCCEEDED,
|
|||
|
|
status_version=current.status_version + 1,
|
|||
|
|
attempt=1,
|
|||
|
|
max_attempts=1,
|
|||
|
|
created_at=_OLD,
|
|||
|
|
started_at=_OLD,
|
|||
|
|
ended_at=datetime.now(UTC),
|
|||
|
|
),
|
|||
|
|
expect_version=current.status_version,
|
|||
|
|
)
|
|||
|
|
return found
|
|||
|
|
|
|||
|
|
store.query = query_with_interference # type: ignore[method-assign]
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
|
|||
|
|
assert [item.outcome for item in report.skipped] == [ReconcileOutcome.SKIPPED]
|
|||
|
|
restored = store.get("racy")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.SUCCEEDED, "对方的状态被覆盖了"
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 预演与报告
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_dry_run_reports_without_writing() -> None:
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("zombie", TaskStatus.RUNNING))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile(dry_run=True)
|
|||
|
|
assert report.dry_run is True
|
|||
|
|
assert [item.task_id for item in report.lost] == ["zombie"]
|
|||
|
|
|
|||
|
|
restored = store.get("zombie")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.RUNNING, "预演不该写任何状态"
|
|||
|
|
assert store.history("zombie") == []
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_report_properties_partition_the_entries() -> None:
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("zombie", TaskStatus.RUNNING))
|
|||
|
|
store.create(_orphan("waiting", TaskStatus.QUEUED))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert len(report.lost) == 1
|
|||
|
|
assert len(report.adoptable) == 1
|
|||
|
|
assert report.skipped == ()
|
|||
|
|
assert bool(report) is True
|
|||
|
|
assert ReconcileReport().lost == ()
|
|||
|
|
assert bool(ReconcileReport()) is False
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_reconcile_is_safe_to_run_twice() -> None:
|
|||
|
|
"""第二次已经没有僵尸了（第一个已落终态）。"""
|
|||
|
|
store = InMemoryTaskStore()
|
|||
|
|
store.create(_orphan("zombie", TaskStatus.RUNNING))
|
|||
|
|
controller = Controller(store=store)
|
|||
|
|
|
|||
|
|
assert len(controller.reconcile().lost) == 1
|
|||
|
|
assert controller.reconcile().entries == ()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# 与真实进程配合
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_leases_are_renewed_while_a_task_runs() -> None:
|
|||
|
|
"""监督者在跑就必须持续续租，否则它自己会被别人判成僵尸。"""
|
|||
|
|
with tempfile.TemporaryDirectory(prefix="chongming-reconcile-") as directory:
|
|||
|
|
store = SqliteTaskStore(Path(directory) / "tasks.db")
|
|||
|
|
controller = Controller(
|
|||
|
|
store=store,
|
|||
|
|
command_poll_interval=0.05,
|
|||
|
|
lease_ttl=30.0,
|
|||
|
|
owner="worker-1",
|
|||
|
|
)
|
|||
|
|
controller.register("block", _block)
|
|||
|
|
run = controller.submit("block", task_id="t1")
|
|||
|
|
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
|
|||
|
|
|
|||
|
|
assert _wait_until(
|
|||
|
|
lambda: (store.get("t1") or TaskSnapshot("x", "x")).lease_owner == "worker-1"
|
|||
|
|
), "监督者没有续租"
|
|||
|
|
leased = store.get("t1")
|
|||
|
|
assert leased is not None
|
|||
|
|
assert leased.lease_until is not None
|
|||
|
|
assert leased.lease_until > datetime.now(UTC)
|
|||
|
|
|
|||
|
|
controller.stop("t1")
|
|||
|
|
controller.wait_all(timeout=_WINDOW)
|
|||
|
|
controller.close(timeout=_WINDOW, close_store=True)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_a_reconcile_run_after_a_crash_cleans_up(tmp_path: Path) -> None:
|
|||
|
|
"""端到端：worker 崩了（连接关闭、租约过期），新进程启动时判它丢失。"""
|
|||
|
|
path = tmp_path / "tasks.db"
|
|||
|
|
store = SqliteTaskStore(path)
|
|||
|
|
store.create(
|
|||
|
|
_orphan(
|
|||
|
|
"crashed",
|
|||
|
|
TaskStatus.RUNNING,
|
|||
|
|
lease_owner="worker-1",
|
|||
|
|
lease_until=datetime.now(UTC) - timedelta(seconds=5),
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
store.close()
|
|||
|
|
|
|||
|
|
# 新进程：另一个连接，同一个文件。
|
|||
|
|
reopened = SqliteTaskStore(path)
|
|||
|
|
try:
|
|||
|
|
controller = Controller(store=reopened, owner="worker-2")
|
|||
|
|
report = controller.reconcile()
|
|||
|
|
assert [item.task_id for item in report.lost] == ["crashed"]
|
|||
|
|
|
|||
|
|
restored = reopened.get("crashed")
|
|||
|
|
assert restored is not None
|
|||
|
|
assert restored.status is TaskStatus.LOST
|
|||
|
|
assert controller.persist_failures == ()
|
|||
|
|
finally:
|
|||
|
|
reopened.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _block(ctx: Any) -> str:
|
|||
|
|
for _ in range(1_000_000):
|
|||
|
|
ctx.checkpoint()
|
|||
|
|
time.sleep(0.005)
|
|||
|
|
return "done"
|