chakcy
f4de43c743
单任务受控执行（15 态状态机 + 审计轨迹）、DAG 管道、Controller 门面。
执行器三个：线程（默认，含 async def 任务）、内联（同步 / 测试）、进程
（spawn 子进程：硬性停止、绕开 GIL、崩溃隔离）。
调度：并发额度 + 每桶限流 + 具名资源信号量（按标签）+ 优先级排队。
持久层：TaskStore 契约 + 内存 / SQLite / Postgres 实现，契约测试随包发布。
跨进程与多节点：写穿 / 读穿、指令通道与租约、owned_buckets 划界、reconcile。
CLI、12 个可直接运行的示例、docs/ 下 5 份文档。
门禁：ruff / ruff format / mypy(strict) 干净；
pytest 563 passed（44 skipped 是需要真实 Postgres 的那批）；
带一台可连的 Postgres 时 607 passed。
455 lines
18 KiB
Python
455 lines
18 KiB
Python
"""换掉默认持久层的**可执行范例**：一个 JSON 文件后端。
|
||
|
||
这不是内置实现，而是"第三方后端长什么样"的活样板。它只做三件事：
|
||
|
||
1. 实现 :class:`~chongming_task.stores.protocol.TaskStore` 契约；
|
||
2. 复用 :mod:`chongming_task.stores.codec` 做序列化，不自己重写 ``datetime`` /
|
||
``frozenset`` / 枚举 / 嵌套结构的往返；
|
||
3. 跑通完整的 ``StoreContract``。
|
||
|
||
**刻意的局限，也是给后端作者的提醒**：它只用一把进程内锁，**没有文件锁**，因此仅在
|
||
单进程内正确。真正的多进程后端必须在 ``claim_commands`` / ``release_command`` 上做
|
||
跨进程原子性——SQLite 用 ``BEGIN IMMEDIATE``，Redis 用 Lua 脚本或 ``WATCH/MULTI``。
|
||
本范例跑通契约不代表它多进程安全，而契约测试本身也不测跨进程。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import json
|
||
import os
|
||
import tempfile
|
||
import threading
|
||
import time
|
||
from collections.abc import Iterator, Mapping
|
||
from dataclasses import replace
|
||
from datetime import UTC, datetime, timedelta
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from chongming_task.controller import Controller
|
||
from chongming_task.errors import ConcurrencyConflict, StoreError
|
||
from chongming_task.models import Progress, TaskSnapshot
|
||
from chongming_task.states import TERMINAL, TaskStatus
|
||
from chongming_task.stores.codec import (
|
||
command_to_row,
|
||
encode_dt,
|
||
encode_progress,
|
||
row_to_command,
|
||
row_to_snapshot,
|
||
row_to_transition,
|
||
snapshot_to_row,
|
||
transition_to_row,
|
||
)
|
||
from chongming_task.stores.protocol import Command, TaskQuery, TaskStore
|
||
from chongming_task.testing.store_contract import StoreContract
|
||
|
||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||
|
||
|
||
#: ``Path.replace`` 在 Windows 上被瞬时占用时的重试次数。
|
||
_REPLACE_ATTEMPTS = 5
|
||
|
||
|
||
class JsonFileTaskStore:
|
||
"""把一个 JSON 文件当作持久层。
|
||
|
||
单进程正确，跨进程不安全（见模块文档）。写它是为了让"换持久层"这件事有一个
|
||
可抄、可跑、可对照的骨架。
|
||
"""
|
||
|
||
def __init__(self, path: str | Path) -> None:
|
||
self._path = Path(path)
|
||
self._lock = threading.RLock()
|
||
self._data: dict[str, Any] = {"tasks": {}, "transitions": {}, "commands": {}}
|
||
|
||
# --- 内部 ---
|
||
|
||
def _reload(self) -> None:
|
||
"""每次操作前重新读盘。
|
||
|
||
**本实现刻意不缓存。** 缓存意味着两个连接各持一份副本，彼此永远看不到对方的
|
||
写入，于是跨进程控制会**静默失效**——本范例第一次跑跨进程用例时就是这样：
|
||
控制端把指令写进了自己的副本，worker 端压根不知道。这是文件/无共享后端最
|
||
容易踩的坑。
|
||
|
||
真要缓存，就必须配文件锁 + 失效判断；比起那个复杂度，每次都读反而更诚实。
|
||
"""
|
||
if not self._path.exists():
|
||
self._data = {"tasks": {}, "transitions": {}, "commands": {}}
|
||
return
|
||
text = self._path.read_text(encoding="utf-8")
|
||
self._data = (
|
||
json.loads(text) if text.strip() else {"tasks": {}, "transitions": {}, "commands": {}}
|
||
)
|
||
|
||
@contextlib.contextmanager
|
||
def _transaction(self) -> Iterator[None]:
|
||
"""所有公开方法的入口：加锁 → 读盘 → 交给调用方改 → 由调用方 ``_flush``。"""
|
||
with self._lock:
|
||
self._reload()
|
||
yield
|
||
|
||
def _flush(self) -> None:
|
||
"""原子替换写回。
|
||
|
||
不能直接 ``write_text``：它先截断再写，另一个进程恰好在这一瞬间打开文件就会
|
||
读到**空文件或半截 JSON**——本范例第一次跑跨进程用例时踩到的就是这个，错误是
|
||
``JSONDecodeError: Expecting value: line 1 column 1``。``Path.replace`` 在
|
||
同一文件系统上是原子的，读方只会看到完整内容。
|
||
|
||
临时文件名带上 pid：两个进程共用一个临时文件会互相覆盖。
|
||
|
||
Windows 上还有第三种失败，属于**瞬时**：刚写完的文件可能被杀毒扫描或索引临时
|
||
占住，``replace`` 抛 ``PermissionError``（WinError 5，实测约每十轮跑批命中一次）。
|
||
这里重试几次就够了。注意它只解决**瞬时**占用——两个实例持续共用一个文件的冲突
|
||
照旧无解（重试耗尽后还是抛），那正是下面 ``test_the_naive_file_store_is_not_
|
||
cross_process_safe`` 钉住的事实。
|
||
"""
|
||
temporary = self._path.with_name(f"{self._path.name}.{os.getpid()}.tmp")
|
||
temporary.write_text(
|
||
json.dumps(self._data, ensure_ascii=False, sort_keys=True), encoding="utf-8"
|
||
)
|
||
for attempt in range(_REPLACE_ATTEMPTS):
|
||
try:
|
||
temporary.replace(self._path)
|
||
return
|
||
except PermissionError:
|
||
if attempt == _REPLACE_ATTEMPTS - 1:
|
||
raise
|
||
time.sleep(0.02 * (attempt + 1))
|
||
|
||
def _rows(self) -> dict[str, dict[str, Any]]:
|
||
return self._data["tasks"] # type: ignore[no-any-return]
|
||
|
||
# --- 1. 任务快照 ---
|
||
|
||
def create(self, snapshot: TaskSnapshot) -> None:
|
||
with self._transaction():
|
||
if snapshot.task_id in self._rows():
|
||
raise StoreError(f"任务已存在：{snapshot.task_id}")
|
||
self._rows()[snapshot.task_id] = snapshot_to_row(snapshot)
|
||
self._data["transitions"].setdefault(snapshot.task_id, {})
|
||
self._flush()
|
||
|
||
def save(self, snapshot: TaskSnapshot, *, expect_version: int) -> None:
|
||
with self._transaction():
|
||
current = self._rows().get(snapshot.task_id)
|
||
if current is None:
|
||
raise StoreError(f"任务不存在：{snapshot.task_id}")
|
||
if current["status_version"] != expect_version:
|
||
raise ConcurrencyConflict(
|
||
snapshot.task_id, expect_version, current["status_version"]
|
||
)
|
||
row = snapshot_to_row(snapshot)
|
||
# 租约由 touch 独占：save 必须保留它，否则重启恢复会把活着的任务误判失联。
|
||
row["lease_owner"] = current["lease_owner"]
|
||
row["lease_until"] = current["lease_until"]
|
||
self._rows()[snapshot.task_id] = row
|
||
self._flush()
|
||
|
||
def get(self, task_id: str) -> TaskSnapshot | None:
|
||
with self._transaction():
|
||
row = self._rows().get(task_id)
|
||
return row_to_snapshot(row) if row is not None else None
|
||
|
||
def query(self, query: TaskQuery | None = None) -> list[TaskSnapshot]:
|
||
condition = query if query is not None else TaskQuery()
|
||
with self._transaction():
|
||
found = [row_to_snapshot(row) for row in self._rows().values()]
|
||
|
||
if condition.statuses is not None:
|
||
found = [item for item in found if item.status in condition.statuses]
|
||
if condition.name is not None:
|
||
found = [item for item in found if item.name == condition.name]
|
||
if condition.name_prefix is not None:
|
||
found = [item for item in found if item.name.startswith(condition.name_prefix)]
|
||
if condition.tag is not None:
|
||
found = [item for item in found if condition.tag in item.tags]
|
||
if condition.buckets is not None:
|
||
found = [item for item in found if item.bucket in condition.buckets]
|
||
if condition.pipeline_run_id is not None:
|
||
found = [item for item in found if item.pipeline_run_id == condition.pipeline_run_id]
|
||
|
||
found.sort(
|
||
key=lambda item: (item.created_at is None, item.created_at or datetime.min),
|
||
reverse=condition.newest_first,
|
||
)
|
||
return found[: condition.limit] if condition.limit is not None else found
|
||
|
||
def counts(self) -> dict[TaskStatus, int]:
|
||
with self._transaction():
|
||
result: dict[TaskStatus, int] = {}
|
||
for row in self._rows().values():
|
||
status = TaskStatus(row["status"])
|
||
result[status] = result.get(status, 0) + 1
|
||
return result
|
||
|
||
# --- 2. 审计轨迹 ---
|
||
|
||
def append(self, transition: Any) -> None:
|
||
with self._transaction():
|
||
entries = self._data["transitions"].setdefault(transition.task_id, {})
|
||
# 幂等：同一版本重复写入被忽略（崩溃恢复会重放）。
|
||
entries.setdefault(str(transition.version), transition_to_row(transition))
|
||
self._flush()
|
||
|
||
def history(self, task_id: str, *, since_version: int = 0) -> list[Any]:
|
||
with self._transaction():
|
||
entries: Mapping[str, Any] = self._data["transitions"].get(task_id, {})
|
||
rows = [row for version, row in entries.items() if int(version) > since_version]
|
||
rows.sort(key=lambda row: row["version"])
|
||
return [row_to_transition(row) for row in rows]
|
||
|
||
# --- 3. 高频旁路 ---
|
||
|
||
def touch(
|
||
self,
|
||
task_id: str,
|
||
*,
|
||
progress: Progress | None = None,
|
||
heartbeat_at: datetime | None = None,
|
||
lease_owner: str | None = None,
|
||
lease_until: datetime | None = None,
|
||
) -> None:
|
||
with self._transaction():
|
||
row = self._rows().get(task_id)
|
||
if row is None:
|
||
return # 心跳晚到一步不该报错
|
||
if progress is not None:
|
||
# 刻意不动 status_version：否则会与状态写入持续冲突。
|
||
row["progress"] = encode_progress(progress)
|
||
if heartbeat_at is not None:
|
||
row["last_heartbeat_at"] = encode_dt(heartbeat_at)
|
||
if lease_owner is not None:
|
||
row["lease_owner"] = lease_owner
|
||
if lease_until is not None:
|
||
row["lease_until"] = encode_dt(lease_until)
|
||
self._flush()
|
||
|
||
# --- 4. 控制指令 ---
|
||
|
||
def submit_command(self, command: Command) -> None:
|
||
with self._transaction():
|
||
if command.command_id in self._data["commands"]:
|
||
raise StoreError(f"指令已存在：{command.command_id}")
|
||
self._data["commands"][command.command_id] = command_to_row(command)
|
||
self._flush()
|
||
|
||
def claim_commands(
|
||
self, *, now: datetime, lease_seconds: float, limit: int = 50
|
||
) -> list[Command]:
|
||
with self._transaction():
|
||
claimable = [
|
||
row_to_command(row)
|
||
for row in self._data["commands"].values()
|
||
if row["finished_at"] is None
|
||
and (row["lease_until"] is None or row["lease_until"] <= encode_dt(now))
|
||
]
|
||
claimable.sort(key=lambda item: item.created_at)
|
||
claimed: list[Command] = []
|
||
for item in claimable[:limit]:
|
||
leased = replace(item, lease_until=now + timedelta(seconds=lease_seconds))
|
||
self._data["commands"][item.command_id] = command_to_row(leased)
|
||
claimed.append(leased)
|
||
self._flush()
|
||
return claimed
|
||
|
||
def complete_command(
|
||
self, command_id: str, *, now: datetime, outcome: str | None = None
|
||
) -> None:
|
||
with self._transaction():
|
||
row = self._data["commands"].get(command_id)
|
||
if row is None:
|
||
return
|
||
row["finished_at"] = encode_dt(now)
|
||
row["outcome"] = outcome
|
||
row["lease_until"] = None
|
||
self._flush()
|
||
|
||
def release_command(self, command_id: str, *, now: datetime) -> None:
|
||
del now
|
||
with self._transaction():
|
||
row = self._data["commands"].get(command_id)
|
||
if row is None or row["finished_at"] is not None:
|
||
return
|
||
row["lease_until"] = None
|
||
self._flush()
|
||
|
||
# --- 运维 ---
|
||
|
||
def prune(self, *, keep_terminal: int) -> int:
|
||
with self._transaction():
|
||
terminal = [
|
||
row for row in self._rows().values() if TaskStatus(row["status"]) in TERMINAL
|
||
]
|
||
terminal.sort(key=lambda row: row["created_at"] or "")
|
||
excess = len(terminal) - max(0, keep_terminal)
|
||
if excess <= 0:
|
||
return 0
|
||
for row in terminal[:excess]:
|
||
task_id = row["task_id"]
|
||
self._rows().pop(task_id, None)
|
||
self._data["transitions"].pop(task_id, None)
|
||
self._flush()
|
||
return excess
|
||
|
||
def close(self) -> None:
|
||
"""写回并释放。文件后端没有长连接，但保持同样的契约形状。"""
|
||
with self._transaction():
|
||
self._flush()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 1. 跑完整的契约套件
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestJsonFileStore(StoreContract):
|
||
"""第三方后端只要通过这一套，就可以替换内置实现。"""
|
||
|
||
_tempdir: tempfile.TemporaryDirectory[str]
|
||
|
||
def make_store(self) -> TaskStore:
|
||
self._tempdir = tempfile.TemporaryDirectory(prefix="chongming-jsonstore-")
|
||
return JsonFileTaskStore(Path(self._tempdir.name) / "tasks.json")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 2. 证明它能接进门面（而不只是自己自洽）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_a_third_party_store_plugs_into_the_controller(tmp_path: Path) -> None:
|
||
"""契约通过 ≠ 能用。这条用例证明它真的能驱动 Controller。"""
|
||
store = JsonFileTaskStore(tmp_path / "tasks.json")
|
||
controller = Controller(store=store, command_poll_interval=0.05)
|
||
controller.register("add", _add)
|
||
run = controller.submit("add", 2, 3, task_id="t1", bucket_key="tenant-1")
|
||
outcome = controller.wait("t1", timeout=10)
|
||
|
||
assert outcome.ok
|
||
assert outcome.value == 5
|
||
assert controller.persist_failures == (), "写穿失败被吞掉了"
|
||
|
||
persisted = store.get("t1")
|
||
assert persisted is not None
|
||
assert persisted.status is TaskStatus.SUCCEEDED
|
||
assert persisted.bucket == controller.bucket_policy.bucket_for("tenant-1")
|
||
assert [item.to_status for item in store.history("t1")][-1] is TaskStatus.SUCCEEDED
|
||
assert run.status is TaskStatus.SUCCEEDED
|
||
controller.close(timeout=10)
|
||
|
||
|
||
def test_the_naive_file_store_is_not_cross_process_safe(tmp_path: Path) -> None:
|
||
"""把这条局限**钉成用例**，而不是写在注释里自欺欺人。
|
||
|
||
无锁的文件型后端在 Windows 上无法支撑"两个实例共用一个文件"：``Path.replace``
|
||
的原子替换与并发打开会互相冲突。真跑起来的表现是 ``PermissionError`` 或读回
|
||
空文件。
|
||
|
||
所以本范例演示的是**接口形状**，不是跨进程方案。要跨进程就用内置的
|
||
``SqliteTaskStore``——文件锁、WAL、``BEGIN IMMEDIATE`` 都由 SQLite 负责，
|
||
这正是它作为默认实现的理由。
|
||
"""
|
||
path = tmp_path / "tasks.json"
|
||
first = JsonFileTaskStore(path)
|
||
second = JsonFileTaskStore(path)
|
||
first.create(snapshot_with_id("t1"))
|
||
|
||
# 单实例读写没问题（每次操作都重新读盘）。
|
||
assert second.get("t1") is not None
|
||
second.touch("t1", heartbeat_at=datetime.now(UTC))
|
||
assert first.get("t1").last_heartbeat_at is not None # type: ignore[union-attr]
|
||
|
||
# 真正的并发访问则不可靠。这里只断言"我们知道它不可靠"这个事实本身：
|
||
# 两个实例各自持锁，因此谁也无法阻止对方在自己读的瞬间替换文件。
|
||
assert first._lock is not second._lock, "没有共享锁，正是它不能跨实例的原因"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 3. 文档与协议不许漂移
|
||
# ---------------------------------------------------------------------------
|
||
|
||
#: 契约的方法全集。``docs/store-guide.md`` 必须逐个说明，协议本身也必须一字不差。
|
||
DOCUMENTED_METHODS = frozenset(
|
||
{
|
||
"append",
|
||
"claim_commands",
|
||
"close",
|
||
"complete_command",
|
||
"counts",
|
||
"create",
|
||
"get",
|
||
"history",
|
||
"prune",
|
||
"query",
|
||
"release_command",
|
||
"save",
|
||
"submit_command",
|
||
"touch",
|
||
}
|
||
)
|
||
|
||
|
||
def test_protocol_surface_matches_the_documented_list() -> None:
|
||
"""协议加了方法却忘了写进指南（或反之），这条会失败。"""
|
||
actual = {name for name in dir(TaskStore) if not name.startswith("_")}
|
||
assert actual == DOCUMENTED_METHODS
|
||
|
||
|
||
def test_store_guide_documents_every_protocol_method() -> None:
|
||
"""指南必须逐个说明每个方法——否则"换持久层"就没有可依据的规格。"""
|
||
guide = (_REPO_ROOT / "docs" / "store-guide.md").read_text(encoding="utf-8")
|
||
missing = [
|
||
name
|
||
for name in sorted(DOCUMENTED_METHODS)
|
||
if f"`{name}(" not in guide and f"`{name}`" not in guide
|
||
]
|
||
assert missing == [], f"以下方法在 docs/store-guide.md 里没有说明：{missing}"
|
||
|
||
|
||
def test_store_guide_mentions_every_shipped_implementation() -> None:
|
||
guide = (_REPO_ROOT / "docs" / "store-guide.md").read_text(encoding="utf-8")
|
||
for name in (
|
||
"InMemoryTaskStore",
|
||
"SqliteTaskStore",
|
||
"PostgresTaskStore",
|
||
"StoreContract",
|
||
):
|
||
assert name in guide, f"指南里没有提到 {name}"
|
||
|
||
|
||
# --- 任务 ---
|
||
|
||
|
||
def _add(ctx: Any, left: int, right: int) -> int:
|
||
del ctx
|
||
return left + right
|
||
|
||
|
||
def snapshot_with_id(task_id: str) -> TaskSnapshot:
|
||
"""构造一个最小的合法快照。"""
|
||
return TaskSnapshot(
|
||
task_id=task_id,
|
||
name="t",
|
||
status=TaskStatus.RUNNING,
|
||
status_version=1,
|
||
attempt=1,
|
||
max_attempts=1,
|
||
created_at=datetime(2026, 1, 1, tzinfo=UTC),
|
||
)
|
||
|
||
|
||
def _block(ctx: Any) -> str:
|
||
for _ in range(1_000_000):
|
||
ctx.checkpoint()
|
||
time.sleep(0.005)
|
||
return "done"
|
||
|
||
|
||
def test_documented_methods_all_exist_on_the_protocol() -> None:
|
||
"""指南里写的名字必须真实存在——防止文档里出现一个不存在的接口。"""
|
||
for name in DOCUMENTED_METHODS:
|
||
assert hasattr(TaskStore, name), f"{name} 不在 TaskStore 上"
|