"""存储接线测试：写穿、读穿、跨进程指令。 这里用**两个 ``SqliteTaskStore`` 连接同一个数据库文件**来模拟两个进程。这不是 过度讲究：跨进程真正难的部分就是"两个连接看不到彼此的内存"，用同一个 store 对象 测是测不出来的。 用例的重点是那些**会静默出错**的地方：写穿漏一次、指令被错误的进程吃掉、指令被 错误地标记完成——它们都不会抛异常，只会让看板上的状态悄悄变成错的。 """ from __future__ import annotations import tempfile import time from datetime import UTC, datetime from pathlib import Path from typing import Any from chongming_task.controller import Controller from chongming_task.errors import StoreError from chongming_task.models import TaskSnapshot from chongming_task.states import TaskStatus from chongming_task.stores.memory import InMemoryTaskStore from chongming_task.stores.protocol import Command, CommandKind from chongming_task.stores.sqlite import SqliteTaskStore _WINDOW = 10.0 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 slow_loop(ctx: Any) -> str: """持续小步推进直到被停止。""" for _ in range(1_000_000): ctx.checkpoint() time.sleep(0.005) return "done" def quick_add(ctx: Any, left: int, right: int) -> int: del ctx return left + right def _shared_store(tmp: Path) -> tuple[SqliteTaskStore, SqliteTaskStore]: """同一个数据库文件上的两个独立连接 = 两个"进程"。""" return ( SqliteTaskStore(tmp / "tasks.db"), SqliteTaskStore(tmp / "tasks.db"), ) # --------------------------------------------------------------------------- # 写穿 # --------------------------------------------------------------------------- def test_initial_snapshot_is_persisted_before_any_transition() -> None: """必须先建行：之后每次迁移都是按版本更新的 UPDATE。 顺序如果反了，第一次 ``save`` 会因"任务不存在"而失败。但注意：写穿失败是被 捕获并记进 ``persist_failures`` 的（存储故障不该让任务失败），所以**光断言行 存在是不够的**——必须同时断言没有任何写穿失败，否则这条用例会在实现有 bug 时 依然通过。 """ store = InMemoryTaskStore() controller = Controller(store=store, max_parallel=1) controller.register("block", slow_loop) controller.submit("block", task_id="t1", label="占住额度") assert _wait_until( lambda: (store.get("t1") or TaskSnapshot("x", "x")).status is TaskStatus.RUNNING ) # 额度已被占满，这个任务停在 queued。 controller.submit("block", task_id="t2", label="排队中") persisted = store.get("t2") assert persisted is not None assert persisted.status is TaskStatus.QUEUED assert persisted.name == "排队中" assert controller.persist_failures == (), "建行与更新的顺序不对，写穿失败了" controller.stop_all() controller.wait_all(timeout=_WINDOW) assert store.get("t2").status is TaskStatus.STOPPED # type: ignore[union-attr] def test_every_transition_is_written_through_and_is_auditable() -> None: store = InMemoryTaskStore() controller = Controller(store=store) controller.register("add", quick_add) controller.submit("add", 2, 3, task_id="t1") controller.wait_all(timeout=_WINDOW) snapshot = store.get("t1") assert snapshot is not None assert snapshot.status is TaskStatus.SUCCEEDED assert snapshot.outcome == {"value": 5} assert snapshot.validate() == () history = store.history("t1") assert [item.to_status for item in history] == [ TaskStatus.QUEUED, TaskStatus.STARTING, TaskStatus.RUNNING, TaskStatus.SUCCEEDED, ] assert [item.version for item in history] == [1, 2, 3, 4] def test_bucket_and_key_survive_the_write_through() -> None: store = InMemoryTaskStore() controller = Controller(store=store) controller.register("add", quick_add) controller.submit("add", 1, 1, task_id="t1", bucket_key="tenant-42") controller.wait_all(timeout=_WINDOW) snapshot = store.get("t1") assert snapshot is not None assert snapshot.bucket == controller.bucket_policy.bucket_for("tenant-42") assert snapshot.bucket_key == "tenant-42" def test_progress_and_heartbeat_reach_the_store_without_bumping_the_version() -> None: """进度走高频旁路：它若参与乐观并发，会与状态写入持续冲突。""" store = InMemoryTaskStore() controller = Controller(store=store, command_poll_interval=0.05) controller.register("block", slow_loop) controller.submit("block", task_id="t1") assert _wait_until( lambda: (store.get("t1") or TaskSnapshot("x", "x")).status is TaskStatus.RUNNING ) before = store.get("t1") assert before is not None version = before.status_version assert _wait_until( lambda: (store.get("t1").last_heartbeat_at if store.get("t1") else None) is not None, timeout=3.0, ) after = store.get("t1") assert after is not None assert after.status_version == version, "心跳回写不应改动版本号" controller.stop("t1") controller.wait_all(timeout=_WINDOW) # --------------------------------------------------------------------------- # 读穿 # --------------------------------------------------------------------------- def test_a_second_controller_sees_tasks_it_does_not_own() -> None: """这就是存储接线存在的全部理由：跨进程可见。""" with tempfile.TemporaryDirectory(prefix="chongming-wire-") as directory: writer_store, reader_store = _shared_store(Path(directory)) try: writer = Controller(store=writer_store) writer.register("add", quick_add) writer.submit("add", 1, 1, task_id="t1", label="跨进程可见") writer.wait_all(timeout=_WINDOW) reader = Controller(store=reader_store) assert reader.runs() == (), "读方没有本地句柄" snapshot = reader.get("t1") assert snapshot is not None assert snapshot.name == "跨进程可见" assert snapshot.status is TaskStatus.SUCCEEDED assert [s.task_id for s in reader.list_tasks()] == ["t1"] assert reader.stats().of(TaskStatus.SUCCEEDED) == 1 writer.close(timeout=_WINDOW) finally: writer_store.close() reader_store.close() def test_reader_filters_are_pushed_down_to_the_store() -> None: with tempfile.TemporaryDirectory(prefix="chongming-wire-") as directory: writer_store, reader_store = _shared_store(Path(directory)) try: writer = Controller(store=writer_store) writer.register("add", quick_add) writer.submit("add", 1, 1, task_id="a1", label="抽取-甲", bucket_key="tenant-A") writer.submit("add", 1, 1, task_id="a2", label="加载", bucket_key="tenant-B") writer.wait_all(timeout=_WINDOW) reader = Controller(store=reader_store) assert [s.task_id for s in reader.list_tasks(name_prefix="抽取")] == ["a1"] assert [ s.task_id for s in reader.list_tasks(bucket=writer.bucket_policy.bucket_for("tenant-B")) ] == ["a2"] assert reader.list_tasks(status=TaskStatus.FAILED) == [] writer.close(timeout=_WINDOW) finally: writer_store.close() reader_store.close() # --------------------------------------------------------------------------- # 跨进程控制 # --------------------------------------------------------------------------- def test_stop_crosses_the_process_boundary_via_a_command() -> None: """一个进程无法给另一个进程的执行载体发信号，只能写指令让持有者取走。""" with tempfile.TemporaryDirectory(prefix="chongming-wire-") as directory: owner_store, client_store = _shared_store(Path(directory)) try: owner = Controller(store=owner_store, command_poll_interval=0.05, owner="owner-process") owner.register("block", slow_loop) owner.submit("block", task_id="t1", label="长任务") assert _wait_until( lambda: ( (owner_store.get("t1") or TaskSnapshot("x", "x")).status is TaskStatus.RUNNING ) ) client = Controller(store=client_store, owner="ops-client") assert client.stop("t1", reason="跨进程停止") is True assert client.runs() == (), "客户端并不持有这个任务" assert _wait_until( lambda: ( (owner_store.get("t1") or TaskSnapshot("x", "x")).status is TaskStatus.STOPPED ) ), "持有者没有取走指令" final = client.get("t1") assert final is not None assert final.status is TaskStatus.STOPPED assert final.failure is not None owner.close(timeout=_WINDOW) finally: owner_store.close() client_store.close() def test_command_records_who_requested_it() -> None: store = InMemoryTaskStore() owner = Controller(store=store, command_poll_interval=0.05, owner="owner-1") owner.register("block", slow_loop) run = owner.submit("block", task_id="t1") assert _wait_until(lambda: run.status is TaskStatus.RUNNING) client = Controller(store=store, owner="ops@laptop") client.stop("t1", reason="发布窗口") assert _wait_until(lambda: run.status is TaskStatus.STOPPED) owner.close(timeout=_WINDOW) def test_pause_of_a_remote_task_reports_unconfirmed() -> None: """跨进程没有回执通道，因此 `pause` 只能如实返回 False（未确认）。""" store = InMemoryTaskStore() owner = Controller(store=store, command_poll_interval=0.05) owner.register("block", slow_loop) run = owner.submit("block", task_id="t1") assert _wait_until(lambda: run.status is TaskStatus.RUNNING) client = Controller(store=store) assert client.pause("t1") is False assert _wait_until(lambda: run.status is TaskStatus.PAUSED), "指令没有被执行" owner.close(timeout=_WINDOW) def test_a_client_that_never_submitted_does_not_steal_commands() -> None: """纯查询客户端若也去取指令，会把别人的停止请求抢走。 即便它把指令交还，持续轮询也会让真正的持有者反复错过——所以轮询线程只在 "提交过任务"之后才启动。 """ store = InMemoryTaskStore() owner = Controller(store=store, command_poll_interval=0.05) owner.register("block", slow_loop) run = owner.submit("block", task_id="t1") assert _wait_until(lambda: run.status is TaskStatus.RUNNING) # 一个只读客户端：只查询，不提交。 client = Controller(store=store) client.list_tasks() client.get("t1") client.stats() assert client._worker is None, "没有提交过任务的实例不该启动指令轮询" owner.stop("t1") assert _wait_until(lambda: run.status is TaskStatus.STOPPED) owner.close(timeout=_WINDOW) def test_a_command_for_a_task_we_do_not_own_is_released_not_completed() -> None: """取到不属于自己的指令必须**交还**，否则真正的持有者永远收不到它。 真实场景：另一个进程提交的任务，它的停止指令被**本**进程的轮询线程取到了—— 本进程没有那个任务的本地句柄。所以这里直接往存储里写"别人的任务"和它的指令， 而不是走 ``Controller.request``（后者会先校验任务存在，见下一条用例）。 """ store = InMemoryTaskStore() controller = Controller(store=store, command_poll_interval=0.05) controller.register("add", quick_add) # 先提交一个自己的任务，让轮询线程启动起来。 controller.submit("add", 1, 1, task_id="mine") controller.wait_all(timeout=_WINDOW) store.create( TaskSnapshot( task_id="someone-elses", name="别人的任务", status=TaskStatus.RUNNING, status_version=1, attempt=1, max_attempts=1, ) ) store.submit_command( Command( command_id="c-foreign", task_id="someone-elses", kind=CommandKind.STOP, created_at=datetime(2026, 1, 1, tzinfo=UTC), requested_by="other-process", ) ) def released() -> bool: found = [item for item in store.commands() if item.command_id == "c-foreign"] return bool(found) and found[0].is_open and found[0].lease_until is None assert _wait_until(released), "指令没有被交还，真正的持有者将永远收不到它" controller.close(timeout=_WINDOW) def test_request_refuses_a_task_that_does_not_exist() -> None: """拼错的 task_id 不该写出一条**永远不会有人执行**的指令。 没有这道校验时，``stop typo`` 会"成功"，于是存储里留下一条孤儿指令，而调用方 还以为停止已经下发。 """ store = InMemoryTaskStore() controller = Controller(store=store) controller.register("add", quick_add) controller.submit("add", 1, 1, task_id="mine") controller.wait_all(timeout=_WINDOW) assert controller.request("typo-task-id", CommandKind.STOP) is False assert store.commands() == [] controller.close(timeout=_WINDOW) def test_request_without_a_store_is_rejected() -> None: """没有共享媒介就不可能跨进程控制，如实返回 False 而不是假装成功。""" controller = Controller() assert controller.request("t1", CommandKind.STOP) is False # --------------------------------------------------------------------------- # 存储故障不该让任务失败 # --------------------------------------------------------------------------- class _FlakyStore(InMemoryTaskStore): """``save`` 总是失败，用来验证存储故障不会打断任务。""" def save(self, snapshot: TaskSnapshot, *, expect_version: int) -> None: raise StoreError("模拟磁盘故障") def test_store_failure_does_not_fail_the_task_but_is_recorded() -> None: """任务本身没问题，坏的只是可观测性。 但也绝不能静默——否则"看板上的状态是旧的"这件事毫无线索。 """ controller = Controller(store=_FlakyStore()) controller.register("add", quick_add) run = controller.submit("add", 2, 3, task_id="t1") outcome = controller.wait("t1", timeout=_WINDOW) assert outcome.ok, "存储故障不该让任务失败" assert outcome.value == 5 assert run.status is TaskStatus.SUCCEEDED failures = controller.persist_failures assert failures, "写穿失败必须留痕" assert "StoreError" in failures[0] assert "t1" in failures[0] def test_persist_failures_are_bounded() -> None: controller = Controller(store=_FlakyStore()) controller.register("add", quick_add) for index in range(30): controller.submit("add", index, 1, task_id=f"t{index}") controller.wait_all(timeout=_WINDOW) assert len(controller.persist_failures) <= 20 # --------------------------------------------------------------------------- # 与分桶、并发额度的协同 # --------------------------------------------------------------------------- def test_bucket_counts_read_through_the_store() -> None: with tempfile.TemporaryDirectory(prefix="chongming-wire-") as directory: writer_store, reader_store = _shared_store(Path(directory)) try: writer = Controller(store=writer_store) writer.register("add", quick_add) writer.submit("add", 1, 1, task_id="a1", bucket_key="tenant-A") writer.submit("add", 1, 1, task_id="a2", bucket_key="tenant-A") writer.submit("add", 1, 1, task_id="b1") writer.wait_all(timeout=_WINDOW) reader = Controller(store=reader_store) counts = reader.bucket_counts() assert counts[writer.bucket_policy.bucket_for("tenant-A")] == 2 assert counts[None] == 1 writer.close(timeout=_WINDOW) finally: writer_store.close() reader_store.close() def test_local_concurrency_limit_still_applies_with_a_store() -> None: """接了存储也不能绕过本地的额度：排队中的任务在存储里应当是 queued。""" store = InMemoryTaskStore() controller = Controller(store=store, max_parallel=1) controller.register("block", slow_loop) first = controller.submit("block", task_id="t1") assert _wait_until(lambda: first.status is TaskStatus.RUNNING) second = controller.submit("block", task_id="t2") assert second.status is TaskStatus.QUEUED assert store.get("t2").status is TaskStatus.QUEUED # type: ignore[union-attr] controller.stop_all() controller.wait_all(timeout=_WINDOW) assert store.get("t2").status is TaskStatus.STOPPED # type: ignore[union-attr]