"""重启恢复(``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"