456 lines
17 KiB
Python
456 lines
17 KiB
Python
|
"""存储接线测试：写穿、读穿、跨进程指令。
|
||
|
|
|
||
|
|
这里用**两个 ``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]
|