773 KiB
773 KiB
./pyproject.toml
[build-system]
build-backend = "hatchling.build"
requires = ["hatchling"]
[project]
dependencies = []
description = "精细化状态控制的 Python 管道任务控制工具"
name = "chongming-task"
readme = "README.md"
requires-python = ">=3.11"
version = "0.1.0"
classifiers = [
"Development Status :: 3 - Alpha",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Programming Language :: Python :: 3.14",
]
[project.scripts]
chongming-task = "chongming_task.cli.main:main"
[project.optional-dependencies]
postgres = [
"psycopg-pool>=3.2",
"psycopg[binary]>=3.2",
]
# 核心包运行时零第三方依赖（CLI 也只用标准库 argparse）。
# 观测后端等适配器在各自阶段追加为 [project.optional-dependencies]。
[dependency-groups]
dev = [
"pytest>=8.0",
"hypothesis>=6.100",
"ruff>=0.6",
"mypy>=1.11",
"psycopg[binary]>=3.2",
"psycopg-pool>=3.2",
"twine>=7.0.0",
]
[tool.hatch.build.targets.wheel]
packages = ["src/chongming_task"]
# ---------------------------------------------------------------------------
# 工具链
# ---------------------------------------------------------------------------
[tool.ruff]
line-length = 100
target-version = "py311"
src = ["src", "tests"]
# pytest 的临时目录默认落在项目根下，测试往里写的临时 .py 文件不该被 lint 扫到
# （其中不乏**故意写坏**的样本，例如"导入一个不存在的依赖"）。
extend-exclude = ["pytest-of-*"]
[tool.ruff.lint]
select = ["E", "F", "I", "UP", "B", "SIM", "RUF", "C4", "PTH"]
ignore = [
"RUF001", # 允许注释/字符串中的中文全角标点
"RUF002",
"RUF003",
]
[tool.ruff.lint.per-file-ignores]
"tests/**" = ["B011"]
[tool.ruff.lint.isort]
known-first-party = ["chongming_task"]
# 内核必须保持可被状态机属性测试穷举验证的纯粹性，故 src 全量 strict。
[tool.mypy]
python_version = "3.11"
strict = true
warn_unreachable = true
warn_redundant_casts = true
show_error_codes = true
files = ["src"]
[[tool.mypy.overrides]]
module = "tests.*"
disallow_untyped_defs = false
disallow_incomplete_defs = false
[tool.pytest.ini_options]
testpaths = ["tests"]
addopts = "-ra --strict-markers --strict-config"
markers = [
"slow: 耗时较长的集成测试",
"windows_only: 仅在 Windows 上有意义的用例",
"posix_only: 仅在 POSIX 上有意义的用例",
]
filterwarnings = ["error"]
./examples\01_quickstart.py
"""示例 1：跑一个任务，并观察它在生命周期中经过的状态。
运行::
uv run python examples/01_quickstart.py
"""
from __future__ import annotations
import time
from chongming_task.run import TaskRun
from chongming_task.runtime.cancellation import CancellationToken
def extract(ctx: CancellationToken, source: str) -> list[int]:
"""模拟一次数据抽取。
要点：长耗时循环里必须周期性地调用 ``ctx.checkpoint()``。它同时做三件事：
响应停止请求、响应暂停请求、刷新心跳。
"""
rows: list[int] = []
for index in range(5):
ctx.checkpoint()
time.sleep(0.05)
rows.append(index)
print(f" 抽取完成：{source} -> {len(rows)} 行")
return rows
def main() -> None:
run = TaskRun(extract, ("s3://demo/rows",), task_id="demo-1", name="抽取")
run.start()
outcome = run.wait(timeout=30)
print("\n状态迁移历史（只追加、可审计）：")
for record in run.transitions:
print(
f" {record.from_status.value:>10} -> {record.to_status.value:<10} by {record.trigger}"
)
print(f"\n最终状态：{outcome.status.value}")
print(f"返回值 ：{outcome.value}")
elapsed = (outcome.snapshot.ended_at - outcome.snapshot.started_at).total_seconds()
print(f"耗时 ：{elapsed:.2f}s")
if __name__ == "__main__":
main()
./examples\02_control.py
"""示例 2：运行期干预 —— 暂停、恢复、停止。
演示协作式暂停（D6 定案的语义）与"停不掉时会如实上报"。
运行::
uv run python examples/02_control.py
"""
from __future__ import annotations
import time
from chongming_task.run import TaskRun
def long_running(ctx) -> str:
"""持续做小步工作，每一步都经过协作点。"""
for _ in range(10_000):
ctx.checkpoint()
time.sleep(0.01)
return "跑完了"
def uncooperative(ctx) -> str:
"""刻意**不**调用 checkpoint：这样的任务既停不掉也暂停不了。
它约 4 秒后自行结束，以免在示例里留下后台线程。
"""
for _ in range(400):
time.sleep(0.01)
return "无视了停止请求"
def demo_pause_resume() -> None:
print("== 暂停 / 恢复 ==")
run = TaskRun(long_running, task_id="demo-pause")
run.start()
time.sleep(0.2)
print(f" 暂停请求已受理：{run.pause(timeout=5)}；当前状态 {run.status.value}")
time.sleep(0.2)
print(f" 暂停期间状态保持：{run.status.value}")
run.resume()
print(f" 恢复后状态 ：{run.status.value}")
run.stop(reason="演示主动停止")
outcome = run.wait(timeout=10)
print(f" 停止后终态 ：{outcome.status.value}")
print(f" 失败归类 ：{outcome.failure.kind.value if outcome.failure else '-'}")
print(f" 可重试 ：{outcome.failure.retryable if outcome.failure else '-'}")
def demo_unstoppable() -> None:
print("\n== 停不掉的任务：如实上报 ==")
run = TaskRun(uncooperative, task_id="demo-stubborn", graceful_stop=0.3)
run.start()
time.sleep(0.2)
run.stop(reason="试一下")
outcome = run.wait(timeout=10)
print(f" 终态 ：{outcome.status.value}")
print(f" 失败归类 ：{outcome.failure.kind.value if outcome.failure else '-'}")
print(" 含义 ：线程无法强杀，任务可能仍在后台产生副作用，因此不敢说'已停止'")
def main() -> None:
demo_pause_resume()
demo_unstoppable()
if __name__ == "__main__":
main()
./examples\03_retry_and_timeout.py
"""示例 3：重试与超时。
运行::
uv run python examples/03_retry_and_timeout.py
"""
from __future__ import annotations
import time
from chongming_task.policies import ExponentialBackoff, RetryPolicy, Timeout
from chongming_task.run import run_task
_attempts = 0
def flaky(ctx, succeed_on_attempt: int) -> str:
"""前几次故意失败，用来演示重试与退避。"""
global _attempts
_attempts += 1
if _attempts < succeed_on_attempt:
raise RuntimeError(f"第 {_attempts} 次尝试遇到暂时性故障")
return f"第 {_attempts} 次尝试成功"
def slow(ctx) -> str:
"""跑得很慢，用来演示超时后会被请求停止。"""
for _ in range(10_000):
ctx.checkpoint()
time.sleep(0.01)
return "跑完了"
def demo_retry() -> None:
print("== 重试：失败两次后成功 ==")
_attempts = 0
outcome = run_task(
flaky,
3,
retry=RetryPolicy(max_attempts=5, backoff=ExponentialBackoff(base=0.05, factor=2)),
task_id="demo-retry",
name="会抖的任务",
)
print(f" 终态 ：{outcome.status.value}")
print(f" 返回值 ：{outcome.value}")
print(f" 尝试次数 ：{outcome.snapshot.attempt}/{outcome.snapshot.max_attempts}")
def demo_retry_exhausted() -> None:
print("\n== 重试：预算耗尽 ==")
global _attempts
_attempts = 0
outcome = run_task(
flaky,
99, # 永远不会成功
retry=RetryPolicy(max_attempts=2, backoff=ExponentialBackoff(base=0.05, factor=1)),
task_id="demo-exhausted",
name="必定失败",
)
print(f" 终态 ：{outcome.status.value}")
print(f" 尝试次数 ：{outcome.snapshot.attempt}/{outcome.snapshot.max_attempts}")
print(f" 失败归类 ：{outcome.failure.kind.value if outcome.failure else '-'}")
def demo_timeout() -> None:
print("\n== 超时：单次尝试不超过 0.5 秒 ==")
outcome = run_task(
slow,
timeout=Timeout(per_attempt=0.5),
task_id="demo-timeout",
name="慢任务",
)
print(f" 终态 ：{outcome.status.value}")
print(f" 失败归类 ：{outcome.failure.kind.value if outcome.failure else '-'}")
print(f" 可重试 ：{outcome.failure.retryable if outcome.failure else '-'}")
def main() -> None:
demo_retry()
demo_retry_exhausted()
demo_timeout()
if __name__ == "__main__":
main()
./examples\04_controller.py
"""示例 4：用 Controller 管理一批任务（注册 / 提交 / 排队 / 查询 / 控制）。
这是管理后台的核心：**一个统一入口**持有任务定义、运行中的任务与并发额度。
将来的 CLI 或看板都只是这一层的薄适配。
运行::
uv run python examples/04_controller.py
"""
from __future__ import annotations
import time
from chongming_task.controller import Controller, render_table
def transform(ctx, rows: int) -> int:
"""转换：逐行推进。"""
for _ in range(30):
ctx.checkpoint()
time.sleep(0.02)
return rows * 2
def extract(ctx, source: str) -> str:
"""抽取：比转换慢一倍。"""
for _ in range(60):
ctx.checkpoint()
time.sleep(0.02)
return f"{source} -> 1200 行"
def main() -> None:
controller = Controller(max_parallel=2)
controller.register("transform", transform, tags={"cpu"})
controller.register("extract", extract, tags={"io"})
controller.register("load", transform, description="写回目标库")
print("已注册任务：", "、".join(controller.registry.names()))
print("并发上限：2\n")
for index in range(4):
controller.submit(
"transform", 100 * (index + 1), task_id=f"t{index}", label=f"转换-{index + 1}"
)
print(f"提交了 4 个任务；此刻排队 {controller.queued_count} 个（超出并发上限的部分）")
time.sleep(0.4)
print("\n--- 查询：运行中的视图 ---")
print(render_table(controller.list_tasks()))
stats = controller.stats()
print(
f" 未结束 {stats.unfinished}"
f"（其中排队 {controller.queued_count}，占用执行额度 {controller.active_count}）"
)
controller.wait_all(timeout=30)
print("\n--- 查询：全部结束 ---")
print(render_table(controller.list_tasks()))
print("\n>>> 提交一个长任务，然后中途停止它")
controller.submit("extract", "s3://demo/big", task_id="t-long", label="抽取(大)")
time.sleep(0.3)
print(f" 当前状态：{controller.get('t-long').status.value}")
controller.stop("t-long", reason="手工干预")
outcome = controller.wait("t-long", timeout=10)
print(
f" 终态：{outcome.status.value}"
f"（失败归类 {outcome.failure.kind.value}，可重试 {outcome.failure.retryable}）"
)
print("\n--- 最终统计 ---")
final = controller.stats()
for status, count in sorted(final.by_status.items(), key=lambda item: item[0].value):
print(f" {status.value:<10} {count}")
print("\n提示：另一个进程里的 chongming-task ps 看不到这些任务——内存不共享，")
print(" 跨进程查询需要共享存储（stdlib sqlite3），那是下一步。")
if __name__ == "__main__":
main()
./examples\05_buckets.py
"""示例 5：任务分桶 —— 本地并发隔离，以及为分布式预留的分配单位。
运行::
uv run python examples/05_buckets.py
"""
from __future__ import annotations
import time
from chongming_task.controller import Controller, render_table
def sync_tenant(ctx, tenant: str) -> str:
"""同步一个租户的数据。"""
for _ in range(20):
ctx.checkpoint()
time.sleep(0.02)
return f"{tenant} 同步完成"
def main() -> None:
controller = Controller(max_parallel=10, max_parallel_per_bucket=1)
controller.register("sync", sync_tenant)
policy = controller.bucket_policy
print(f"分桶策略：{policy.num_buckets} 个桶")
print("并发限制：全局 10，但每个桶最多 1 个\n")
for index in range(4):
controller.submit(
"sync", "租户A", task_id=f"a{index}", label="A-同步", bucket_key="tenant-A"
)
for index in range(2):
controller.submit(
"sync", "租户B", task_id=f"b{index}", label="B-同步", bucket_key="tenant-B"
)
print("业务键 -> 桶号（同一个键在任何进程、任何机器上都会落到同一个桶）：")
for key in ("tenant-A", "tenant-B"):
print(f" {key:<10} -> bucket {policy.bucket_for(key)}")
time.sleep(0.5)
print("\n--- 查询：A 桶积压，但不影响 B 桶 ---")
print(render_table(controller.list_tasks()))
print(
f" 占用额度 {controller.active_count}，排队 {controller.queued_count}，"
f"A 桶在用 {controller.active_in_bucket(policy.bucket_for('tenant-A'))}"
)
controller.wait_all(timeout=30)
print("\n--- 查询：全部结束 ---")
print(render_table(controller.list_tasks()))
print("\n各桶任务数（分布式调度要靠它决定把桶分给谁）：")
for bucket, count in sorted(
controller.bucket_counts().items(), key=lambda kv: (kv[0] is None, kv[0])
):
print(f" bucket {bucket}: {count}")
bucket_a = policy.bucket_for("tenant-A")
print(f"\n按桶查询（只取 A 桶）：{[s.task_id for s in controller.list_tasks(bucket=bucket_a)]}")
print("\n分布式预留：桶是分配单位——一台机器拥有一组桶，只取自己那组桶的任务即可。")
if __name__ == "__main__":
main()
./examples\06_cross_process.py
"""示例 6：跨进程 —— 一个进程跑任务，另一个进程查询与控制。
这里用**同一个 SQLite 文件上的两个独立连接**模拟两个进程。真正的两个进程行为完全
一样：它们看不到彼此的内存，只能通过存储交换状态与指令。
运行::
uv run python examples/06_cross_process.py
"""
from __future__ import annotations
import tempfile
import time
from pathlib import Path
from chongming_task import Controller, SqliteTaskStore, render_table
def sync_tenant(ctx, tenant: str) -> str:
"""稳步推进直到被停止。"""
for _ in range(300):
ctx.checkpoint()
time.sleep(0.02)
return f"{tenant} 同步完成"
def _wait_for_status(client: Controller, task_id: str, wanted: str, timeout: float = 20.0) -> bool:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
snapshot = client.get(task_id)
if snapshot is not None and snapshot.status.value == wanted:
return True
time.sleep(0.02)
return False
def main() -> None:
with tempfile.TemporaryDirectory(prefix="chongming-demo-") as directory:
path = Path(directory) / "tasks.db"
# ---- 进程 A：真正跑任务的 worker ----
worker = Controller(
store=SqliteTaskStore(path),
max_parallel=2,
command_poll_interval=0.1,
owner="worker-1",
)
worker.register("sync", sync_tenant)
worker.submit("sync", "租户A", task_id="job-1", label="A-同步", bucket_key="tenant-A")
worker.submit("sync", "租户B", task_id="job-2", label="B-同步", bucket_key="tenant-B")
print("worker-1 已提交 2 个任务（并发上限 2）\n")
time.sleep(0.5)
# ---- 进程 B：管理客户端，只查询与控制，不执行任何任务 ----
client = Controller(store=SqliteTaskStore(path), owner="ops-console")
print("--- ops-console 看到的状态（它并不持有这些任务）---")
print(render_table(client.list_tasks()))
print(f" 它本地的任务句柄数：{len(client.runs())}")
print(f" 统计：{ {s.value: n for s, n in client.stats().by_status.items()} }")
print("\n>>> ops-console 跨进程停止 job-1")
print(f" 下发结果：{client.stop('job-1', reason='运维从后台停止')}")
print(" （客户端无法给另一个进程的线程发信号，只能写一条指令）")
if _wait_for_status(client, "job-1", "stopped"):
print(" worker-1 已取走指令并停止该任务")
else:
print(" !! 超时：指令没有被取走")
print("\n--- job-1 的迁移历史（从存储读，可审计）---")
for record in client.store.history("job-1"): # type: ignore[union-attr]
print(
f" {record.version:>2} {record.from_status.value:>10} -> "
f"{record.to_status.value:<10} by {record.trigger}"
)
print("\n--- 让 job-2 自己跑完，再看最终状态 ---")
worker.wait_all(timeout=30)
print(render_table(client.list_tasks()))
# 收尾：两个 Controller 各自关掉自己的存储连接。没关的 SQLite 连接会泄漏
# 文件句柄，在 Windows 上还会让数据库文件无法删除。
worker.close(timeout=10, close_store=True)
client.close(timeout=5, close_store=True)
print("\n要点：worker 负责执行，console 负责查询与控制，两者只通过存储耦合。")
if __name__ == "__main__":
main()
./examples\07_pipeline.py
"""示例 7：DAG 管道 —— 一个源 → 两个分支 → 汇合。
管道把多个任务按依赖关系串起来，由一个调度器推进：节点先停在 ``blocked`` 等上游，
上游满足后才放行；注定跑不了的下游会被 ``skipped``。
运行::
uv run python examples/07_pipeline.py
"""
from __future__ import annotations
import time
from typing import Any
from chongming_task import Controller, Pipeline, render_pipeline
from chongming_task.policies import AggregationPolicy, JoinMode
def fetch(ctx: Any, day: str) -> str:
"""取数：模拟一次耗时操作。"""
for _ in range(10):
ctx.checkpoint()
time.sleep(0.03)
return f"{day} 的数据"
def transform(ctx: Any, day: str, shard: int) -> str:
"""分片转换。"""
for _ in range(8):
ctx.checkpoint()
time.sleep(0.02)
return f"{day}/shard{shard}"
def load(ctx: Any, day: str) -> str:
"""入库：依赖全部分片完成。"""
del ctx
return f"{day} 已入库"
def audit(ctx: Any, day: str) -> str:
"""审计：**不管上游成败都要跑**（ALL_SETTLED）。"""
del ctx
return f"{day} 审计完成"
def main() -> None:
controller = Controller(max_parallel=4)
# ---- 定义 ----
pipeline = Pipeline("nightly-etl", aggregation=AggregationPolicy.all_settled())
controller.register("fetch", fetch, tags={"io"})
controller.register("transform", transform)
controller.register("load", load, tags={"io"})
controller.register("audit", audit)
fetch_node = pipeline.node("fetch", "2026-10-02", node_id="fetch")
shards = [
pipeline.node(
"transform",
"2026-10-02",
index,
node_id=f"shard{index}",
depends_on=[fetch_node],
)
for index in range(3)
]
# 入库：等所有分片成功（默认 ALL_SUCCEEDED 屏障）
load_node = pipeline.node("load", "2026-10-02", node_id="load", depends_on=shards)
# 审计：等 load 落终态即可，不管成败
pipeline.node(
"audit",
"2026-10-02",
node_id="audit",
depends_on=[load_node],
join=JoinMode.ALL_SETTLED,
)
print(f"管道 {pipeline.name!r} 共 {len(pipeline)} 个节点，拓扑序：")
print(" " + " -> ".join(node.node_id for node in pipeline.topological_order()))
print(" 注意：位置参数是传给任务的，节点名由 node_id= 指定\n")
# ---- 跑起来 ----
run = controller.submit_pipeline(pipeline)
print(f"已提交，pipeline_run_id = {run.run_id}\n")
time.sleep(0.05)
print("--- 刚提交时：下游都停在 blocked ---")
print(render_pipeline(run))
print("\n--- 等待收场 ---")
outcome = run.wait(timeout=60)
print(render_pipeline(run))
print(f"\n管道结论：{outcome.status.value}")
for node_id, status in outcome.nodes.items():
print(f" {node_id:<8} {status.value}")
print(f"\n每个任务都带 pipeline_run_id：{_run_id_of(controller, run, 'load')}")
controller.close(timeout=10)
print("\n要点：依赖由状态机表达（blocked -> queued），调度器只负责放行。")
def _run_id_of(controller: Controller, run: Any, node_id: str) -> str:
snapshot = controller.get(run.task_id_of(node_id))
return "?" if snapshot is None else str(snapshot.pipeline_run_id)
if __name__ == "__main__":
main()
./examples\08_task_context.py
"""示例 8：``@task`` 装饰器与 ``TaskContext``。
演示四件事：
* ``@task`` 把元数据写在函数旁边，``register_module`` 一次登记整个模块；
* ``ctx.progress(...)`` 上报进度——它会进快照、进存储，并出现在 ``ps`` 的 PROGRESS 列；
* ``ctx.log(...)`` 走标准库 ``logging``，logger 名字是 ``chongming_task.task.<task_id>``；
* ``ctx.input(节点id)`` 取上游节点的返回值——管道之间真正能传值。
运行::
uv run python examples/08_task_context.py
"""
from __future__ import annotations
import logging
import sys
import time
from typing import Any
from chongming_task import Controller, Pipeline, RetryPolicy, render_table, task
@task(tags={"io"}, retry=RetryPolicy(max_attempts=2))
def extract(ctx: Any, day: str) -> list[str]:
"""按天抽取分区列表。"""
ctx.log(f"开始抽取 {day}")
partitions = [f"{day}/part{i}" for i in range(6)]
for index, part in enumerate(partitions, start=1):
ctx.checkpoint()
ctx.progress(index, total=len(partitions), unit="partitions", message=part)
time.sleep(0.05)
return partitions
@task
def transform(ctx: Any, day: str) -> int:
"""把上游抽出的分区逐个转换，返回处理条数。"""
partitions = ctx.input("extract", [])
ctx.log(f"待转换分区：{len(partitions)} 个")
rows = 0
for index, part in enumerate(partitions, start=1):
ctx.checkpoint()
rows += 100
ctx.progress(index, total=len(partitions), unit="partitions", message=part)
time.sleep(0.05)
return rows
@task
def load(ctx: Any, day: str) -> str:
"""入库：条数来自上游的返回值，而不是自己再算一遍。"""
rows = ctx.input("transform", 0)
ctx.log(f"{day} 共 {rows} 行待入库")
for index in range(1, 5):
ctx.checkpoint()
ctx.progress(index, total=4, unit="batch")
time.sleep(0.04)
return f"{day}: {rows} 行已入库"
def main() -> None:
logging.basicConfig(
level=logging.INFO,
format="%(levelname)-7s %(name)s %(message)s",
stream=sys.stdout,
)
controller = Controller(max_parallel=2)
# ---- 一次登记整个模块 ----
registered = controller.register_module(sys.modules[__name__])
print("从 @task 元数据登记的：")
for item in registered:
print(f" {item.name:<10} tags={sorted(item.tags) or '-'} {item.description}")
# ---- 串成管道：数据沿依赖流动 ----
pipeline = Pipeline("daily-etl")
extract_node = pipeline.node("extract", "2026-10-02", node_id="extract")
transform_node = pipeline.node(
"transform", "2026-10-02", node_id="transform", depends_on=[extract_node]
)
pipeline.node("load", "2026-10-02", node_id="load", depends_on=[transform_node])
run = controller.submit_pipeline(pipeline)
print(f"\n已提交管道 {run.run_id}\n")
# ---- 跑的过程中看进度 ----
deadline = time.monotonic() + 30
while time.monotonic() < deadline and not run.settled:
snapshots = controller.list_tasks(pipeline_run_id=run.run_id)
cells: list[str] = []
for snapshot in snapshots:
node = snapshot.name.split("/", 1)[-1]
progress = "-"
if snapshot.progress is not None:
progress = f"{snapshot.progress.done}/{snapshot.progress.total}"
cells.append(f"{node:<10} {snapshot.status.value:<10} {progress:<8}")
sys.stdout.write("\r" + " | ".join(cells).ljust(120))
sys.stdout.flush()
time.sleep(0.1)
outcome = run.wait(timeout=30)
print("\n\n--- 结束 ---")
print(render_table(controller.list_tasks(pipeline_run_id=run.run_id)))
print(f"\n管道结论：{outcome.status.value}")
final = controller.get(run.task_id_of("load"))
print(f"load 的返回值：{final.outcome if final else None}") # type: ignore[union-attr]
controller.close(timeout=10)
print(
"\n要点：ctx 是任务与运行时之间唯一的协作入口——停止、暂停、进度、日志、上游数据都在这里。"
)
if __name__ == "__main__":
main()
./examples\09_long_running.py
"""示例 9：无限循环任务的**启与停**。
一个"永远跑下去"的任务该怎么写、怎么启动、怎么停，以及三种收尾方式的终态差异。
运行::
uv run python examples/09_long_running.py
## 三条规则
1. **任务体每轮必须调 ``ctx.checkpoint()``。** 停止是**协作式**的：只有调用了它，任务
才会看到停止请求。不调用 = 停不掉（见文件最后的对照）。
2. **停止请求有两种收尾方式，终态不同**：
================================================ =========== ==========
写法 终态 返回值
================================================ =========== ==========
``while True: ctx.checkpoint()``（让它抛） ``stopped`` 无
检查 ``ctx.is_stop_requested`` 后正常 ``return`` ``succeeded`` ✅ 有
完全不调 ``checkpoint()`` ``lost`` 无
================================================ =========== ==========
"停止"在第一种写法里是"立刻中止"，在第二种里是"**请收工**"——后者会把已经处理好的
结果留下来，对"跑到下一个安全点再退出"这类需求更合适。
3. **停止是终态。** 目前没有 ``retry``，停止/失败之后想再跑要换一个 ``task_id``
重新提交（``create`` 会拒绝重复 id）。所以日常"先停一下"应该用 ``pause`` / ``resume``，
``stop`` 留给"我要它彻底结束"。
"""
from __future__ import annotations
import sys
import threading
import time
from typing import Any
from chongming_task import Controller, InMemoryTaskStore, render_table
def worker(ctx: Any, source: str) -> int:
"""无限循环：持续处理，直到收到停止请求后**优雅收尾**。"""
processed = 0
while not ctx.is_stop_requested:
ctx.checkpoint()
processed += 1
ctx.progress(processed, unit="batch", message=f"{source} 第 {processed} 批")
if processed % 20 == 0:
ctx.log(f"已处理 {processed} 批")
time.sleep(0.02)
# 走到这里说明收到停止请求了：做收尾，然后正常返回——
# 返回值会被保留下来，终态是 succeeded 而不是 stopped。
ctx.log(f"收到停止请求，收尾：本次共处理 {processed} 批")
return processed
def main() -> None:
store = InMemoryTaskStore()
controller = Controller(store=store, owner="node-a")
controller.register("worker", worker)
# ---- 启动 ----
run = controller.submit("worker", "kafka://orders", task_id="orders-worker")
print(f"已启动 {run.snapshot.task_id} 状态={run.status.value}")
time.sleep(0.5)
print("\n--- 运行中 ---")
print(render_table(controller.list_tasks()))
# ---- 日常先停一下：pause / resume（不是 stop） ----
print("\n--- pause（就地挂起，仍占用执行槽）---")
print(f" pause 是否确认: {controller.pause('orders-worker', timeout=5)}")
print(f" 状态: {run.status.value}")
print("\n--- resume ---")
controller.resume("orders-worker")
time.sleep(0.2)
print(f" 状态: {run.status.value}")
# ---- 让它彻底结束：stop ----
print("\n--- stop（终态，不可逆）---")
controller.stop("orders-worker", reason="发布窗口")
outcome = controller.wait("orders-worker", timeout=10)
print(f" 终态: {outcome.status.value}")
print(f" 返回值: {outcome.value} ← 优雅收尾把已处理的数量带了出来")
print("\n--- 结束后 ---")
print(render_table(controller.list_tasks()))
print("\n--- 审计轨迹（谁按的、为什么，都在这里）---")
for record in store.history("orders-worker"):
reason = f" 原因={record.reason}" if record.reason else ""
print(
f" {record.version:>2} {record.from_status.value:>9} -> "
f"{record.to_status.value:<9} {record.trigger}{reason}"
)
# ---- 停止是终态：同一个 task_id 不能复活 ----
print("\n--- 想再跑一次？同一个 task_id 会被拒绝 ---")
try:
controller.submit("worker", "kafka://orders", task_id="orders-worker")
except Exception as exc:
print(f" {type(exc).__name__}: {exc}")
print(" 换一个 task_id 重新提交即可（或让进程一直活着，用 pause/resume 做日常启停）。")
controller.close(timeout=5)
_contrast_uncooperative()
_print_cli_usage()
def _contrast_uncooperative() -> None:
"""对照：不调 ``checkpoint()`` 的任务**停不掉**。
这不是缺陷，是诚实：线程无法被强杀，所以系统只能如实报 ``lost``（"结果未知"），
而任务体仍在后台继续跑。真正终止它只能靠结束进程。
"""
print("\n" + "=" * 70)
print("对照：一个从不调 checkpoint() 的无限任务")
print("=" * 70)
release = threading.Event()
ticks = {"n": 0}
def uncooperative(ctx: Any) -> int:
while not release.is_set():
ticks["n"] += 1
time.sleep(0.02)
return ticks["n"]
controller = Controller(store=InMemoryTaskStore(), graceful_stop=0.5)
controller.register("stubborn", uncooperative)
controller.submit("stubborn", task_id="stubborn")
time.sleep(0.3)
print(" 停止请求已发出，等待载体响应……")
controller.stop("stubborn")
outcome = controller.wait("stubborn", timeout=8)
print(f" 终态: {outcome.status.value} ← 不是 stopped，而是 lost（结果未知）")
print(f" 失败类型: {outcome.failure.kind.value if outcome.failure else None}")
before = ticks["n"]
time.sleep(0.4)
alive = [t.name for t in threading.enumerate() if t.is_alive() and t.name.startswith("task-")]
print(f" 任务体的线程还在吗: {alive}")
print(f" 它还在继续干活吗: {ticks['n'] > before}（{before} -> {ticks['n']}）")
print(" → 线程无法被强杀，只能靠结束进程来真正终止它。")
release.set()
controller.close(timeout=3)
def _print_cli_usage() -> None:
"""运维视角：命令行怎么启、怎么停。"""
print("\n" + "=" * 70)
print("命令行（另一个进程，通过共享存储启停）")
print("=" * 70)
print(
"""
# 启动：前台执行，交给 systemd / nohup 保持存活
# 注意 --bucket-key：限定了 --owned-buckets 时它是**必需**的，否则任务没有归属桶、
# 不会有任何节点认领它（CLI 会在提交时报错）。
chongming-task --db /var/lib/tasks.db --num-buckets 64 --owned-buckets 0-15 \\
run --tasks myapp.tasks worker --arg source=kafka://orders \\
--task-id orders-worker --bucket-key orders
# 另一个终端：看状态、停止
chongming-task --db /var/lib/tasks.db --owned-buckets 0-15 ps --all --mine
chongming-task --db /var/lib/tasks.db --owned-buckets 0-15 status orders-worker
chongming-task --db /var/lib/tasks.db --owned-buckets 0-15 stop orders-worker --reason 发布窗口
# Ctrl-C 也可以：main() 会走 finally 里的 controller.close()，把在跑的任务停干净
""".rstrip()
)
print(" 注意：`run` 是前台进程——它退出时任务就没人管了，请交给进程管理器。")
sys.stdout.write("")
if __name__ == "__main__":
main()
./examples\10_worker_walkthrough.py
"""示例 10：常驻任务的完整启停流程 —— **一个脚本跑通，不需要第二个终端**。
前面几轮把桶归属、跨进程、多终端全搅在一起了，所以这个示例刻意只用一条主线：
启动 → 边跑边看（进度/心跳/租约）→ 暂停 → 恢复 → 停止 → 校验
每一步都打印"你该看什么"，最后**断言**进程内的结论与存储里的状态一致——不一致会
失败退出，不会让你盯着两段输出自己猜。
运行::
uv run python examples/10_worker_walkthrough.py
命令行版本（需要第二个终端，因为 ``run`` 是前台进程）见文件末尾的说明。
"""
from __future__ import annotations
import sys
import time
from typing import Any
from chongming_task import Controller, SqliteTaskStore, render_table
def consume(ctx: Any, source: str, batch: int = 1000) -> int:
"""常驻消费：无限循环 + 批次进度 + 收到停止请求后优雅收尾。"""
total = 0
while not ctx.is_stop_requested:
ctx.checkpoint() # ← 不调用它 = 停不掉
time.sleep(0.05)
total += batch
ctx.progress(total, unit="msg", message=source)
return total
def main() -> int:
# 用**文件**存储：这样存储里的状态与进程内的结论是两份数据，能对得上才有意义。
store = SqliteTaskStore(":memory:") # 想留下文件就把 ":memory:" 换成 "tasks.db"
controller = Controller(store=store, owner="demo-node")
controller.register("consume", consume)
# ---------------------------------------------------------------- 1. 启动
print("=" * 72)
print("1) 启动")
print("=" * 72)
run = controller.submit("consume", "kafka://orders", task_id="orders-worker")
print(" controller.submit(...) 立即返回，任务在后台跑：")
print(f" {run.snapshot.task_id} 状态={run.status.value}")
print("\n 注意：`submit` 之后**必须有人让进程活着**（这里的 main 就是那个人）。")
# ------------------------------------------------- 2. 边跑边看（两个视角）
print("\n" + "=" * 72)
print("2) 边跑边看：进程内的句柄 vs 存储里的事实")
print("=" * 72)
first = _watch(store, seconds=1.0)
second = _watch(store, seconds=1.0)
assert second["progress"] > first["progress"], "进度没有推进——回写链路有问题"
print(f"\n ✅ 进度在推进：{first['progress']} → {second['progress']}")
print(' ✅ 租约新鲜 → 有监督者在管它（进程崩溃时这里会显示"已过期"）')
# ---------------------------------------------------------------- 3. 暂停
print("\n" + "=" * 72)
print('3) pause：日常"先停一下"用这个，不要用 stop')
print("=" * 72)
confirmed = controller.pause("orders-worker", timeout=5)
print(f" pause 返回 {confirmed}（True = 已确认挂起；False = 任务没到检查点，状态已回滚）")
print(f" 状态: {store.get('orders-worker').status.value}") # type: ignore[union-attr]
assert confirmed is True
frozen = _watch(store, seconds=0.6, show_progress=False)
print(f"\n 挂起期间进度不再推进（{frozen['progress']}）—— 但注意：它**仍占着执行槽**")
print("\n resume 后状态: ", end="")
controller.resume("orders-worker")
time.sleep(0.2)
print(store.get("orders-worker").status.value) # type: ignore[union-attr]
# ---------------------------------------------------------------- 4. 停止
print("\n" + "=" * 72)
print('4) stop：终态、不可逆。返回的是"请求已受理"，不是"已停止"')
print("=" * 72)
controller.stop("orders-worker", reason="发布窗口")
outcome = controller.wait("orders-worker", timeout=15)
print(f" 终态: {outcome.status.value} 返回值: {outcome.value}")
print(" ↑ succeeded 而不是 stopped：因为任务体看到 stop 请求后**正常返回**了，")
print(" 已经处理好的数量被完整保留。（让它抛 TaskCancelled 才是 stopped。）")
# ---------------------------------------------------------------- 5. 校验
print("\n" + "=" * 72)
print("5) 校验：进程内的结论 与 存储里的状态 必须一致")
print("=" * 72)
print(render_table(controller.list_tasks()))
stored = store.get("orders-worker")
assert stored is not None
assert stored.status is outcome.status, (
f"存储={stored.status.value} 与进程内={outcome.status.value} 不一致！"
)
assert stored.status_version == outcome.snapshot.status_version, "版本号不一致"
assert stored.outcome == {"value": outcome.value}, "返回值没有落库"
print(
f"\n ✅ 一致：status={stored.status.value} "
f"version={stored.status_version} outcome={stored.outcome}"
)
print("\n 审计轨迹（终态为什么是 succeeded，看这里就清楚了）：")
for record in store.history("orders-worker"):
reason = f" 原因={record.reason}" if record.reason else ""
print(
f" {record.version} {record.from_status.value:>9} -> "
f"{record.to_status.value:<9} {record.trigger}{reason}"
)
failures = controller.persist_failures or "无"
print(f"\n 写穿失败（应当为空；非空说明存储落后于内存）：{failures}")
print(" 永久分叉的任务（应当为空）：", controller.desynced_tasks or "无")
assert controller.persist_failures == (), "有写穿失败——存储不可信"
assert controller.desynced_tasks == {}, "有任务永久分叉"
controller.close(timeout=10, close_store=True)
print("\n" + "=" * 72)
print("全部通过 ✅")
print("=" * 72)
_cli_recipe()
return 0
def _watch(store: SqliteTaskStore, *, seconds: float, show_progress: bool = True) -> dict[str, Any]:
"""从**存储**采一次样（模拟另一个终端跑 `status`），返回关键字段。"""
deadline = time.monotonic() + seconds
snapshot = None
while time.monotonic() < deadline:
snapshot = store.get("orders-worker")
time.sleep(0.1)
assert snapshot is not None
lease = "无"
if snapshot.lease_until is not None:
from chongming_task import utcnow
remaining = (snapshot.lease_until - utcnow()).total_seconds()
lease = f"剩余 {remaining:.0f}s" if remaining > 0 else f"已过期 {-remaining:.0f}s"
progress = snapshot.progress.done if snapshot.progress is not None else 0
if show_progress:
print(
f" 存储: status={snapshot.status.value:<9} v={snapshot.status_version} "
f"progress={progress:<7} 心跳={_hhmmss(snapshot)} 租约={lease}"
)
return {"progress": progress, "status": snapshot.status.value}
def _hhmmss(snapshot: Any) -> str:
return "-" if snapshot.last_heartbeat_at is None else str(snapshot.last_heartbeat_at)[11:19]
def _cli_recipe() -> None:
print(
"""
命令行版本（`run` 是前台进程，所以要**两个终端**）
────────────────────────────────────────────────────────
# 终端 1：启动（前台阻塞；生产上交给 systemd / nohup）
set PYTHONPATH=examples
uv run chongming-task --db ./tasks.db run --tasks worker_tasks consume ^
--arg source=kafka://orders --task-id orders-worker
# 终端 2：看状态 → 停止 → 再确认
uv run chongming-task --db ./tasks.db status orders-worker
uv run chongming-task --db ./tasks.db stop orders-worker --reason 发布窗口
uv run chongming-task --db ./tasks.db status orders-worker <- 这一步看到 succeeded 才算停干净
""".rstrip()
)
print(" 提示：Ctrl-C 也能停（CLI 会走 finally 里的 close()），但 status 确认仍然要做。\n")
if __name__ == "__main__":
sys.exit(main())
./examples\11_best_practice.py
"""示例 11：最佳实践 —— 一份可以直接照抄的多节点任务服务模板。
示例 01–10 各自演示一个机制；这一个把它们收拢成**一个完整、可运行、可运维的
最小形态**，回答集成时的四个必答题：
① 桶怎么划（含节点划分）—— 业务键稳定映射到桶；桶按连续区间分给节点；
一处计算，代码与 CLI 共用。
② 谁跑谁的任务 —— 每个节点只接受自己桶里的任务；配错桶在**提交那一刻**
报错，而不是产出一个永远没有节点执行的任务。
③ 进程启停做什么 —— 启动先 ``reconcile()``（只清理自己桶里的僵尸）；
退出用 ``close(close_store=True)`` 把连接关干净。
④ 怎么观测与控制 —— 所有节点共用一套存储：一个控制台看全部；
跨节点停止走指令通道（写指令 -> 持有者轮询取走）。
文件按"照抄顺序"分四部分：
# 第一部分 部署参数与桶划分
# 第二部分 任务定义（业务代码）
# 第三部分 接线：build_node()（唯一集成点）
# 第四部分 演示：启动 -> 提交 -> 观测 -> 控制 -> 收尾
运行::
uv run python examples/11_best_practice.py
本示例在一个进程里模拟"两个节点 + 一个运维控制台"，是为了**一条命令跑通全流程**；
生产上它们是各自独立的进程（同一份代码，只有 NODE_INDEX 不同），见文件末尾的部署说明。
注意：SQLite 文件存储要求所有节点在同一台机器上；**跨机部署**把 ``build_node()`` 里的
``SqliteTaskStore(deployment.db_path)`` 换成 ``PostgresTaskStore(dsn)``
（``pip install "chongming-task[postgres]"``）即可，业务代码一行都不用改。
"""
from __future__ import annotations
import logging
import os
import sys
import tempfile
import time
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import timedelta
from pathlib import Path
from chongming_task import (
BucketPolicy,
ConcurrencyPolicy,
ConfigurationError,
Controller,
ExponentialBackoff,
RetryPolicy,
SqliteTaskStore,
TaskContext,
TaskSnapshot,
TaskStatus,
render_table,
task,
utcnow,
)
_log = logging.getLogger("best_practice")
# =============================================================================
# 第一部分：部署参数与桶划分（照抄进你的配置模块）
# =============================================================================
#: 桶数。默认 64；这里取 16 让演示输出可读。
#: 它是**部署范围的约定**，不是调优旋钮：所有节点与所有运维命令都必须用同一个数。
#: 经验值：节点数的 4~16 倍（给节点增减留出重新划分的空间）。
NUM_BUCKETS = 16
#: 本示例部署的节点数。生产上每个节点是一个独立进程。
NODE_COUNT = 2
def buckets_for_node(node_index: int, node_count: int, num_buckets: int = NUM_BUCKETS) -> range:
"""把桶**连续**分给各节点：第 ``node_index`` 个节点负责 ``[start, stop)`` 区间。
为什么连续划分，而不是取模交错（``bucket % node_count``）：
* 人可读可核对：``--owned-buckets 0-7`` 与代码里的 ``range(0, 8)`` 是同一件事；
* 负载不偏：桶号本身就是 blake2b 的均匀散列，等长区间的期望负载相同；
* 迁移简单：节点增减时整段区间搬运，不用逐桶列举。
除不尽时余数归靠前的节点，保证**不重不漏**——漏桶是静默积压、重桶是重复执行，
两者都不报错，所以划分必须在一处算出来、处处引用。
"""
if node_count < 1:
raise ValueError(f"node_count 至少为 1，收到 {node_count}")
if num_buckets < node_count:
raise ValueError(
f"桶数（{num_buckets}）必须不小于节点数（{node_count}）："
"否则会有节点分不到桶。经验上桶数取节点数的 4~16 倍。"
)
if not 0 <= node_index < node_count:
raise ValueError(f"node_index 必须在 [0, {node_count}) 内，收到 {node_index}")
size, remainder = divmod(num_buckets, node_count)
start = node_index * size + min(node_index, remainder)
length = size + (1 if node_index < remainder else 0)
return range(start, start + length)
@dataclass(frozen=True)
class Deployment:
"""一个节点进程的部署参数。生产环境从环境变量 / 配置中心来，不要写死在代码里。"""
node_index: int
node_count: int
db_path: Path
num_buckets: int = NUM_BUCKETS
@classmethod
def from_env(cls) -> Deployment:
"""从环境变量读取部署参数（systemd / 容器编排注入）。"""
return cls(
node_index=int(os.environ.get("TASK_NODE_INDEX", "0")),
node_count=int(os.environ.get("TASK_NODE_COUNT", "1")),
db_path=Path(os.environ.get("TASK_DB", "tasks.db")),
)
@property
def owner(self) -> str:
"""写进存储的节点身份。排障时能回答"是谁在管这个任务"。"""
return f"node-{self.node_index}"
@property
def owned_buckets(self) -> range:
"""本节点负责的桶。三个地方必须用同一份结果：Controller 参数、CLI 参数、日志。"""
return buckets_for_node(self.node_index, self.node_count, self.num_buckets)
@property
def cli_buckets(self) -> str:
"""同一个划分的 CLI 写法（``--owned-buckets 0-7``），供运维直接复制。"""
owned = self.owned_buckets
return f"{owned.start}-{owned.stop - 1}"
def owner_of(bucket: int, deployments: Sequence[Deployment]) -> Deployment:
"""桶号反查节点。查不到意味着划分有空洞或配置不一致——直接炸，不静默。"""
for deployment in deployments:
if bucket in deployment.owned_buckets:
return deployment
raise AssertionError(
f"没有任何节点负责桶 {bucket}：请检查 num_buckets / node_count 是否各处一致"
)
# =============================================================================
# 第二部分：任务定义（业务代码：只依赖 ctx 约定，照抄进你的 tasks 模块）
# =============================================================================
@task(
tags={"sync", "db"},
retry=RetryPolicy(max_attempts=2, backoff=ExponentialBackoff(base=0.05, factor=2)),
)
def sync_tenant(ctx: TaskContext, tenant: str, batches: int = 24) -> dict[str, object]:
"""同步一个租户的数据。
标签 ``db`` 不只是检索用的：它还是**具名资源信号量**的键——每个节点到库里的
并发连接数是有限的，靠它限流比调小全局并发更准（见 ``build_node``）。
长耗时循环里必须周期性调用 ``ctx.checkpoint()``——停止、暂停、心跳都在那里生效；
不调用它的任务既停不掉，也会被 idle 超时判定为疑似卡死。
"""
ctx.log(f"{tenant}：开始同步（{batches} 批）")
done = 0
for index in range(1, batches + 1):
if ctx.is_stop_requested:
# 优雅收尾：正常 return -> 终态 succeeded，返回值保留（"请收工"语义）。
# 若希望"停止即中止"，去掉这个检查、让 checkpoint() 抛 TaskCancelled
# -> 终态 stopped、无返回值。两种都如实上报，见示例 09。
break
ctx.checkpoint()
time.sleep(0.05)
done = index
ctx.progress(done, total=batches, unit="批", message=tenant)
ctx.log(f"{tenant}：本次同步 {done}/{batches} 批")
return {"tenant": tenant, "batches": done}
# =============================================================================
# 第三部分：接线 —— build_node()（唯一集成点，应用启动时调用一次）
# =============================================================================
def register(controller: Controller) -> None:
"""把任务定义登记进白名单。
两种用法共用同一份定义：应用嵌入时由 :func:`build_node` 调用；
CLI 用 ``--tasks 本模块`` 时，由 CLI 调用（``register(controller)``
正是 CLI 的约定函数名）。
"""
controller.register_all(sync_tenant)
def build_node(deployment: Deployment) -> Controller:
"""建好一个节点：存储 + 身份 + 桶归属 + 并发上限 + 任务定义 + 启动恢复。
这是应用**唯一**需要接进本库的地方：你的服务（Web / 定时器 / 消费者）启动时
调用一次，把返回的 ``Controller`` 持有到进程结束；退出前 ``close(close_store=True)``。
"""
controller = Controller(
store=SqliteTaskStore(deployment.db_path),
owner=deployment.owner,
bucket_policy=BucketPolicy(num_buckets=deployment.num_buckets),
# 桶归属 = 划分的落地处：本节点只接受、只清理这些桶里的任务。
owned_buckets=deployment.owned_buckets,
# 并发与资源约束集中在这一处（ConcurrencyPolicy）：
# max_parallel=4 —— 本进程同时最多执行 4 个；
# limits={"db": 3} —— 其中最多 3 个可以同时压数据库（带 db 标签的任务），
# 保护的是"本节点到库里的连接数"这份预算；
# priority=... —— 默认优先级，submit(priority=...) 可逐条覆盖。
concurrency=ConcurrencyPolicy(max_parallel=4, limits={"db": 3}),
max_parallel_per_bucket=1, # 同一个租户（同桶）串行，避免自己挤自己
command_poll_interval=0.1, # 仅为让演示输出得快；生产用默认 0.5s 即可
)
register(controller)
# 启动恢复：把上次进程留下的、占着资源却没人管的僵尸任务如实标为 lost。
# 只看**自己负责的桶**——别的节点的任务由它自己的 reconcile 处理。
report = controller.reconcile()
owned = deployment.owned_buckets
_log.info(
"[%s] 启动：负责桶 %d-%d；reconcile：%d 个判失联，%d 个待重新提交",
deployment.owner,
owned.start,
owned.stop - 1,
len(report.lost),
len(report.adoptable),
)
return controller
# =============================================================================
# 第四部分：演示（启动 -> 提交 -> 观测 -> 控制 -> 收尾；这一段不用照抄）
# =============================================================================
#: 演示用的租户（业务键）。生产上它来自你的业务清单。
TENANTS = ("tenant-acme", "tenant-umbrella", "tenant-globex", "tenant-hooli")
#: 给"跨节点停止"演示留足运行窗口：80 批 × 0.05s ≈ 4 秒。
SLOW_BATCHES = {"tenant-globex": 80}
#: 模拟"上一个进程崩溃"遗留的僵尸任务的业务键（bucket 7，归 node-0）。
ZOMBIE_KEY = "tenant-vandelay"
def _plant_zombie(db_path: Path, policy: BucketPolicy, key: str) -> str:
"""往存储里放一条"崩溃进程留下的"任务：RUNNING 且租约已过期。
真实场景里这是崩溃的进程自己留下的（写穿 + 租约续不上）；这里手工造一条，
让 reconcile 有东西可清、让"如实标记而不是假装恢复"看得见。
"""
now = utcnow()
task_id = f"sync-{key}"
snapshot = TaskSnapshot(
task_id=task_id,
name=f"同步-{key}",
bucket=policy.bucket_for(key),
bucket_key=key,
status=TaskStatus.RUNNING,
status_version=1,
attempt=1,
max_attempts=1,
created_at=now - timedelta(minutes=12),
queued_at=now - timedelta(minutes=12),
started_at=now - timedelta(minutes=11),
last_heartbeat_at=now - timedelta(minutes=11),
executor="thread",
lease_owner="node-0@dead-process",
lease_until=now - timedelta(minutes=6),
)
store = SqliteTaskStore(db_path)
try:
store.create(snapshot)
finally:
store.close()
return task_id
def _print_plan(deployments: Sequence[Deployment], policy: BucketPolicy, db_path: Path) -> None:
"""打印部署计划、桶划分与路由表——集成前先把"谁负责什么"看清楚。"""
print("=" * 72)
print(f"部署：{NUM_BUCKETS} 个桶，{NODE_COUNT} 个节点，共享一套存储（{db_path.name}）")
print("=" * 72)
print(" 节点 负责桶 等价的 CLI 参数")
print(" -------- ---------- -----------------------------------")
for deployment in deployments:
owned = deployment.owned_buckets
spec = f"--owned-buckets {deployment.cli_buckets}"
print(f" {deployment.owner:<8} {owned.start:>2}-{owned.stop - 1:<8} {spec}")
print(" 注意：--num-buckets 也必须一致，否则同一个键会落到不同的桶，且无任何报错。")
print("\n 路由表（blake2b 稳定哈希：同一个键在任何进程、任何机器上都落到同一个桶）：")
for tenant in TENANTS:
bucket = policy.bucket_for(tenant)
print(f" {tenant:<18} -> 桶 {bucket:>2} -> {owner_of(bucket, deployments).owner}")
def main() -> int:
# 输出到 stdout（默认是 stderr）：这样日志与 print 的顺序一致，
# 管道重定向时也不会看到"日志跑到前面去了"。
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s", stream=sys.stdout)
with tempfile.TemporaryDirectory(prefix="chongming-best-practice-") as directory:
db_path = Path(directory) / "tasks.db"
deployments = [
Deployment(node_index=index, node_count=NODE_COUNT, db_path=db_path)
for index in range(NODE_COUNT)
]
policy = BucketPolicy(num_buckets=NUM_BUCKETS)
_print_plan(deployments, policy, db_path)
zombie_id = _plant_zombie(db_path, policy, ZOMBIE_KEY)
zombie_bucket = policy.bucket_for(ZOMBIE_KEY)
assert zombie_bucket in deployments[0].owned_buckets, "僵尸任务应落在 node-0 的桶里"
print(
f"\n[准备] 模拟上一个进程崩溃：存储里留有一条 RUNNING 且租约过期的任务 "
f"{zombie_id}（桶 {zombie_bucket}，归 node-0）"
)
# ---------------------------------------------------------------- 启动
print("\n" + "=" * 72)
print("一、启动：每个节点先 reconcile()，只清理**自己桶里**的僵尸任务")
print("=" * 72)
nodes = [build_node(deployment) for deployment in deployments]
# 运维控制台：一个不持有任务的 Controller，只查询与控制（不声明 owned_buckets）。
console = Controller(store=SqliteTaskStore(db_path), owner="ops-console")
zombie = console.get(zombie_id)
assert zombie is not None and zombie.status is TaskStatus.LOST
print(
f" 控制台看到 {zombie_id} 现在是 {zombie.status.value}"
"（被如实标记为丢失，而不是永远显示 running；它不会被自动重跑——"
"存储里只有状态、没有配方）"
)
# ---------------------------------------------------------------- 提交
print("\n" + "=" * 72)
print("二、触发：按桶把任务交给拥有它的节点（配错桶 = 永远不会有人执行）")
print("=" * 72)
for tenant in TENANTS:
bucket = policy.bucket_for(tenant)
node = nodes[owner_of(bucket, deployments).node_index]
run = node.submit(
"sync_tenant",
tenant,
task_id=f"sync-{tenant}",
label=f"同步-{tenant}",
bucket_key=tenant,
batches=SLOW_BATCHES.get(tenant, 24),
)
print(f" {tenant:<18} 桶 {bucket:>2} -> {node.owner} 提交 {run.snapshot.task_id}")
foreign_tenant = next(
tenant
for tenant in TENANTS
if owner_of(policy.bucket_for(tenant), deployments).node_index != 0
)
foreign_bucket = policy.bucket_for(foreign_tenant)
foreign_owner = owner_of(foreign_bucket, deployments).owner
print(
f"\n >>> 反例：把 {foreign_tenant}（桶 {foreign_bucket}，归 {foreign_owner}）"
"提交给 node-0："
)
try:
nodes[0].submit(
"sync_tenant",
foreign_tenant,
task_id="sync-misrouted",
bucket_key=foreign_tenant,
)
except ConfigurationError as exc:
print(f" ✅ 提交那一刻就拦住了：{exc}")
else:
raise AssertionError("提交给错误的节点居然成功了——桶归属校验失效")
# ---------------------------------------------------------------- 观测
time.sleep(0.3)
print("\n" + "=" * 72)
print("三、观测：一套存储 = 一个看板看全部；按桶过滤 = 只看我这一份")
print("=" * 72)
print("--- 全局视图（控制台不在任何节点上，也能看到全部任务）---")
print(render_table(console.list_tasks()))
owned_0 = deployments[0].owned_buckets
print(f"\n--- node-0 的视角（list_tasks(buckets={owned_0.start}-{owned_0.stop - 1})）---")
print(render_table(nodes[0].list_tasks(buckets=owned_0)))
for node in nodes:
print(f" {node.owner}：占用执行额度 {node.active_count}，排队 {node.queued_count}")
# ---------------------------------------------------------------- 控制
target_id = f"sync-{foreign_tenant}"
target = console.get(target_id)
assert target is not None and not target.is_terminal, "演示假设：停止目标仍在运行"
print("\n" + "=" * 72)
print("四、控制：跨节点停止（不能跨进程发信号 -> 写一条指令，由持有者取走）")
print("=" * 72)
print(
f" ops-console 请求停止 {target_id}（持有者是 {foreign_owner}）："
f"下发结果 {console.stop(target_id, reason='发布窗口')}"
)
outcome = nodes[1].wait(target_id, timeout=15)
print(f" 终态：{outcome.status.value}；返回值：{outcome.value}")
print(" 说明：任务体看到停止请求后优雅收尾 -> succeeded 且数值保留；")
print(" 让 checkpoint() 抛 TaskCancelled 则为 stopped。两种都如实上报。")
print("\n 审计轨迹（从共享存储读；谁按的、为什么，都可追溯）：")
for record in console.store.history(target_id): # type: ignore[union-attr]
reason = f" reason={record.reason}" if record.reason else ""
print(
f" {record.version:>2} {record.from_status.value:>9} -> "
f"{record.to_status.value:<9} {record.trigger}{reason}"
)
# ---------------------------------------------------------------- 收尾
for node in nodes:
node.wait_all(timeout=30)
print("\n" + "=" * 72)
print("五、收尾：校验一致性 -> close(close_store=True)")
print("=" * 72)
print(render_table(console.list_tasks()))
stats = console.stats()
print(f" 统计：共 {stats.total}，未结束 {stats.unfinished}，已结束 {stats.finished}")
for tenant in TENANTS:
snapshot = console.get(f"sync-{tenant}")
assert snapshot is not None and snapshot.status in (
TaskStatus.SUCCEEDED,
TaskStatus.STOPPED,
), f"{tenant} 的终态不对：{snapshot.status if snapshot else '不存在'}"
for node in nodes:
assert node.persist_failures == (), f"{node.owner} 有写穿失败——存储不可信"
assert node.desynced_tasks == {}, f"{node.owner} 有永久分叉的任务"
print(" ✅ 结论与存储一致：没有写穿失败，没有永久分叉，僵尸任务如实落 lost")
for node in nodes:
node.close(timeout=30, close_store=True)
console.close(timeout=5, close_store=True)
_print_ops_notes()
print("\n全部通过 ✅")
return 0
def _print_ops_notes() -> None:
"""部署与运维说明。"""
print(
f"""
生产部署：把它变成一个真正的服务
────────────────────────────────────────────────────────────────────────
# 每台机器 / 每个容器跑一个节点进程：同一份代码，只有 NODE_INDEX 不同
TASK_NODE_INDEX=0 TASK_NODE_COUNT={NODE_COUNT} TASK_DB=/srv/tasks/tasks.db \\
python -m your_service.node
TASK_NODE_INDEX=1 TASK_NODE_COUNT={NODE_COUNT} TASK_DB=/srv/tasks/tasks.db \\
python -m your_service.node
# 运维（任意一台机器；--num-buckets 与代码里的 NUM_BUCKETS 必须一致）
uv run chongming-task --db /srv/tasks/tasks.db --num-buckets {NUM_BUCKETS} \\
ps --all
uv run chongming-task --db /srv/tasks/tasks.db --num-buckets {NUM_BUCKETS} \\
--owned-buckets 0-7 ps --all --mine
uv run chongming-task --db /srv/tasks/tasks.db --num-buckets {NUM_BUCKETS} \\
--owned-buckets 0-7 reconcile --dry-run
uv run chongming-task --db /srv/tasks/tasks.db --num-buckets {NUM_BUCKETS} \\
stop sync-tenant-globex --reason "发布窗口"
# 本模块也能直接当 CLI 任务模块用（register(controller) 约定）
PYTHONPATH=examples uv run chongming-task --db ./demo.db --num-buckets {NUM_BUCKETS} \\
--owned-buckets 0-7 run --tasks 11_best_practice sync_tenant \\
--arg tenant=tenant-acme --task-id sync-tenant-acme --bucket-key tenant-acme
集成三步（照抄本文件的三个"部分"）
────────────────────────────────────────────────────────────────────────
# 1. 第一部分 + 第二部分 -> 你的配置模块与任务模块；
# 2. 应用启动时 build_node(Deployment.from_env())，退出前 close(close_store=True)；
# 3. 触发处 submit(任务名, ..., bucket_key=业务键)。
# 两条硬约束：任务只能提交给拥有它那个桶的节点；bucket_key 必须给
# （限定了 owned_buckets 时不给会被当场拒绝——没有归属桶的任务没人认领）。
#
# 跨机部署：SQLite 文件不能跨机共享，把第三部分的 store 换成
# PostgresTaskStore("postgresql://user:pw@db.internal:5432/chongming")
# （pip install "chongming-task[postgres]"）即可，其余部分不用改。""".rstrip()
)
if __name__ == "__main__":
sys.exit(main())
./examples\12_process_executor.py
"""示例 12：进程执行器 —— 硬性停止、绕开 GIL、崩溃隔离。
uv run python examples/12_process_executor.py
绝大多数任务该用默认的**线程执行器**：共享内存、没有 pickle 约束、能直接复用连接池。
但有三件事线程做不到，只有子进程能：
1. **硬性停止保证**：任务不响应停止（不调 ``ctx.checkpoint()``）时，线程只能放弃等待
并把结果记成 ``lost``（未知）；子进程会被 ``terminate()`` → ``kill()`` 真的停掉；
2. **绕开 GIL**：CPU 密集的计算能真正并行；
3. **崩溃隔离**：任务里加载的原生扩展段错误，只会带走子进程。
本示例把 1 和 3 跑给你看，并把代价说清楚。
这个文件本身也是"目标函数必须是模块级"的演示：下面的任务都定义在模块顶层，因为
spawn 起的子进程要通过"模块 + 函数名"把它们找回来。
"""
from __future__ import annotations
import logging
import os
import sys
import time
from typing import Any
from chongming_task import (
ConfigurationError,
FailureKind,
ProcessExecutor,
RetryPolicy,
TaskStatus,
Timeout,
run_task,
)
from chongming_task.runtime.executors.base import ExecutionRequest
# 子进程不继承父进程的 logging 配置（spawn 起的是全新解释器），日志想看就得自己配。
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stdout)
# -----------------------------------------------------------------------------
# 任务定义（模块级：spawn 的子进程要靠"模块 + 函数名"找回来）
# -----------------------------------------------------------------------------
def cpu_burn(ctx: Any, rounds: int = 6) -> dict[str, object]:
"""CPU 密集的活：子进程里不受父进程 GIL 的影响。"""
total = 0
for _ in range(rounds):
ctx.checkpoint()
for index in range(400_000):
total += index % 7
return {"rounds": rounds, "checksum": total, "pid": os.getpid()}
def deaf_loop(ctx: Any) -> str:
"""**不**调用 checkpoint 的任务：停止请求它根本看不见。
这正是不配合的任务的样子——拿线程执行器它只会落 ``lost``；拿进程执行器它会被真的杀掉。
"""
del ctx
end = time.monotonic() + 60
while time.monotonic() < end:
pass
return "never"
def crash_hard(ctx: Any) -> str:
"""模拟原生扩展段错误：直接消失，一句话都不留。"""
del ctx
os._exit(139)
return "never" # pragma: no cover - 走不到这里
# -----------------------------------------------------------------------------
# 演示
# -----------------------------------------------------------------------------
def demo_cpu_bound() -> None:
print("=" * 72)
print("一、CPU 密集：任务跑在子进程里（父进程的 GIL 管不着它）")
print("=" * 72)
outcome = run_task(
cpu_burn,
6,
task_id="cpu-1",
executor=ProcessExecutor(),
timeout=Timeout(per_attempt=30),
)
assert outcome.ok, outcome.failure
payload = outcome.value
print(f" 结果：{payload}")
print(f" 监督进程 pid={os.getpid()}，任务进程 pid={outcome.snapshot.pid}")
print(f" 快照里的 executor={outcome.snapshot.executor}（写进存储，事后可查）")
print()
def demo_hard_stop() -> None:
print("=" * 72)
print("二、硬性停止：不响应停止的任务被真的杀掉（线程执行器只能放弃等待）")
print("=" * 72)
started = time.monotonic()
outcome = run_task(
deaf_loop,
task_id="deaf-1",
executor=ProcessExecutor(),
timeout=Timeout(idle=1), # 一秒没有心跳就判卡死
retry=RetryPolicy(max_attempts=1),
graceful_stop=0.5, # 优雅窗口：0.5 秒内不响应就强杀
)
elapsed = time.monotonic() - started
print(f" 任务自己的循环要跑 60 秒，实际用了 {elapsed:.1f} 秒")
print(f" 终态：{outcome.status.value}；失败类型：{outcome.failure.kind.value}")
print(f" 子进程 pid={outcome.snapshot.pid} —— 它已经不在了（ps 查不到）")
assert outcome.status is TaskStatus.FAILED
assert outcome.failure.kind is FailureKind.TIMEOUT_IDLE, "超时如实上报，不谎报'已停止'"
assert elapsed < 30, "不该等到任务自己结束"
print(" 注意：超时仍然记成超时（不是'已停止'）——但资源确实释放了。")
print()
def demo_crash_isolation() -> None:
print("=" * 72)
print("三、崩溃隔离：子进程崩了，监督进程活着并如实记一笔")
print("=" * 72)
outcome = run_task(
crash_hard,
task_id="crash-1",
executor=ProcessExecutor(),
timeout=Timeout(per_attempt=30),
)
assert not outcome.ok
print(f" 监督进程还活着（pid={os.getpid()}），任务终态：{outcome.status.value}")
print(f" 失败类型：{outcome.failure.kind.value}；退出码：{outcome.failure.exit_code}")
print(f" 说明：{outcome.failure.message}")
print(" 换成线程执行器，这个崩溃会带走整个监督进程。")
print()
def demo_picklable_guard() -> None:
print("=" * 72)
print("四、代价：不可 pickle 的目标在**提交那一刻**就被拦下")
print("=" * 72)
def local_closure(ctx: Any) -> str: # 局部函数：子进程里找不到它
del ctx
return "nope"
# 1) 直接交给执行器：当场抛错，并告诉你该改成什么。
request = ExecutionRequest(task_id="bad-0", name="bad", target=local_closure, args=())
try:
ProcessExecutor().launch(request)
except ConfigurationError as exc:
print(f" launch() 直接报错 {type(exc).__name__}：")
for line in str(exc).splitlines():
print(f" {line}")
else: # pragma: no cover - 应当抛错
raise AssertionError("不可 pickle 的目标必须被拦下")
# 2) 走 run_task：任务**落 failed 并带上这条诊断**，调用方不会干等，也不会被
# 说成"框架内部错误"——那是把任务侧的问题指向错误的排查方向。
outcome = run_task(local_closure, task_id="bad-1", executor=ProcessExecutor())
assert outcome.status is TaskStatus.FAILED
assert outcome.failure is not None
print(f"\n 走 run_task：终态 {outcome.status.value}，retryable={outcome.failure.retryable}")
print(f" 失败信息首行：{outcome.failure.message.splitlines()[0]}")
print(" 配置类错误不重试——重试多少次都是同一个结果，只会白烧重试预算。")
assert outcome.failure.retryable is False
print()
def main() -> int:
print()
demo_cpu_bound()
demo_hard_stop()
demo_crash_isolation()
demo_picklable_guard()
print("=" * 72)
print("选型：默认线程，需要这三种能力时才换进程")
print("=" * 72)
print(
"""
线程（默认） 进程（ProcessExecutor）
------------------------ ------------------------------------------
共享内存、可复用连接池 参数与返回值必须可 pickle
闭包 / lambda 都能用 目标必须是模块级函数，调用方要有 __main__ 守卫
停止/暂停即时生效 停止/暂停有约 20ms 轮询延迟
不配合的任务只能放弃（lost） 不配合的任务被真的杀掉
受 GIL 限制 真并行
原生扩展段错误带走整个进程 只带走子进程
"""
)
print("全部通过 ✅")
return 0
if __name__ == "__main__":
sys.exit(main())
./examples\tasks_demo.py
"""给 CLI 用的任务模块。
CLI **只传任务名**，定义必须来自这样一个模块——一个能启动任意可调用对象的命令行
入口等于任意代码执行，所以"任务定义从哪里来"必须先有答案。约定就是本文件里的
``register(controller)``。
用法（把本目录加进 ``PYTHONPATH``，或把模块放进你自己的包里）::
chongming-task --db /tmp/tasks.db tasks --tasks tasks_demo
chongming-task --db /tmp/tasks.db run --tasks tasks_demo add --arg left=2 --arg right=3
"""
from __future__ import annotations
import time
from typing import Any
def sync_tenant(ctx: Any, tenant: str, batches: int = 120) -> str:
"""同步一个租户的数据。"""
for index in range(batches):
ctx.checkpoint()
if index % 20 == 0:
time.sleep(0.05)
return f"{tenant} 同步完成（{batches} 批）"
def add(ctx: Any, left: int, right: int) -> int:
"""两数相加。"""
del ctx
return left + right
def flaky(ctx: Any, fail_times: int = 2) -> str:
"""前几次故意失败，用来演示重试。"""
state = getattr(flaky, "_attempts", 0) + 1
flaky._attempts = state # type: ignore[attr-defined]
if state <= fail_times:
raise RuntimeError(f"第 {state} 次尝试遇到暂时性故障")
return f"第 {state} 次尝试成功"
def register(controller: Any) -> None:
"""把任务登记进白名单。CLI 导入本模块后就会调用它。"""
controller.register("sync", sync_tenant, tags={"io"})
controller.register("add", add, tags={"cpu"})
controller.register("flaky", flaky)
./examples\worker_tasks.py
"""一个可以直接用的常驻任务模块（CLI 的 ``--tasks`` 就指向它）。
它是给命令行用的示例模块，也是你自己模块的模板。用法::
PYTHONPATH=examples uv run chongming-task --db ./tasks.db \
run --tasks worker_tasks consume --arg source=kafka://orders --task-id orders-worker
## 三件必须做对的事
1. **每轮调 ``ctx.checkpoint()``**——停止是协作式的，不调用它就停不掉（会落 ``lost``）。
2. **用 ``ctx.is_stop_requested`` 收尾**——这样任务终态是 ``succeeded``，返回值保留；
让 ``checkpoint()`` 抛则是 ``stopped``、无返回值。
3. **进度按批次上报，不要每条一次**——``ctx.progress()`` 每次都要取运行锁，它是给
"人看得懂"的频率设计的。
"""
from __future__ import annotations
import time
from typing import Any
from chongming_task import RetryPolicy, task
@task(tags={"kafka"}, retry=RetryPolicy(max_attempts=1))
def consume(ctx: Any, source: str, batch: int = 1000) -> int:
"""常驻消费 source；收到停止请求后收尾并返回本次共处理多少条。"""
ctx.log(f"开始消费 {source}")
total = 0
while not ctx.is_stop_requested:
ctx.checkpoint()
# 假装取了一批消息。真实实现里这里是一次 poll / 一次拉取。
time.sleep(0.05)
total += batch
# 按批次上报进度（不是每条一次）。
ctx.progress(total, unit="msg", message=source)
ctx.log(f"收到停止请求，收尾：本次共处理 {total} 条")
return total
@task
def once(ctx: Any, value: int) -> int:
"""一个短任务，用来对照"有始有终"的普通形态。"""
del ctx
return value * 2
./myapp\tasks.py
# myapp/tasks.py
from chongming_task import task
@task(tags={"io"})
def worker(ctx, source: str) -> int:
"""常驻消费。"""
processed = 0
while not ctx.is_stop_requested: # ← 优雅收尾（终态 succeeded，返回值保留）
ctx.checkpoint() # ← 不调用它 = 停不掉
processed += 1
ctx.progress(processed, unit="batch")
return processed
def register(controller): # 或者不写它，只靠 @task 也行
controller.register_module(__import__(__name__))
./myapp\__init__.py
./pytest-of-chakcy\pytest-77\test_a_module_found_but_failin0\broken_tasks.py
import no_such_dependency_here
./pytest-of-chakcy\pytest-77\test_a_module_with_only_task_d0\decorated_only.py
from chongming_task import task
@task(tags={'io'})
def doubled(ctx, value: int) -> int:
'''把数字翻倍。'''
del ctx
return value * 2
./pytest-of-chakcy\pytest-92\test_a_module_found_but_failin0\broken_tasks.py
import no_such_dependency_here
./pytest-of-chakcy\pytest-92\test_a_module_with_only_task_d0\decorated_only.py
from chongming_task import task
@task(tags={'io'})
def doubled(ctx, value: int) -> int:
'''把数字翻倍。'''
del ctx
return value * 2
./pytest-of-chakcy\pytest-93\test_a_module_found_but_failin0\broken_tasks.py
import no_such_dependency_here
./pytest-of-chakcy\pytest-93\test_a_module_with_only_task_d0\decorated_only.py
from chongming_task import task
@task(tags={'io'})
def doubled(ctx, value: int) -> int:
'''把数字翻倍。'''
del ctx
return value * 2
./pytest-of-chakcy\pytest-94\test_a_module_found_but_failin0\broken_tasks.py
import no_such_dependency_here
./pytest-of-chakcy\pytest-94\test_a_module_with_only_task_d0\decorated_only.py
from chongming_task import task
@task(tags={'io'})
def doubled(ctx, value: int) -> int:
'''把数字翻倍。'''
del ctx
return value * 2
./pytest-of-chakcy\pytest-95\test_a_module_found_but_failin0\broken_tasks.py
import no_such_dependency_here
./pytest-of-chakcy\pytest-95\test_a_module_with_only_task_d0\decorated_only.py
from chongming_task import task
@task(tags={'io'})
def doubled(ctx, value: int) -> int:
'''把数字翻倍。'''
del ctx
return value * 2
./src\chongming_task\clock.py
"""时钟抽象。
时间必须可注入：重试退避、超时判定、心跳陈旧度这些逻辑一旦依赖真实时间，
测试就只能靠 ``sleep``，既慢又不稳定。因此运行时的每一处取时都走 :data:`Clock`。
生产用 :func:`utcnow`；测试用 ``chongming_task.testing.clock.FakeClock``。
"""
from __future__ import annotations
from collections.abc import Callable
from datetime import UTC, datetime
__all__ = ["Clock", "utcnow"]
def utcnow() -> datetime:
"""默认取时，返回带 UTC 时区的当前时间。"""
return datetime.now(UTC)
#: 时钟。返回带时区的当前时间。
#:
#: 注意：这里取的是**墙钟**，可能因 NTP 校时发生跳变。超时判定基于两个时刻之差，
#: 若对跳变敏感，可注入基于 ``time.monotonic`` 的自定义实现（代价是失去可读时间戳）。
Clock = Callable[[], datetime]
./src\chongming_task\controller.py
"""任务管理门面：注册、提交、查询、控制。
这是管理后台（CLI，将来可能的看板）**唯一需要对接的层**。它统一持有三样东西：
* **任务定义注册表**——启动时把可调用对象登记成名字，之后一切提交都只传名字；
* **运行中的任务**——``task_id -> TaskRun``，控制指令据此路由；
* **并发额度**——超出上限的提交会排队，状态为 ``queued`` 而不是立刻起线程。
安全底线：**只允许启动已注册的任务，绝不接受代码。** 一个能启动任意可调用对象的
入口等于任意代码执行——将来无论接 CLI 还是 HTTP，这条都不能破。
跨进程说明
----------
本类管理的是**当前进程内**的任务。另一个进程里的 ``chongming-task ps`` 看不到它们，
因为内存不共享。跨进程查询与持久化要等共享存储（``sqlite3`` 属于标准库，不引入
第三方依赖）落地——那正是 CLI 的媒介。
"""
from __future__ import annotations
import os
import socket
import threading
import time
import unicodedata
import uuid
from collections import deque
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass, field
from datetime import UTC, datetime, timedelta
from enum import StrEnum
from pathlib import Path
from typing import Any
from .clock import Clock, utcnow
from .core.state_machine import StateMachine, Transition
from .decorators import get_task_spec
from .errors import ConcurrencyConflict, ConfigurationError, InvalidTransition, StoreError
from .models import Failure, FailureKind, Progress, TaskSnapshot
from .pipeline import Pipeline, PipelineRun, start_pipeline
from .policies import BucketPolicy, ConcurrencyPolicy, RetryPolicy, Timeout
from .run import TaskRun, TaskRunOutcome
from .runtime.executors.base import Executor
from .runtime.executors.thread import ThreadExecutor
from .states import NON_TERMINAL, RESOURCE_HOLDING, TERMINAL, TaskStatus
from .stores.protocol import Command, CommandKind, TaskQuery, TaskStore
__all__ = [
"Controller",
"ReconcileEntry",
"ReconcileOutcome",
"ReconcileReport",
"Registration",
"StatusCounts",
"TaskRegistry",
"render_table",
]
#: 表格列之间的分隔。
_COLUMN_SEPARATOR = " "
def _validated_tags(tags: Iterable[str]) -> frozenset[str]:
"""把标签收成集合，并拒绝非字符串元素。
标签会被写进存储（JSON 数组），所以非字符串元素既不能正确往返、又只会在**写穿**
路径上才炸——而那条路径的失败不会让任务失败，问题会以"看板上的标签不对"的形式
出现。因此按真实类型契约（``frozenset[str]``）在入口拒绝。
注册与提交两条路径都要过这里：``submit`` 会把注册时的标签与本次传入的合并，
只校验其中一条会漏。
"""
resolved = frozenset(tags)
bad = sorted(repr(tag) for tag in resolved if not isinstance(tag, str))
if bad:
raise ConfigurationError(f"标签必须是字符串，收到：{', '.join(bad)}")
return resolved
#: 每个任务最多缓冲多少条未落库的迁移。
#:
#: 必须有界：存储长期不可用时，内存不能跟着无限增长。超出后丢弃最旧的——审计轨迹会
#: 出现缺口，所以丢弃这件事本身也会被记录（见 Controller._flush_pending）。
_PENDING_TRANSITIONS = 100
def _first_doc_line(target: Callable[..., Any]) -> str:
"""取函数 docstring 的第一行，作为注册项的默认描述。"""
doc = (target.__doc__ or "").strip()
return doc.splitlines()[0] if doc else ""
@dataclass(frozen=True, slots=True)
class Registration:
"""一个已注册的任务定义。"""
name: str
target: Callable[..., Any]
retry: RetryPolicy
timeout: Timeout
tags: frozenset[str] = frozenset()
description: str = ""
class TaskRegistry:
"""任务定义的白名单。
只按**名字**提交任务，因此运行期的任何控制面都不需要、也无法传递代码。
"""
def __init__(self) -> None:
self._items: dict[str, Registration] = {}
self._lock = threading.RLock()
def register(
self,
name: str,
target: Callable[..., Any],
*,
retry: RetryPolicy | None = None,
timeout: Timeout | None = None,
tags: Iterable[str] = (),
description: str = "",
) -> Registration:
"""登记一个任务定义。
重名会抛 :class:`ConfigurationError` 而不是静默覆盖——静默覆盖会让"我明明
改了实现，跑的还是旧的"这类问题变得极难发现。
"""
if not name or not name.strip():
raise ConfigurationError("任务名不能为空")
if not callable(target):
raise ConfigurationError(f"任务 {name!r} 的 target 必须是可调用对象")
resolved_tags = _validated_tags(tags)
with self._lock:
if name in self._items:
raise ConfigurationError(
f"任务名 {name!r} 已注册；如需替换请先注销（当前不支持静默覆盖）"
)
registration = Registration(
name=name,
target=target,
retry=retry if retry is not None else RetryPolicy(),
timeout=timeout if timeout is not None else Timeout(),
tags=resolved_tags,
description=description or _first_doc_line(target),
)
self._items[name] = registration
return registration
def register_module(self, module: Any) -> tuple[Registration, ...]:
"""把一个模块里所有被 ``@task`` 装饰过的函数登记进来。
显式调用，而不是"import 时自动注册"：后者会让导入产生副作用，而 import 可能
发生在任何地方（包括测试收集阶段），那种隐式行为极难排查。
"""
registered: list[Registration] = []
seen: set[str] = set()
for member in vars(module).values():
if not callable(member):
continue
spec = get_task_spec(member)
if spec is None:
continue
name = spec.name if spec.name is not None else getattr(member, "__name__", None)
if not name or name in seen:
continue
seen.add(name)
registered.append(
self.register(
name,
member,
retry=spec.retry,
timeout=spec.timeout,
tags=spec.tags,
description=spec.description,
)
)
if not registered:
module_name = getattr(module, "__name__", repr(module))
raise ConfigurationError(f"模块 {module_name!r} 里没有用 @task 装饰过的函数")
return tuple(registered)
def register_decorated(self, target: Callable[..., Any]) -> Registration:
"""登记一个被 ``@task`` 装饰过的函数。"""
spec = get_task_spec(target)
if spec is None:
raise ConfigurationError(f"{getattr(target, '__name__', target)!r} 没有用 @task 装饰过")
name = spec.name if spec.name is not None else getattr(target, "__name__", None)
if not name:
raise ConfigurationError("无法确定任务名：请用 @task(name=...) 显式给出")
return self.register(
name,
target,
retry=spec.retry,
timeout=spec.timeout,
tags=spec.tags,
description=spec.description,
)
def get(self, name: str) -> Registration:
"""取一个已注册的定义；不存在则抛 :class:`ConfigurationError`。"""
with self._lock:
try:
return self._items[name]
except KeyError:
known = ", ".join(sorted(self._items)) or "(空)"
raise ConfigurationError(f"未注册的任务 {name!r}；已注册：{known}") from None
def names(self) -> tuple[str, ...]:
"""已注册的任务名。"""
with self._lock:
return tuple(sorted(self._items))
def all(self) -> tuple[Registration, ...]:
"""全部已注册定义。"""
with self._lock:
return tuple(self._items[name] for name in sorted(self._items))
def __contains__(self, name: object) -> bool:
with self._lock:
return name in self._items
def __len__(self) -> int:
with self._lock:
return len(self._items)
@dataclass(frozen=True, slots=True)
class StatusCounts:
"""各状态的计数，供看板首页使用。"""
by_status: Mapping[TaskStatus, int]
@property
def total(self) -> int:
"""任务总数。"""
return sum(self.by_status.values())
def of(self, status: TaskStatus) -> int:
"""某个状态的数量。"""
return self.by_status.get(status, 0)
@property
def unfinished(self) -> int:
"""尚未落终态的数量。
包含排队（``queued``）、暂停（``paused``）与重试等待（``retry_wait``）——
它们都不等于"正在跑"。要区分"真的在占用执行资源"，请用
:attr:`Controller.active_count`。
"""
return sum(count for status, count in self.by_status.items() if status not in TERMINAL)
@property
def finished(self) -> int:
"""已落终态的数量。"""
return sum(count for status, count in self.by_status.items() if status in TERMINAL)
class ReconcileOutcome(StrEnum):
"""重启恢复对单个任务的处置结果。"""
MARKED_LOST = "marked-lost"
"""占着执行资源却没有监督者——载体已随进程消失，判 ``LOST``。"""
ADOPTABLE = "adoptable"
"""没有执行载体（排队/暂停/退避中），只是没人接手。**不改状态**，需要重新提交。"""
SKIPPED = "skipped"
"""有人正在写这个任务，或状态机不允许该转移——不动它。"""
@dataclass(frozen=True, slots=True)
class ReconcileEntry:
"""一个任务的重启恢复结论。"""
task_id: str
previous_status: TaskStatus
outcome: ReconcileOutcome
detail: str
@dataclass(frozen=True, slots=True)
class ReconcileReport:
"""一次 :meth:`Controller.reconcile` 的完整结论。"""
entries: tuple[ReconcileEntry, ...] = ()
dry_run: bool = False
@property
def lost(self) -> tuple[ReconcileEntry, ...]:
"""被判为丢失的任务。"""
return tuple(item for item in self.entries if item.outcome is ReconcileOutcome.MARKED_LOST)
@property
def adoptable(self) -> tuple[ReconcileEntry, ...]:
"""无人接手、但保住状态的任务（排队/暂停/退避）。"""
return tuple(item for item in self.entries if item.outcome is ReconcileOutcome.ADOPTABLE)
@property
def skipped(self) -> tuple[ReconcileEntry, ...]:
"""被跳过的任务。"""
return tuple(item for item in self.entries if item.outcome is ReconcileOutcome.SKIPPED)
def __bool__(self) -> bool:
"""是否有任何需要关注的结论。"""
return bool(self.entries)
@dataclass(frozen=True, slots=True)
class _ActiveSlot:
"""一个正在占用并发额度的任务：它的桶与标签。
记桶是为了算每桶额度；记标签是为了算具名资源额度（``ConcurrencyPolicy.limits``）。
两者都必须记在**占用当时**——任务的标签与桶不会变，但重新从 ``_runs`` 里查一遍
会让"额度记账"依赖另一个数据结构的状态，得不偿失。
"""
bucket: int | None
tags: frozenset[str]
@dataclass(slots=True)
class _PersistState:
"""一个任务的写穿状态。
``last_persisted_version`` 是**我们上次确实写进存储的版本**——注意不是"上次算出来的
版本"。二者一旦分叉（写穿失败），用后者会导致此后每一次写入都因版本对不上而失败，
也就是一次瞬时故障把该任务的存储状态**永久冻结**。
"""
lock: threading.RLock = field(default_factory=threading.RLock)
last_persisted_version: int = 0
pending: deque[tuple[TaskSnapshot, Transition]] = field(default_factory=deque)
desynced: str | None = None
dropped: int = 0
@dataclass(slots=True)
class _Submission:
"""Controller 内部记录：任务运行 + 它的提交参数。"""
run: TaskRun
registration: Registration
args: tuple[Any, ...] = ()
kwargs: Mapping[str, Any] = field(default_factory=dict)
class Controller:
"""任务管理门面。"""
def __init__(
self,
*,
registry: TaskRegistry | None = None,
executor: Executor | None = None,
concurrency: ConcurrencyPolicy | None = None,
max_parallel: int | None = None,
max_parallel_per_bucket: int | None = None,
bucket_policy: BucketPolicy | None = None,
clock: Clock | None = None,
graceful_stop: float = 5.0,
poll: float = 0.05,
store: TaskStore | None = None,
owner: str | None = None,
command_poll_interval: float = 0.5,
command_lease: float = 30.0,
lease_ttl: float = 30.0,
owned_buckets: Iterable[int] | None = None,
) -> None:
"""``concurrency`` 是并发与资源约束的**唯一配置口**
（见 :class:`~chongming_task.policies.ConcurrencyPolicy`）：全局
``max_parallel``、按标签的具名信号量、默认优先级。
``max_parallel`` 是它的简写，给只调一个数字的场合用；两者同时给出会报错——
两份配置谁说了算必须明确，不能靠"后者覆盖前者"这种猜法。
``max_parallel_per_bucket`` 额外给每个桶各自限流：一个桶把额度占满不会饿死
其他桶。这是分桶在**本地**就有价值的地方，不必等到分布式。
``store`` 一旦给出，本实例同时成为**可跨进程的一环**：每次状态迁移写穿到
存储；查询走存储（因此能看到别的进程的任务）；停止/暂停/恢复对不属于本进程
的任务会写成一条指令，由持有者轮询取走。
``owned_buckets`` 是**本节点负责的桶**，也是"自治节点 + 统一平面"架构里唯一的
划分单位：每个节点有自己的任务定义与触发方式，但共用同一套存储与控制层，
靠桶区分"哪一份是我的"。
* ``None``（默认）= **不限定**，适合单机，或负责全部桶的节点；
* 给定集合 = 本节点只负责这些桶。此时：
- :meth:`submit` 只接受落在这些桶里的任务（配错桶会立刻报错，而不是产出一个
**永远没人认领**的任务）；
- :meth:`reconcile` 只清理这些桶里的僵尸任务，不去动别的节点的任务。
它的作用不是"抢占"，而是**划界**：节点之间不互相执行对方的任务。
"""
if concurrency is not None and max_parallel is not None:
raise ConfigurationError(
"concurrency 与 max_parallel 只能给一个：前者是完整的并发策略，"
"后者只是它的简写。两个都给会让'到底哪个生效'变成一个需要猜的问题。"
)
if max_parallel is not None and max_parallel < 1:
raise ConfigurationError(f"max_parallel 至少为 1，收到 {max_parallel}")
if max_parallel_per_bucket is not None and max_parallel_per_bucket < 1:
raise ConfigurationError(
f"max_parallel_per_bucket 至少为 1，收到 {max_parallel_per_bucket}"
)
self._registry = registry if registry is not None else TaskRegistry()
self._executor: Executor = executor if executor is not None else ThreadExecutor()
# 简写在这里被翻译成策略对象，之后调度器只认策略——不再有第二条读配置的路径。
if concurrency is not None:
self._concurrency = concurrency
else:
self._concurrency = ConcurrencyPolicy(max_parallel=max_parallel)
self._max_parallel_per_bucket = max_parallel_per_bucket
self._bucket_policy = bucket_policy if bucket_policy is not None else BucketPolicy()
self._owned_buckets = self._validated_owned_buckets(owned_buckets)
self._clock: Clock = clock if clock is not None else utcnow
self._graceful_stop = graceful_stop
self._poll = poll
self._runs: dict[str, _Submission] = {}
self._pending: deque[TaskRun] = deque()
self._active: dict[str, _ActiveSlot] = {}
"""占用并发额度的任务：``task_id -> 桶与标签``。
用映射而不是计数器：被 :meth:`close` 直接启动的排队任务从未计入额度，
若用计数器，它们结束时会把别人的名额一并扣掉。记下桶号是为了算每桶额度，
记下标签是为了算具名资源额度。
"""
self._closed = False
self._lock = threading.RLock()
self._wake = threading.Event()
self._dispatcher: threading.Thread | None = None
# ---- 跨进程（可选） ----
self._store = store
self._owner = owner if owner is not None else f"{socket.gethostname()}:{os.getpid()}"
self._command_poll_interval = command_poll_interval
self._command_lease = command_lease
self._lease_ttl = lease_ttl
self._persist_failures: list[str] = []
self._persist_states: dict[str, _PersistState] = {}
self._listener_failures: list[str] = []
self._listeners: list[Callable[[TaskRunOutcome], None]] = []
self._pipelines: dict[str, PipelineRun] = {}
self._worker: threading.Thread | None = None
if self._store is not None and self._owned_buckets is not None:
# 声明了"我负责哪些桶"= 这是一个真 worker，从启动就轮询指令。
#
# 否则会有一个延迟生效的坑：别人 `stop` 一个属于本节点、当前还排在队列里的
# 任务时，指令已经写进存储了，但本节点要等到**第一次提交**才会启动轮询线程
# ——对一个定时任务节点来说，那可能是几分钟甚至几小时之后。
#
# 没有声明归属的实例（纯查询客户端）仍然不轮询，避免抢走别人的指令。
self._ensure_worker()
@property
def owner(self) -> str:
"""本实例的身份，写进它下发的指令里，便于事后追溯是谁按的停止。"""
return self._owner
@property
def store(self) -> TaskStore | None:
"""绑定的存储；未接线时为 ``None``。"""
return self._store
@property
def persist_failures(self) -> tuple[str, ...]:
"""最近的持久化失败摘要（最多 20 条）。
存储故障**不该让任务失败**——任务本身没问题，坏的只是可观测性。但也绝不能
静默：写穿失败会记在这里，让"看板上的状态是旧的"这件事有迹可循。
"""
with self._lock:
return tuple(self._persist_failures)
@property
def listener_failures(self) -> tuple[str, ...]:
"""最近的监听器异常摘要（最多 20 条）。
监听器是外部代码（管道调度器就是其中之一）。它抛异常不该影响任务本身的结论，
但也绝不能静默——否则管道会停在半路而没有任何线索。
"""
with self._lock:
return tuple(self._listener_failures)
# ------------------------------------------------------------------
# 注册
# ------------------------------------------------------------------
@property
def registry(self) -> TaskRegistry:
"""任务定义注册表。"""
return self._registry
def register(
self,
name: str,
target: Callable[..., Any],
*,
retry: RetryPolicy | None = None,
timeout: Timeout | None = None,
tags: Iterable[str] = (),
description: str = "",
) -> Registration:
"""登记任务定义。"""
return self._registry.register(
name,
target,
retry=retry,
timeout=timeout,
tags=tags,
description=description,
)
def register_all(self, *targets: Callable[..., Any]) -> tuple[Registration, ...]:
"""登记若干被 ``@task`` 装饰过的函数，元数据从装饰器上取。"""
return tuple(self._registry.register_decorated(item) for item in targets)
def register_module(self, module: Any) -> tuple[Registration, ...]:
"""把一个模块里所有被 ``@task`` 装饰过的函数登记进来。
``@task`` 只把规格挂在函数对象上，**不产生全局注册表**——谁注册、注册到哪个
控制器由你决定，导入模块也不会有副作用。
"""
return self._registry.register_module(module)
# ------------------------------------------------------------------
# 提交
# ------------------------------------------------------------------
def submit(
self,
task: str,
/,
*args: Any,
task_id: str | None = None,
label: str | None = None,
tags: Iterable[str] = (),
bucket_key: str | None = None,
priority: int | None = None,
retry: RetryPolicy | None = None,
timeout: Timeout | None = None,
artifact_dir: Path | None = None,
pipeline_run_id: str | None = None,
deferred: bool = False,
**kwargs: Any,
) -> TaskRun:
"""提交一个**已注册**的任务。
返回 :class:`~chongming_task.run.TaskRun`，它自身就带 ``stop`` / ``pause``
/ ``resume`` / ``wait`` / ``snapshot``，因此句柄不需要另造一个类型。
超出并发额度（``ConcurrencyPolicy.max_parallel``）、桶额度
（``max_parallel_per_bucket``）或某个具名资源额度（``limits``）时，任务进入
队列，状态为 ``queued``——它是"已被受理、等待执行槽"，与 ``pending``
（尚未受理）不同。**排队不是失败**：额度一有富余，调度器就会把它取走。
``tags`` 除了用于筛选查询，还决定具名资源：带 ``gpu`` 标签的任务会占用
``limits["gpu"]`` 里的一个名额。
``priority`` 是排队顺序：数值越大越先被执行，同优先级按提交顺序。不给则用
``ConcurrencyPolicy.priority`` 的默认值。
``bucket_key`` 是分桶的业务键（租户 id、数据源、日期分区……）。同一个键在
任何进程、任何机器上都落到同一个桶，因此它既是本地并发隔离的单位，也是将来
分布式分配的单位——一台机器拥有一组桶。
``deferred=True`` 表示"先受理但**不要**排队执行"，任务停在 ``blocked``。
管道调度器用它创建下游节点：等上游满足后再 :meth:`_release`。
"""
registration = self._registry.get(task)
resolved_id = task_id if task_id is not None else uuid.uuid4().hex[:12]
if priority is not None and not isinstance(priority, int):
# 在这里拦下：非整数优先级会在调度器比较时抛 TypeError，而那个异常发生
# 在**派发线程**里——线程会直接死掉，之后所有排队任务都不再启动，且没有
# 任何报错指向真正的原因。
raise ConfigurationError(f"priority 必须是整数，收到 {priority!r}")
merged_tags = _validated_tags((*registration.tags, *tags))
bucket = self._bucket_for(bucket_key)
self._check_bucket_ownership(bucket, task_id=resolved_id, bucket_key=bucket_key)
with self._lock:
if self._closed:
raise ConfigurationError("Controller 已关闭，不能再提交任务")
if resolved_id in self._runs:
raise ConfigurationError(f"task_id 已存在：{resolved_id}")
run = TaskRun(
registration.target,
args,
kwargs=kwargs,
task_id=resolved_id,
name=label if label is not None else registration.name,
tags=merged_tags,
bucket=bucket,
bucket_key=bucket_key,
pipeline_run_id=pipeline_run_id,
priority=self._concurrency.priority if priority is None else priority,
executor=self._executor,
retry=retry if retry is not None else registration.retry,
timeout=timeout if timeout is not None else registration.timeout,
artifact_dir=artifact_dir,
clock=self._clock,
graceful_stop=self._graceful_stop,
poll=self._poll,
on_finished=self._on_run_finished,
on_change=self._make_observer(resolved_id),
)
if self._store is not None:
# 必须先建行：之后每一次迁移都是按版本更新的 UPDATE。
self._store.create(run.snapshot)
# 基准：存储里现在是这个版本，写穿的 expect_version 以它为准。
self._seed_persist_state(resolved_id, run.snapshot.status_version)
if deferred:
# 等待上游：停在 blocked，**不**进待执行队列，执行器取不到它。
run.block(reason="等待上游任务")
else:
# 先落 queued：它表示"已被受理"，与 pending（尚未受理）不同。
run.enqueue()
self._runs[resolved_id] = _Submission(
run=run,
registration=registration,
args=args,
kwargs=dict(kwargs),
)
if deferred:
start_now = False
else:
start_now = self._can_start_locked(run.snapshot)
if start_now:
self._active[resolved_id] = _ActiveSlot(bucket, run.snapshot.tags)
else:
self._pending.append(run)
if start_now:
run.start()
elif not deferred:
self._ensure_dispatcher()
# 只有真正提交过任务的实例才去轮询指令：否则一个纯查询客户端会把别人的
# 停止请求抢走（见 TaskStore.release_command 的说明）。
self._ensure_worker()
return run
def _bucket_for(self, bucket_key: str | None) -> int | None:
"""把业务键映射成桶号；未给键则返回 ``None``（该任务不参与分桶）。"""
if bucket_key is None:
return None
return self._bucket_policy.bucket_for(bucket_key)
def _release(self, run: TaskRun) -> None:
"""把一个等待上游的任务放进待执行队列（``blocked`` → ``queued``）。
由管道调度器在上游满足后调用。若任务已经落终态（例如刚被连坐取消），什么也
不做——把它塞进队列只会让调度器取到一个不该启动的任务。
"""
with self._lock:
if self._closed or run.snapshot.is_terminal:
return
run.unblock()
self._pending.append(run)
self._ensure_dispatcher()
self._wake.set()
# ------------------------------------------------------------------
# 管道（DAG）
# ------------------------------------------------------------------
def submit_pipeline(
self,
pipeline: Pipeline,
*,
pipeline_run_id: str | None = None,
) -> PipelineRun:
"""提交一条管道，返回可等待的运行句柄。
定义会在**创建任何任务之前**全部校验完（任务名已注册、无环、依赖存在、结算
方式可用）：否则一个写错的 DAG 会在存储里留下一堆半成品任务，而它们看起来
与正常提交的没有区别。
"""
run = start_pipeline(self, pipeline, pipeline_run_id=pipeline_run_id)
with self._lock:
self._pipelines[run.run_id] = run
return run
def pipelines(self) -> tuple[PipelineRun, ...]:
"""本进程提交过的管道运行。"""
with self._lock:
return tuple(self._pipelines.values())
def get_pipeline(self, run_id: str) -> PipelineRun | None:
"""按 id 取管道运行；不存在返回 ``None``。"""
with self._lock:
return self._pipelines.get(run_id)
def _watch_pipeline(self, run: PipelineRun) -> None:
"""让管道运行订阅任务落终态事件，然后立刻评估一次。
立刻评估是必须的：根节点在上面的循环里已经被提交，它们可能在订阅之前就跑完
了，那样监听器永远收不到通知，管道会一直停在 ``running``。
"""
run_id = run.run_id
def listener(outcome: TaskRunOutcome) -> None:
if outcome.snapshot.pipeline_run_id != run_id:
return
run._on_task_finished(outcome.snapshot.task_id)
self.add_listener(listener)
run._evaluate()
# ------------------------------------------------------------------
# 写穿到存储
# ------------------------------------------------------------------
def _make_observer(self, task_id: str) -> Callable[[TaskSnapshot, Transition], None]:
"""构造挂在状态机上的观察者。
挂在状态机的 ``_apply`` 上而不是散落在各业务方法里，是为了保证**任何**一次
迁移都不会漏写——漏一次就意味着存储与内存悄悄分叉，而且不会有任何报错。
"""
def observe(snapshot: TaskSnapshot, transition: Transition) -> None:
self._persist(snapshot, transition)
del task_id
return observe
def _persist_state(self, task_id: str) -> _PersistState:
"""取（或创建）一个任务的写穿状态。"""
with self._lock:
state = self._persist_states.get(task_id)
if state is None:
state = _PersistState()
self._persist_states[task_id] = state
return state
def _seed_persist_state(self, task_id: str, version: int) -> None:
"""记录"存储里当前是哪个版本"，作为后续写穿的 ``expect_version`` 基准。"""
self._persist_state(task_id).last_persisted_version = version
def _persist(self, snapshot: TaskSnapshot, transition: Transition) -> None:
"""把一次状态迁移写穿到存储。
## 为什么 ``expect_version`` 用"上次**确实写进去**的版本"
曾经的写法是 ``transition.version - 1``，它**假设此前每一次迁移都成功落库**。
一旦某次 ``save`` 失败（返回值不可序列化、网络抖动、连接被断……），存储就停在
k-1 而内存已到 k；此后的每一次写入都会因为 ``expect_version`` 与存储对不上而
冲突——**一次瞬时故障让该任务的存储状态永久冻结**，最后被 ``reconcile()`` 判成
``LOST``，把一个真正跑完的任务报成"结果未知"。
现在记的是"我上次确实写进去的版本"：漏写一笔不会污染后续写入，下一次迁移会把
存储一次追上（版本跳号是允许的，``status_version`` 只需要单调）。
## 存储故障不让任务失败
任务本身没问题，坏的只是可观测性。所以失败会进缓冲区、等下次再试，并记进
:attr:`persist_failures`。但绝不静默——持续分叉会出现在 :attr:`desynced_tasks`。
"""
store = self._store
if store is None:
return
state = self._persist_state(snapshot.task_id)
with state.lock:
if state.desynced is not None:
# 已经放弃这个任务的写入（存在外部写入者），不再刷屏。
return
state.pending.append((snapshot, transition))
if len(state.pending) > _PENDING_TRANSITIONS:
excess = len(state.pending) - _PENDING_TRANSITIONS
for _ in range(excess):
state.pending.popleft()
state.dropped += excess
self._flush_pending(snapshot.task_id, state)
def _flush_pending(self, task_id: str, state: _PersistState) -> None:
"""尝试把缓冲的写穿补上。须持有 ``state.lock``。
顺序是：**只写最新快照**（存储行只需要反映当前状态），再按序补写全部迁移
（审计轨迹需要完整）。``append`` 按 ``(task_id, version)`` 幂等，所以中途失败
可以安全重来。
"""
store = self._store
if store is None or not state.pending:
return
latest = state.pending[-1][0]
try:
store.save(latest, expect_version=state.last_persisted_version)
except ConcurrencyConflict as exc:
# 用"上次确实写进去的版本"仍然冲突 → 存储里的版本是**别人**推进的。
# 覆盖它会抹掉对方的状态，所以放手并明确报告，而不是反复重试。
state.desynced = (
f"检测到外部写入者（期望版本 {exc.expected_version}，"
f"实际 {exc.actual_version}），已停止写穿以免覆盖对方"
)
state.pending.clear()
self._remember_failure(f"{task_id}: {state.desynced}")
return
except Exception as exc:
# 瞬时故障：保留缓冲，下一次迁移时再试。任务本身不受影响。
self._remember_failure(
f"{task_id}@{latest.status_version} 写穿失败（已缓存待重试）："
f"{type(exc).__name__}: {exc}"
)
return
state.last_persisted_version = latest.status_version
while state.pending:
_, record = state.pending[0]
try:
store.append(record)
except Exception as exc:
self._remember_failure(
f"{task_id}@{record.version} 迁移轨迹补写失败（已缓存待重试）："
f"{type(exc).__name__}: {exc}"
)
return
state.pending.popleft()
if state.dropped:
self._remember_failure(
f"{task_id}: 写穿积压超过 {_PENDING_TRANSITIONS} 条上限，"
f"丢弃了 {state.dropped} 条迁移记录（轨迹有缺口）"
)
state.dropped = 0
def _flush_all_pending(self) -> None:
"""收尾前把积压的写穿再试一次（尽力而为）。"""
if self._store is None:
return
with self._lock:
items = list(self._persist_states.items())
for task_id, state in items:
with state.lock:
if state.desynced is None:
self._flush_pending(task_id, state)
@property
def desynced_tasks(self) -> Mapping[str, str]:
"""存储状态与内存**永久分叉**、已放弃写入的任务：``task_id -> 原因``。
唯一会让"内存显示成功、存储停在半路"变成永久的情形（通常是检测到外部写入者）。
必须可见——否则只能靠翻 :attr:`persist_failures` 才发现，而那里只留最近 20 条。
"""
with self._lock:
return {
task_id: state.desynced
for task_id, state in self._persist_states.items()
if state.desynced is not None
}
# ------------------------------------------------------------------
# 查询
# ------------------------------------------------------------------
def get(self, task_id: str) -> TaskSnapshot | None:
"""单个任务快照；不存在返回 ``None``。
绑定了存储时**优先读存储**——那才能看到别的进程提交的任务。本地回退是为了
容忍写穿失败：内存里的状态总比没有强。
"""
store = self._store
if store is not None:
snapshot = store.get(task_id)
if snapshot is not None:
return snapshot
with self._lock:
submission = self._runs.get(task_id)
return submission.run.snapshot if submission is not None else None
def runs(self) -> tuple[TaskRun, ...]:
"""全部任务句柄，按提交时间排序。"""
with self._lock:
items = list(self._runs.values())
items.sort(
key=lambda item: item.run.snapshot.created_at or datetime.min.replace(tzinfo=UTC)
)
return tuple(item.run for item in items)
def list_tasks(
self,
*,
status: TaskStatus | Iterable[TaskStatus] | None = None,
name_prefix: str | None = None,
tag: str | None = None,
bucket: int | None = None,
buckets: Iterable[int] | None = None,
pipeline_run_id: str | None = None,
limit: int | None = None,
) -> list[TaskSnapshot]:
"""按条件查询任务快照。
绑定了存储时**走存储查询**——因此能看到别的进程提交的任务，且过滤是下推的
（SQLite 侧有索引），不是取回全部再筛。
``buckets`` 是为将来的分布式预留的：一台机器只关心自己拥有的那组桶。
方法名刻意不叫 ``list``：那会在类作用域里遮蔽内置的 ``list``，使类内所有
``list[...]`` 类型注解（例如 :meth:`wait_all` 的返回值）失效——mypy 会报
"Function is not valid as a type"。
"""
wanted: frozenset[TaskStatus] | None = None
if status is not None:
wanted = frozenset({status}) if isinstance(status, TaskStatus) else frozenset(status)
wanted_buckets: frozenset[int] | None = None
if buckets is not None:
wanted_buckets = frozenset(buckets)
elif bucket is not None:
wanted_buckets = frozenset({bucket})
store = self._store
if store is not None:
return store.query(
TaskQuery(
statuses=wanted,
name_prefix=name_prefix,
tag=tag,
buckets=wanted_buckets,
pipeline_run_id=pipeline_run_id,
limit=limit,
)
)
snapshots = [run.snapshot for run in self.runs()]
if wanted is not None:
snapshots = [s for s in snapshots if s.status in wanted]
if wanted_buckets is not None:
snapshots = [s for s in snapshots if s.bucket in wanted_buckets]
if pipeline_run_id is not None:
snapshots = [s for s in snapshots if s.pipeline_run_id == pipeline_run_id]
if name_prefix is not None:
snapshots = [s for s in snapshots if s.name.startswith(name_prefix)]
if tag is not None:
snapshots = [s for s in snapshots if tag in s.tags]
return snapshots[:limit] if limit is not None else snapshots
def stats(self) -> StatusCounts:
"""各状态计数。绑定了存储时统计的是**全部进程**的任务。"""
store = self._store
if store is not None:
return StatusCounts(by_status=store.counts())
counts: dict[TaskStatus, int] = {}
for run in self.runs():
counts[run.status] = counts.get(run.status, 0) + 1
return StatusCounts(by_status=counts)
def bucket_counts(self) -> dict[int | None, int]:
"""每个桶的任务数（``None`` 键表示未分桶的任务）。
分布式调度要靠它回答"哪些桶是空的、哪些桶积压了"，进而决定把桶分给谁。
绑定了存储时是 O(全表)——查询协议还没有"按桶聚合计"的入口，等真的成为
瓶颈时再加；现在加就是无消费者的死代码。
"""
if self._store is not None:
snapshots = self._store.query()
else:
snapshots = [run.snapshot for run in self.runs()]
counts: dict[int | None, int] = {}
for snapshot in snapshots:
counts[snapshot.bucket] = counts.get(snapshot.bucket, 0) + 1
return counts
@property
def bucket_policy(self) -> BucketPolicy:
"""分桶策略。桶数是集群范围的约定，分布式部署时各处必须一致。"""
return self._bucket_policy
def _validated_owned_buckets(self, buckets: Iterable[int] | None) -> frozenset[int] | None:
"""校验并收下"本节点负责的桶"。
越界的桶号在启动时就报错：如果放过去，症状是"某些任务永远停在 ``queued`` 且
没有任何报错"——那是最难查的一类问题。
"""
if buckets is None:
return None
resolved = frozenset(buckets)
if not resolved:
raise ConfigurationError(
"owned_buckets 不能是空集：那意味着本节点什么都不负责。"
"若想表示'不限定归属'，请传 None（默认）。"
)
limit = self._bucket_policy.num_buckets
# 先逐个检查类型，**再**排序：对混合类型排序会先抛 TypeError，
# 把一条清晰的配置错误变成一句莫名其妙的比较错误。
for bucket in resolved:
if not isinstance(bucket, int) or isinstance(bucket, bool):
raise ConfigurationError(f"桶号必须是整数，收到 {bucket!r}")
if bucket < 0 or bucket >= limit:
raise ConfigurationError(
f"桶号越界：{bucket}；本集群的桶数是 {limit}（合法范围 0..{limit - 1}）"
)
return resolved
@property
def owned_buckets(self) -> frozenset[int] | None:
"""本节点负责的桶；``None`` 表示不限定。
``None`` 与"空集"是两件不同的事：前者是"我不管归属"，后者是"我一个桶都不负责"。
后者通常是配置写漏了，所以构造时会拒绝空集。
"""
return self._owned_buckets
@property
def owns_all_buckets(self) -> bool:
"""是否**不构成划分**：不限定，或显式覆盖了全部桶。
覆盖全部桶时没有任何"别的节点"，因此无桶任务同样会且只会由本节点执行——
此时再强制每个任务写 ``bucket_key`` 就只是纯摩擦。
"""
owned = self._owned_buckets
if owned is None:
return True
return len(owned) == self._bucket_policy.num_buckets
def owns_bucket(self, bucket: int | None) -> bool:
"""某个桶是否归本节点负责（只回答桶的归属，不回答"无桶任务能否提交"）。
未限定时恒为 ``True``；限定时 ``None`` 不算归我。
"""
if self._owned_buckets is None:
return True
return bucket is not None and bucket in self._owned_buckets
def _check_bucket_ownership(
self, bucket: int | None, *, task_id: str, bucket_key: str | None
) -> None:
"""提交时确认任务落在本节点负责的桶里。
判据只有一条：**这个任务将来会不会有人执行**。
为什么是**硬错误**而不是警告：一个落在别人桶里的任务不会有任何节点去跑它，
它会永远停在 ``queued``，而且**没有任何报错**。与其让它静默堆积，不如在提交
那一刻就说清楚。
"""
if self.owns_all_buckets:
return
owned = self._owned_buckets
assert owned is not None # owns_all_buckets 为 False 时必然有限定
if bucket is None:
raise ConfigurationError(
f"本节点只负责桶 {sorted(owned)}，因此提交任务 {task_id!r} 时必须给出 "
"bucket_key——否则它没有归属桶，不会有任何节点认领它。"
)
if bucket not in owned:
raise ConfigurationError(
f"桶 {bucket}（来自 bucket_key={bucket_key!r}）不属于本节点；"
f"本节点负责 {sorted(owned)}。任务若写进别人的桶，不会有任何节点执行它。"
)
@property
def active_count(self) -> int:
"""当前占用并发额度的任务数。"""
with self._lock:
return len(self._active)
def active_in_bucket(self, bucket: int | None) -> int:
"""某个桶当前占用的执行额度。"""
with self._lock:
return sum(1 for slot in self._active.values() if slot.bucket == bucket)
@property
def queued_count(self) -> int:
"""正在排队等待执行槽的任务数。"""
with self._lock:
return len(self._pending)
def _can_start_locked(self, snapshot: TaskSnapshot) -> bool:
"""须在持有 ``_lock`` 时调用。三道闸都要有余量才放行。
顺序由"最便宜且最可能挡住"到"最贵"：全局额度 → 桶额度 → 具名资源。不过
顺序只影响性能，不影响结论——三者是**与**关系，任何一道不过都不启动。
"""
policy = self._concurrency
if policy.max_parallel is not None and len(self._active) >= policy.max_parallel:
return False
if self._max_parallel_per_bucket is not None:
in_bucket = sum(1 for slot in self._active.values() if slot.bucket == snapshot.bucket)
if in_bucket >= self._max_parallel_per_bucket:
return False
return self._resources_available_locked(snapshot.tags)
def _resources_available_locked(self, tags: frozenset[str]) -> bool:
"""具名信号量：任务需要的**每个**受限期资源都还有余量。须持有 ``_lock``。
只统计候选任务自己关心的那几个标签：没写进 ``limits`` 的标签不受限，因此
数它们纯属浪费——标签集可能很大，而受限期通常只有一两个。
"""
limits = self._concurrency.limits
if not limits or not tags:
return True
wanted = [tag for tag in tags if tag in limits]
if not wanted:
return True
in_use: dict[str, int] = dict.fromkeys(wanted, 0)
for slot in self._active.values():
for tag in slot.tags:
if tag in in_use:
in_use[tag] += 1
return all(in_use[tag] < limits[tag] for tag in wanted)
# ------------------------------------------------------------------
# 重启恢复
# ------------------------------------------------------------------
def reconcile(
self,
*,
stale_after: float | None = None,
reason: str | None = None,
dry_run: bool = False,
) -> ReconcileReport:
"""清理上次进程留下的**僵尸任务**。进程启动时应当先跑一次。
## 判据是"还有人在管它吗"，不是"它还在推进吗"
每个监督进程都会给自己负责的任务**续租**（见 :attr:`lease_owner`）。所以：
* 租约未过期 → 有活着的进程在管，**不动它**；
* 租约已过期 / 从来没有租约（且已过宽限期）→ 没有监督者：
- 任务**占着执行资源**（``starting`` / ``running`` / ``pausing`` / ``stopping``）
→ 载体已随进程消失，谁也不会再推进它，判 ``LOST``；
- 任务**不占资源**（``queued`` / ``paused`` / ``retry_wait`` …）
→ 只是没人接手，**不改状态**，报告为 :attr:`ReconcileReport.adoptable`。
这里刻意**不用任务心跳**（``last_heartbeat_at``）当判据。心跳由任务自己的
``checkpoint()`` 驱动，一个长时间不调 checkpoint 的健康任务心跳同样会停滞——
拿它判断进程存活，会把正常任务杀掉。反过来用它做 idle 超时才是对的。
## 它不做什么
**不会重新执行任何任务。** 存储里记的是状态，不是配方：任务的函数与参数从未
落库（序列化任意参数会引入不可反序列化、版本漂移等一堆问题）。所以恢复一个
失去载体的任务在原理上不可能，能做的只是把它如实标记为丢失，让看板不再永远
显示一个假的"运行中"。
真正"重跑"要靠调用方重新 :meth:`submit`，而那需要调用方自己知道参数。
``dry_run=True`` 只出报告、不写存储——这类操作会改状态，先看一眼是值得的。
"""
store = self._store
if store is None:
raise ConfigurationError(
"reconcile 需要共享存储：没有存储就无从判断哪个任务还有监督者。"
)
now = self._clock()
grace = self._lease_ttl * 2 if stale_after is None else stale_after
entries: list[ReconcileEntry] = []
# 只看**本节点负责的桶**。统一存储里可能有多个自治节点的任务；去动别人的任务
# 既越界（那是它的责任）也没有意义（它的租约由它自己续）。
for snapshot in store.query(TaskQuery(statuses=NON_TERMINAL, buckets=self._owned_buckets)):
if self._find(snapshot.task_id) is not None:
continue # 本进程正在管它，当然活着
if self._has_supervisor(snapshot, now=now, grace=grace):
continue
if snapshot.status not in RESOURCE_HOLDING:
entries.append(
ReconcileEntry(
snapshot.task_id,
snapshot.status,
ReconcileOutcome.ADOPTABLE,
"没有执行载体，等待重新提交",
)
)
continue
entries.append(self._reconcile_lost(snapshot, reason=reason, dry_run=dry_run))
return ReconcileReport(entries=tuple(entries), dry_run=dry_run)
@staticmethod
def _has_supervisor(snapshot: TaskSnapshot, *, now: datetime, grace: float) -> bool:
"""是否还有活着的监督者。"""
if snapshot.lease_until is not None:
return snapshot.lease_until > now
# 从未续过租：可能是刚提交、工作线程还没轮到第一次续租，给一个宽限期。
reference = snapshot.started_at or snapshot.created_at
if reference is None:
return False
return (now - reference).total_seconds() < grace
def _reconcile_lost(
self,
snapshot: TaskSnapshot,
*,
reason: str | None,
dry_run: bool,
) -> ReconcileEntry:
"""把一个僵尸任务如实标记为 ``LOST``。"""
store = self._store
assert store is not None # 调用方已确保
message = reason or "监督者已失联（进程重启后未再续租）"
if dry_run:
return ReconcileEntry(
snapshot.task_id, snapshot.status, ReconcileOutcome.MARKED_LOST, message
)
# 仍然经由状态机构造迁移：它是唯一写者，边界校验（终态集合、转移表）都在那里。
# 直接拼一条 Transition 会绕开这些校验，也会让"非法转移"从不该发生的路径进来。
machine = StateMachine(snapshot, clock=self._clock)
try:
record = machine.mark_lost(
reason=message,
failure=Failure(kind=FailureKind.LOST, message=message, retryable=False),
)
store.save(machine.snapshot, expect_version=snapshot.status_version)
store.append(record)
except ConcurrencyConflict:
# 原本的监督者又活过来了（或别人抢先处理了）。这时候**必须放手**：
# 强行覆盖会把对方正在写的状态抹掉。
return ReconcileEntry(
snapshot.task_id,
snapshot.status,
ReconcileOutcome.SKIPPED,
"有人正在写这个任务，已跳过",
)
except (InvalidTransition, StoreError) as exc:
return ReconcileEntry(
snapshot.task_id,
snapshot.status,
ReconcileOutcome.SKIPPED,
f"{type(exc).__name__}: {exc}",
)
return ReconcileEntry(
snapshot.task_id, snapshot.status, ReconcileOutcome.MARKED_LOST, message
)
# ------------------------------------------------------------------
# 控制
# ------------------------------------------------------------------
def stop(self, task_id: str, *, reason: str | None = None) -> bool:
"""请求停止单个任务。返回是否找到了该任务。
本进程持有的任务直接下发；属于**别的进程**的任务则往存储里写一条指令，
由持有者轮询取走。两种情况下返回值都表示"请求已被受理"。
"""
run = self._find(task_id)
if run is not None:
run.stop(reason=reason)
return True
return self.request(task_id, CommandKind.STOP, reason=reason)
def pause(self, task_id: str, *, timeout: float = 5.0) -> bool:
"""请求暂停单个任务；返回是否**确认**已挂起。
对属于别的进程的任务，本方法会下发指令并返回 ``False``——**跨进程无法确认
暂停是否生效**（那需要回执通道）。要观测结果请轮询 :meth:`get`。
返回 ``False`` 在这里的含义是"未确认"，不是"失败"。
"""
run = self._find(task_id)
if run is not None:
return run.pause(timeout=timeout)
if self.request(task_id, CommandKind.PAUSE):
return False
return False
def resume(self, task_id: str) -> bool:
"""恢复单个任务。返回请求是否已被受理。"""
run = self._find(task_id)
if run is not None:
run.resume()
return True
return self.request(task_id, CommandKind.RESUME)
def request(
self,
task_id: str,
kind: CommandKind,
*,
reason: str | None = None,
) -> bool:
"""向**别的进程**持有的任务下发一条控制指令。
这是跨进程控制的唯一通道：一个进程无法给另一个进程的执行载体发信号，只能
把意图写进共享存储，由持有者轮询取走。
任务在本进程内时不要用它——直接调 :meth:`stop` / :meth:`pause` /
:meth:`resume`，那是同步的、能拿到确认的。
"""
store = self._store
if store is None:
return False
# 先确认任务真实存在。否则一个拼错的 task_id 会写出一条**永远不会有人执行**的
# 指令，一直躺在存储里；而调用方还会以为"停止已下发"。
if store.get(task_id) is None:
return False
store.submit_command(
Command(
command_id=uuid.uuid4().hex[:12],
task_id=task_id,
kind=kind,
created_at=self._clock(),
requested_by=self._owner,
reason=reason,
)
)
self._ensure_worker()
self._wake.set()
return True
def stop_all(self, *, reason: str | None = None) -> int:
"""向所有未落终态的任务**发出**停止请求，返回本轮发出的请求数。
是"已发出"而不是"已停止"：请求发出后任务先进入 ``stopping``，落终态是
异步的。要等它们真的结束，请接着调 :meth:`wait_all`。
只作用于**本进程持有**的任务。别的进程的任务要靠共享存储里的指令，而
"全部停止"跨进程需要一份全局任务清单——那是分布式调度该决定的事，不是
一个便利方法顺手能定义的。
"""
count = 0
for run in self.runs():
if not run.snapshot.is_terminal:
run.stop(reason=reason)
count += 1
return count
# ------------------------------------------------------------------
# 等待与收尾
# ------------------------------------------------------------------
def wait(self, task_id: str, timeout: float | None = None) -> TaskRunOutcome:
"""等待单个任务落终态。"""
run = self._find(task_id)
if run is None:
raise ConfigurationError(f"没有这个任务：{task_id}")
return run.wait(timeout=timeout)
def wait_all(self, timeout: float | None = None) -> list[TaskRunOutcome]:
"""等待**所有**任务落终态，返回已落终态的结果。
包含仍在排队的任务——它们会随着并发额度释放而被自动启动。若 ``timeout``
到期仍有任务未结束，只返回已经落终态的那些（不会抛异常，也不会改变任何任务
的状态）。
不能用"逐个 ``run.wait()``"实现：排队中的任务尚未启动，没有可等待的执行。
"""
deadline = None if timeout is None else time.monotonic() + timeout
while True:
pending = [run for run in self.runs() if not run.snapshot.is_terminal]
if not pending:
break
if deadline is not None and time.monotonic() >= deadline:
break
time.sleep(self._poll)
outcomes: list[TaskRunOutcome] = []
for run in self.runs():
if run.started and run.snapshot.is_terminal:
outcomes.append(run.wait())
return outcomes
def close(
self,
*,
reason: str | None = None,
timeout: float = 30.0,
close_store: bool = False,
) -> list[TaskRunOutcome]:
"""停止一切并收尾：排队的任务会被"启动后立刻停止"，不会真的执行任务体。
``close_store`` 默认 ``False``：存储是调用方传进来的，通常由调用方决定它的
生命周期（多个 Controller 可能共用同一个存储）。但**用完一定要关**——一个
没关的 SQLite 连接会泄漏文件句柄，在 Windows 上还会让数据库文件无法删除。
"""
# 先关闸，再停止：否则调度器可能抢在清空队列之前把排队任务跑起来。
with self._lock:
self._closed = True
still_pending = list(self._pending)
self._pending.clear()
self.stop_all(reason=reason)
for run in still_pending:
# 监督循环第一轮就会看到停止请求并落 stopped，任务体不会被执行。
run.start()
self._wake.set()
outcomes = self.wait_all(timeout=timeout)
for thread in (self._dispatcher, self._worker):
if thread is not None:
thread.join(timeout)
# 收尾前再试一次：把因瞬时故障积压的写穿补上。尽力而为——补不上也不会让这里
# 抛异常，原因已经记在 persist_failures / desynced_tasks 里。
self._flush_all_pending()
if close_store and self._store is not None:
self._store.close()
return outcomes
def prune(self, *, keep_terminal: int = 1000) -> int:
"""丢弃最旧的终态记录，防止内存无限增长；返回丢弃的数量。
绑定了存储时**只清内存**：存储里的记录仍然可查。这正是持久化的意义——
内存可以丢，历史不该丢。
"""
with self._lock:
terminal = [item for item in self._runs.values() if item.run.snapshot.is_terminal]
terminal.sort(
key=lambda item: item.run.snapshot.created_at or datetime.min.replace(tzinfo=UTC)
)
excess = len(terminal) - max(0, keep_terminal)
if excess <= 0:
return 0
for item in terminal[:excess]:
task_id = item.run.snapshot.task_id
del self._runs[task_id]
# 写穿状态也跟着丢：留着它只会让字典无限增长。
self._persist_states.pop(task_id, None)
return excess
# ------------------------------------------------------------------
# 内部
# ------------------------------------------------------------------
def _find(self, task_id: str) -> TaskRun | None:
with self._lock:
submission = self._runs.get(task_id)
return submission.run if submission is not None else None
def _on_run_finished(self, outcome: TaskRunOutcome) -> None:
"""任务落终态：归还并发额度并唤醒调度器。"""
with self._lock:
self._active.pop(outcome.snapshot.task_id, None)
listeners = tuple(self._listeners)
self._wake.set()
# 监听器在额度已归还、调度器已唤醒**之后**才被调用，因此它们可以安全地提交
# 新任务。它们的异常绝不能影响任务本身的结论，但也不能静默——见 listener_failures。
for listener in listeners:
try:
listener(outcome)
except Exception as exc:
with self._lock:
self._listener_failures.append(
f"{outcome.snapshot.task_id}: {type(exc).__name__}: {exc}"
)
del self._listener_failures[:-20]
def add_listener(self, listener: Callable[[TaskRunOutcome], None]) -> None:
"""注册一个"任务落终态"的监听器，返回 ``None``。
管道调度器靠它推进——这是它唯一需要知道的时机，不必轮询。监听器在并发额度
已归还之后被调用，因此里面可以安全地提交新任务。
"""
with self._lock:
self._listeners.append(listener)
def _ensure_dispatcher(self) -> None:
with self._lock:
if self._dispatcher is not None:
return
self._dispatcher = threading.Thread(
target=self._dispatch_loop, name="controller-dispatch", daemon=True
)
self._dispatcher.start()
def _take_runnable_locked(self) -> TaskRun | None:
"""取下一个可启动的排队任务；没有则返回 ``None``。须持有 ``_lock``。
必须**扫描**而不是直接取队首：某个桶或某个具名资源满的时候，队首可能正好被
它挡住，而后面不受它约束的任务本可以先跑。直接取队首会让一处积压阻塞所有
任务——那正是分桶与信号量要解决的问题，不能在调度器里又引入一次。
在**都能启动**的任务里取优先级最高的；同优先级取先提交的——优先级机制不该
打乱同优先级的时间序。
扫描途中顺手清掉**已落终态**的排队任务（被取消或跳过）：它们绝不能被启动，
否则监督循环会撞上非法状态转移。
"""
best_index = -1
best_priority = 0
index = 0
while index < len(self._pending):
run = self._pending[index]
if run.snapshot.is_terminal:
del self._pending[index]
continue
if self._can_start_locked(run.snapshot) and (
best_index < 0 or run.priority > best_priority
):
best_index = index
best_priority = run.priority
index += 1
if best_index < 0:
return None
run = self._pending[best_index]
del self._pending[best_index]
return run
def _dispatch_loop(self) -> None:
while True:
self._wake.wait(0.2)
self._wake.clear()
while True:
with self._lock:
# 关闭后不再启动任何新工作：排队的任务由 close() 自己收尾，
# 否则它可能抢在 close() 清空队列之前把任务跑起来。
if self._closed:
return
run = self._take_runnable_locked()
if run is None:
break
self._active[run.snapshot.task_id] = _ActiveSlot(
run.snapshot.bucket, run.snapshot.tags
)
run.start()
# ------------------------------------------------------------------
# 跨进程：指令轮询与心跳回写
# ------------------------------------------------------------------
def _ensure_worker(self) -> None:
"""启动跨进程工作线程。只在真的有存储、且真的提交过任务时启动。
一个纯查询客户端若也去取指令，就会把别人的停止请求抢走；即便它把指令交还，
持续轮询也会让真正的持有者反复错过。所以入口设成"提交过任务才启动"。
"""
if self._store is None:
return
with self._lock:
if self._worker is not None:
return
self._worker = threading.Thread(
target=self._worker_loop, name="controller-worker", daemon=True
)
self._worker.start()
def _worker_loop(self) -> None:
while True:
self._wake.wait(self._command_poll_interval)
self._wake.clear()
if self._closed:
return
self._drain_commands()
self._flush_activity()
def _drain_commands(self) -> None:
store = self._store
if store is None:
return
try:
commands = store.claim_commands(now=self._clock(), lease_seconds=self._command_lease)
except Exception as exc:
self._remember_failure(f"claim_commands: {type(exc).__name__}: {exc}")
return
for item in commands:
run = self._find(item.task_id)
if run is None:
# 不属于本进程：**交还**而不是标记完成，否则真正的持有者永远收不到
# 这次停止请求（一个静默的丢指令）。
store.release_command(item.command_id, now=self._clock())
continue
outcome = self._apply_command(item, run)
store.complete_command(item.command_id, now=self._clock(), outcome=outcome)
def _apply_command(self, item: Command, run: TaskRun) -> str:
"""执行一条指令。返回写回存储的结果摘要，便于事后追溯。"""
if run.snapshot.is_terminal:
return "already-terminal"
if item.kind is CommandKind.STOP:
run.stop(reason=item.reason)
return "stopped"
if item.kind is CommandKind.PAUSE:
confirmed = run.pause(timeout=self._graceful_stop)
return "paused" if confirmed else "pause-not-confirmed"
run.resume()
return "resumed"
def _flush_activity(self) -> None:
"""把本进程任务的进度与心跳推给存储。
走 ``touch``（高频旁路）而不是状态写入：进度可能每秒几十次，若参与乐观并发
会与状态写入持续冲突。
"""
store = self._store
if store is None:
return
lease_until = self._clock() + timedelta(seconds=self._lease_ttl)
for run in self.runs():
# 先让快照跟上载体的真实心跳，否则下一次全量 save 会把 touch 写进去的
# 心跳覆盖回 None（同一行的两条写入路径互相打架）。
snapshot = run.refresh_activity()
if snapshot.is_terminal:
continue
try:
store.touch(
snapshot.task_id,
progress=snapshot.progress,
heartbeat_at=snapshot.last_heartbeat_at,
# 续租回答的是"还有人在管它吗"，与心跳是两件事——用任务心跳判断
# 监督者存活，会把一个长时间不调 checkpoint 的健康任务误判失联。
lease_owner=self._owner,
lease_until=lease_until,
)
except Exception as exc:
self._remember_failure(f"touch {snapshot.task_id}: {type(exc).__name__}: {exc}")
def _remember_failure(self, summary: str) -> None:
with self._lock:
self._persist_failures.append(summary)
del self._persist_failures[:-20]
# ----------------------------------------------------------------------
# 渲染：CLI 与看板共用
# ----------------------------------------------------------------------
def _display_width(text: str) -> int:
"""按终端显示宽度计算：中日韩全角字符占两列。
不这么做，中文任务名会把整张表挤歪——运维看板里这种事非常刺眼。
"""
return sum(2 if unicodedata.east_asian_width(ch) in {"W", "F"} else 1 for ch in text)
def _pad(text: str, width: int) -> str:
return text + " " * max(0, width - _display_width(text))
def _format_elapsed(snapshot: TaskSnapshot, now: datetime) -> str:
"""已执行时长。
尚未开始执行时返回 ``-`` 而不是"距创建多久"——排队中的任务显示一个时长会被
误读成"已经跑了这么久"。
"""
if snapshot.started_at is None:
return "-"
end = snapshot.ended_at or now
return f"{(end - snapshot.started_at).total_seconds():.1f}s"
def _format_progress(progress: Progress | None) -> str:
if progress is None:
return "-"
if progress.total is None:
return str(progress.done)
return f"{progress.done}/{progress.total}"
def _format_bucket(bucket: int | None) -> str:
return "-" if bucket is None else str(bucket)
def render_table(snapshots: Sequence[TaskSnapshot], *, now: datetime | None = None) -> str:
"""把任务快照渲染成等宽文本表。
CLI 直接打印它；将来的看板可以拿同样的数据渲染 HTML，两边不会出现口径差异。
"""
moment = now if now is not None else utcnow()
rows: list[tuple[str, ...]] = [
("TASK_ID", "NAME", "BUCKET", "STATUS", "ATTEMPT", "ELAPSED", "PROGRESS", "FAILURE")
]
for snapshot in snapshots:
failure: Failure | None = snapshot.failure
rows.append(
(
snapshot.task_id,
snapshot.name,
_format_bucket(snapshot.bucket),
snapshot.status.value,
f"{snapshot.attempt}/{snapshot.max_attempts}",
_format_elapsed(snapshot, moment),
_format_progress(snapshot.progress),
failure.kind.value if failure is not None else "-",
)
)
widths = [max(_display_width(row[index]) for row in rows) for index in range(len(rows[0]))]
lines = []
for position, row in enumerate(rows):
line = _COLUMN_SEPARATOR.join(
_pad(cell, width) for cell, width in zip(row, widths, strict=True)
).rstrip()
lines.append(line)
if position == 0:
lines.append(_COLUMN_SEPARATOR.join("-" * width for width in widths))
return "\n".join(lines)
./src\chongming_task\decorators.py
"""``@task`` 装饰器：把任务元数据写在函数定义旁边。
```python
from chongming_task import RetryPolicy, task
@task(tags={"io"}, retry=RetryPolicy(max_attempts=3), timeout=Timeout(per_attempt=60))
def extract(ctx, source: str) -> list[Row]:
\"\"\"从 source 读取数据。\"\"\"
...
然后一次性登记整个模块::
controller.register_module(myapp.tasks)
为什么挂在函数上、而不是写进一个全局注册表
装饰器只把规格附在函数对象上（属性 __chongming_task__），不产生任何全局状态。
好处不只是"好测试"：
- 谁注册、注册到哪个 :class:
~chongming_task.controller.TaskRegistry由调用方决定， 同一个模块可以登记进多个互不干扰的控制器； - 导入一个模块不会有副作用。全局注册表意味着 import 即注册，而 import 可能发生 在任何地方（包括测试收集阶段），那种隐式行为极难排查。 """
from future import annotations
from collections.abc import Callable, Iterable from dataclasses import dataclass, field from typing import Any, TypeVar
from .policies import RetryPolicy, Timeout
all = ["TASK_SPEC_ATTR", "TaskSpec", "get_task_spec", "task"]
#: 规格挂在函数上的属性名。 TASK_SPEC_ATTR = "chongming_task"
_F = TypeVar("_F", bound=Callable[..., Any])
@dataclass(frozen=True, slots=True)
class TaskSpec:
"""@task 声明的任务元数据。"""
name: str | None = None
"""注册名。省略时用函数名。"""
description: str = ""
"""描述。省略时取 docstring 的第一行。"""
tags: frozenset[str] = field(default_factory=frozenset)
retry: RetryPolicy | None = None
timeout: Timeout | None = None
def task( _func: _F | None = None, *, name: str | None = None, description: str | None = None, tags: Iterable[str] = (), retry: RetryPolicy | None = None, timeout: Timeout | None = None, ) -> Any: """声明一个任务。
既可以带参数用（``@task(tags={"io"})``），也可以不带参数用（``@task``）——后者的
写法太自然，不支持会让人反复踩空。
"""
def decorate(func: _F) -> _F:
resolved_description = description
if resolved_description is None:
first_line = (func.__doc__ or "").strip().splitlines()
resolved_description = first_line[0].strip() if first_line else ""
spec = TaskSpec(
name=name if name is not None else getattr(func, "__name__", None),
description=resolved_description,
tags=frozenset(tags),
retry=retry,
timeout=timeout,
)
setattr(func, TASK_SPEC_ATTR, spec)
return func
if _func is not None:
# 裸 @task 的用法：参数直接就是被装饰的函数。
return decorate(_func)
return decorate
def get_task_spec(target: Callable[..., Any]) -> TaskSpec | None:
"""取函数上的任务规格；没被 @task 装饰过则返回 None。"""
spec = getattr(target, TASK_SPEC_ATTR, None)
return spec if isinstance(spec, TaskSpec) else None
`./src\chongming_task\enum.py`
```python
"""兼容层（已弃用）。
早期版本在此定义 ``State``。状态模型已迁移至 :mod:`chongming_task.states`，并且必须
使用语义明确的 ``TaskStatus``——``State`` / ``FINISHED`` 无法区分"成功结束"与
"其它终止"，而这一区分是精细化控制的基础。
本模块仅为避免旧引用断裂而保留。新代码请直接::
from chongming_task.states import TaskStatus
"""
from __future__ import annotations
import warnings
from .states import TaskStatus
__all__ = ["TaskStatus"]
def __getattr__(name: str) -> object:
"""在访问 ``State`` 时给出弃用提示（PEP 562）。
``State`` 刻意不在此处静态定义：一旦写成模块级赋值，就会绕过本函数，
弃用提示将永远不会触发。
"""
if name == "State":
warnings.warn(
"chongming_task.enum.State 已弃用，请改用 chongming_task.states.TaskStatus",
DeprecationWarning,
stacklevel=2,
)
return TaskStatus
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
./src\chongming_task\errors.py
"""错误体系。
约定
----
* 用户代码与执行器抛出的异常必须被包装为 :class:`TaskExecutionError` 并保留
``__cause__``，绝不吞掉 traceback。
* :class:`TaskCancelled` 是**控制流信号**而非错误，刻意继承 ``BaseException``。
"""
from __future__ import annotations
from collections.abc import Iterable, Sequence
from .models import TimeoutScope
from .states import TaskStatus, allowed_from
__all__ = [
"AttemptsExhausted",
"CircularDependency",
"ConcurrencyConflict",
"ConfigurationError",
"DependencyFailed",
"InvalidTransition",
"StoreError",
"TaskCancelled",
"TaskError",
"TaskExecutionError",
"TaskTimeout",
]
class TaskError(Exception):
"""本项目所有错误的基类。
注意它**不**包含 :class:`TaskCancelled`，后者是控制流信号。
"""
class InvalidTransition(TaskError):
"""请求了转移表中不允许的状态转移。
刻意抛异常而不是返回 ``False``：静默失败会让"停止指令没生效"这类问题极难排查。
需要预检的调用方应使用 ``StateMachine.can_apply()``。
"""
def __init__(
self,
current: TaskStatus,
requested: TaskStatus,
*,
allowed: Iterable[TaskStatus] | None = None,
action: str | None = None,
) -> None:
self.current = current
self.requested = requested
self.allowed = frozenset(allowed) if allowed is not None else allowed_from(current)
self.action = action
targets = "、".join(sorted(s.value for s in self.allowed)) or "无（终态）"
parts = [f"非法状态转移：{current.value} -> {requested.value}"]
if action is not None:
parts.append(f"动作 {action!r}")
parts.append(f"当前允许：{targets}")
super().__init__("；".join(parts))
class DependencyFailed(TaskError):
"""上游依赖未成功，本任务无法执行。"""
def __init__(self, upstream_ids: Iterable[str], *, reason: str | None = None) -> None:
self.upstream_ids = tuple(upstream_ids)
self.reason = reason
detail = f"上游任务未成功：{', '.join(self.upstream_ids) or '(未提供)'}"
super().__init__(f"{detail}（{reason}）" if reason else detail)
class CircularDependency(TaskError):
"""管道依赖图中存在环。"""
def __init__(self, cycle: Sequence[str]) -> None:
self.cycle = tuple(cycle)
path = " -> ".join((*self.cycle, self.cycle[0])) if self.cycle else "(空环)"
super().__init__(f"管道中存在循环依赖：{path}")
class TaskCancelled(BaseException):
"""协作式取消信号，由 ``TaskContext.checkpoint()`` 抛出。
刻意继承 ``BaseException`` 而非 :class:`TaskError`：若继承 ``Exception``，
用户代码里常见的 ``try/except Exception`` 会静默吞掉取消信号，任务将无法被
停止——这是协作式取消最典型的失效方式。``asyncio.CancelledError`` 采用同样的
处理方式。
"""
def __init__(self, reason: str | None = None, *, forced: bool = False) -> None:
self.reason = reason
self.forced = forced
"""是否由强制手段（强杀/放弃等待）导致。"""
super().__init__(reason or "任务已被取消")
class TaskTimeout(TaskError):
"""超出某个作用域的时限。"""
def __init__(self, scope: TimeoutScope, limit: float, *, task_id: str | None = None) -> None:
self.scope = scope
self.limit = limit
self.task_id = task_id
where = f"任务 {task_id} " if task_id else ""
super().__init__(f"{where}超时（{scope.value}）：已超过 {limit} 秒")
class TaskExecutionError(TaskError):
"""用户代码/执行器异常的统一包装。原异常保留在 ``__cause__``。"""
def __init__(
self,
message: str,
*,
task_id: str | None = None,
attempt: int | None = None,
) -> None:
self.task_id = task_id
self.attempt = attempt
super().__init__(message)
class AttemptsExhausted(TaskError):
"""重试预算已耗尽，不允许再开始新的尝试。
与 :class:`InvalidTransition` 区分：这不是"状态表不允许"，而是"状态表允许但
预算不够"。若混为一谈，运维看到报错会去查转移表，方向完全错了。
"""
def __init__(self, task_id: str, attempt: int, max_attempts: int) -> None:
self.task_id = task_id
self.attempt = attempt
self.max_attempts = max_attempts
super().__init__(
f"任务 {task_id} 重试预算已耗尽：已完成 {attempt} 次尝试，上限 {max_attempts}"
)
class ConcurrencyConflict(TaskError):
"""乐观并发校验失败：快照在本次写入前已被他处修改。
每个任务的写入者应当唯一（调度器）。出现本错误通常意味着存在第二个写入者，
是设计缺陷的信号，因此**不做自动重试掩盖**。
"""
def __init__(self, task_id: str, expected_version: int, actual_version: int) -> None:
self.task_id = task_id
self.expected_version = expected_version
self.actual_version = actual_version
super().__init__(
f"任务 {task_id} 并发写冲突：期望版本 {expected_version}，实际 {actual_version}"
)
class StoreError(TaskError):
"""存储层故障。"""
class ConfigurationError(TaskError):
"""配置错误，如循环依赖、无效策略组合、执行器名未知。"""
./src\chongming_task\models.py
"""任务快照与正交维度模型。
:class:`TaskSnapshot` 既是状态机的输入/输出，也是持久化的单元。它把生命周期
(``status``) 与其余**正交维度**分开承载：把"重试中""已超时""进度 50%"塞进同一个
枚举是同类工具最常见的建模错误，也是"精细化控制"无法落地的主因。
"""
from __future__ import annotations
import json
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import datetime
from enum import StrEnum
from typing import Any, Final
from .states import REVIVABLE, TERMINAL, TaskStatus
__all__ = [
"AttemptRecord",
"Failure",
"FailureKind",
"Progress",
"RequestedAction",
"TaskSnapshot",
"TimeoutScope",
"describe_serialization_problem",
]
class RequestedAction(StrEnum):
"""已受理但尚未完成的控制请求。
与 ``status`` 分开承载：``stop()`` / ``pause()`` 是**请求**而非瞬时赋值，
运行中的任务需要时间响应，UI 依赖本字段渲染"停止中 / 暂停中"。
"""
STOP = "stop"
PAUSE = "pause"
class TimeoutScope(StrEnum):
"""超时作用域。超时是**失败原因**，不是生命周期状态。"""
START = "start"
"""从提交到真正开始执行的等待时长。"""
ATTEMPT = "attempt"
"""单次尝试的总时长。"""
IDLE = "idle"
"""基于心跳的静默超时，用于发现死锁/卡死——这是唯一能发现"没有报错但不动了"的机制。"""
class FailureKind(StrEnum):
"""失败或未正常完成的归类。"""
EXCEPTION = "exception"
NONZERO_EXIT = "nonzero_exit"
TIMEOUT_START = "timeout_start"
TIMEOUT_ATTEMPT = "timeout_attempt"
TIMEOUT_IDLE = "timeout_idle"
CANCELLED = "cancelled"
LOST = "lost"
DEPENDENCY_FAILED = "dependency_failed"
CONDITION_FALSE = "condition_false"
RESULT_INVALID = "result_invalid"
"""任务违反了框架契约，例如返回值不可持久化。
与 ``EXCEPTION`` 分开是有意的：一个是"任务自己出错了"，一个是"任务的产出框架
接不住"。报表与告警需要能分开——它们的修法完全不同（前者查逻辑，后者改接口）。
"""
@classmethod
def for_timeout(cls, scope: TimeoutScope) -> FailureKind:
"""由超时作用域得到对应归类，避免在两处重复书写映射。"""
return cls(f"timeout_{scope.value}")
#: 一次返回值编码尝试中使用的 JSON 选项。
#:
#: 必须与存储层编码快照时**完全一致**，否则校验通过、写库仍可能失败。
_JSON_OPTIONS: Final[dict[str, Any]] = {"ensure_ascii": False, "sort_keys": True}
def describe_serialization_problem(value: Any) -> str | None:
"""检查一个值能否被持久化：可以返回 ``None``，否则返回原因。
为什么需要**在写库之前**检查：持久化契约要求返回值 JSON 可序列化（存储层要把快照
编码成 JSON）。若不检查，失败会发生在**写穿路径**上——而那条路径的失败只被记录下来、
不会让任务失败。结果是进程内显示 ``succeeded``、存储里停在 ``running``，随后
``reconcile()`` 再把一个真正成功过的任务判成 ``LOST``。
调用方应当把"不可序列化"当成**任务失败**：违反契约是任务的问题，不是存储的问题。
"""
try:
# 与 codec 一样把值包在 outcome 里编码，这样键类型、排序等问题也能被发现。
json.dumps({"value": value}, **_JSON_OPTIONS)
except (TypeError, ValueError) as exc:
return f"{type(value).__name__} 不可 JSON 序列化：{exc}"
return None
@dataclass(frozen=True, slots=True)
class Progress:
"""进度。与 ``status`` 正交：``RUNNING`` 的任务可能进度为 0，也可能已达 99%。"""
done: int = 0
total: int | None = None
unit: str | None = None
message: str | None = None
updated_at: datetime | None = None
def __post_init__(self) -> None:
if self.done < 0:
raise ValueError(f"progress.done 不能为负：{self.done}")
if self.total is not None and self.total < 0:
raise ValueError(f"progress.total 不能为负：{self.total}")
if self.total is not None and self.done > self.total:
raise ValueError(f"progress.done({self.done}) 不能超过 total({self.total})")
@property
def fraction(self) -> float | None:
"""完成比例。
``total`` 未知或为 0 时返回 ``None`` 而非 0.0——0.0 会被渲染成"尚未开始"，
而实际情况是"无法估计"，两者必须区分。
"""
if self.total is None or self.total == 0:
return None
return self.done / self.total
@dataclass(frozen=True, slots=True)
class Failure:
"""一次失败的结构化描述。"""
kind: FailureKind
message: str
traceback: str | None = None
"""完整 traceback。用户代码的异常必须保留 ``__cause__`` 并在此完整记录。"""
exit_code: int | None = None
"""执行载体的退出码。**只有进程执行器会写它**（子进程非零退出、或被信号终止）。
同进程执行器为 ``None``：任务出错走的是异常，没有"退出码"这回事。它存在的价值
是把"子进程为什么没了"从一句猜不出的话变成一个有据可查的事实——配合
:attr:`TaskSnapshot.pid`，运维可以 `ps` 对上号。
"""
retryable: bool = True
"""是否值得重试。由策略与错误类型共同决定。"""
@dataclass(frozen=True, slots=True)
class AttemptRecord:
"""单次尝试的不可变历史记录。
注意"只追加"：重试不会修改既有记录。``TaskSnapshot.attempt`` 始终等于最新
尝试的序号。
.. note::
**当前运行时还不产出它**（没有任何代码路径构造 ``AttemptRecord``）：一次尝试
的完整明细今天由迁移历史承载——``starting`` / ``retry_wait`` 这些迁移本身就带
版本号、时间戳、触发源与原因，看板与 ``history`` 都读它。这个类型留给"把逐次
尝试单独持久化"的将来（例如按尝试维度的耗时统计），**不要**在没有写入方之前
把它当成可依赖的接口。
"""
task_id: str
attempt: int
"""从 1 开始的尝试序号。"""
started_at: datetime
ended_at: datetime | None = None
status: TaskStatus = TaskStatus.RUNNING
failure: Failure | None = None
executor: str | None = None
def __post_init__(self) -> None:
if self.attempt < 1:
raise ValueError(f"attempt 从 1 开始，收到 {self.attempt}")
@dataclass(frozen=True, slots=True)
class TaskSnapshot:
"""某一时刻任务的完整状态。
不可变。任何状态变更都由状态机产生**新的**快照并递增 ``status_version``，
因此快照可以安全地跨线程/协程传递，也可以直接作为持久化记录。
"""
# ---- 身份 ----
task_id: str
name: str
pipeline_run_id: str | None = None
tags: frozenset[str] = field(default_factory=frozenset)
# ---- 分桶 ----
bucket: int | None = None
"""所属桶号。由 :class:`~chongming_task.policies.BucketPolicy` 从 ``bucket_key`` 算出。
它既是本地并发隔离的单位，也是将来分布式分配的单位（一台机器拥有一组桶）。
任务一旦创建就把桶**记下来**，而不是每次按需重算：万一哈希策略或桶数发生变化，
重算会把一个正在跑的任务挪到别的桶，局部性与隔离假设随之悄悄失效。
"""
bucket_key: str | None = None
"""算出桶号所依据的业务键（如租户 id、数据源、日期分区）。
留下它是为了**迁移**：将来若必须改动桶数，只有拿得到原始键才能重新分桶。
"""
# ---- 生命周期 ----
status: TaskStatus = TaskStatus.PENDING
status_version: int = 0
"""乐观并发版本号。每次成功转移 +1；存储层用它检测并发写入冲突。"""
requested_action: RequestedAction | None = None
"""已受理但尚未完成的请求，用于区分"正在停止"与"已停止"。"""
# ---- 尝试 ----
attempt: int = 0
"""已开始的尝试次数；0 表示尚未执行过。"""
max_attempts: int = 1
# ---- 时间 ----
created_at: datetime | None = None
queued_at: datetime | None = None
started_at: datetime | None = None
ended_at: datetime | None = None
last_heartbeat_at: datetime | None = None
"""最近一次心跳。超过 idle 超时即判定为卡死或失联。"""
# ---- 结果 ----
progress: Progress | None = None
failure: Failure | None = None
outcome: Mapping[str, Any] | None = None
# ---- 执行载体 ----
executor: str | None = None
"""执行器标识：``inline``（同步、测试用）、``thread``（默认）或 ``process``。"""
pid: int | None = None
"""执行载体的操作系统进程号。**只有进程执行器会写它**：线程与内联执行器下任务
就跑在监督进程里，再存一个 pid 只是重复（``lease_owner`` 已经带主机与进程信息）。
有了它，事后排障能直接对上号：任务卡住时可以先 `ps` 看这个进程还在不在、在干什么。"""
# ---- 监督者租约 ----
lease_owner: str | None = None
"""当前**监督**这个任务的进程标识。
与 ``last_heartbeat_at`` 是两件不同的事，不能合并：
* ``last_heartbeat_at`` 由任务自己的 ``checkpoint()`` 驱动，回答"**任务在推进吗**"。
它停滞意味着卡死——这是发现死锁的唯一信号。
* ``lease_owner`` / ``lease_until`` 由监督进程定期续租，回答"**还有人在管它吗**"。
若拿任务心跳去判断"进程是否还活着"，一个长时间不调 ``checkpoint()`` 的健康任务
会被误判为失联而被判 ``LOST``——那是把一个正常任务杀掉。反过来拿租约去判断卡死
也不行。所以两个信号必须分开。
"""
lease_until: datetime | None = None
"""租约到期时刻。过期即"没有监督者"，重启恢复据此判定僵尸任务。"""
# ---- 人工干预 ----
overridden: bool = False
"""是否经 ``force_transition()`` 人工改写过。此类任务在报表中应被单独标出。"""
def __post_init__(self) -> None:
if not self.task_id:
raise ValueError("task_id 不能为空")
if not self.name:
raise ValueError("name 不能为空")
if self.status_version < 0:
raise ValueError(f"status_version 不能为负：{self.status_version}")
if self.attempt < 0:
raise ValueError(f"attempt 不能为负：{self.attempt}")
if self.max_attempts < 1:
raise ValueError(f"max_attempts 至少为 1，收到 {self.max_attempts}")
if self.attempt > self.max_attempts:
raise ValueError(f"attempt({self.attempt}) 不能超过 max_attempts({self.max_attempts})")
if self.bucket is not None and self.bucket < 0:
raise ValueError(f"bucket 不能为负：{self.bucket}")
if self.bucket is not None and self.bucket_key is None:
raise ValueError("给出了 bucket 却没有 bucket_key，将无法在改桶数时重新分桶")
if self.lease_until is not None and self.lease_owner is None:
# 没有主人的租约无法追溯"是谁在管"，也无法在接管时判断该不该抢。
raise ValueError("给出了 lease_until 却没有 lease_owner")
@property
def is_terminal(self) -> bool:
"""是否已进入终态。"""
return self.status in TERMINAL
@property
def is_revivable(self) -> bool:
"""是否还能经 ``retry()`` 重新执行。"""
return self.status in REVIVABLE
def validate(self) -> tuple[str, ...]:
"""返回跨字段一致性问题清单，空元组表示一致。
刻意**不**在 ``__post_init__`` 中强制这些规则：它们涉及时间线完整性，
在逐步构造快照时（例如先建终态再补 ``ended_at``）会造成无谓摩擦。
持久化层与恢复流程应在写入/加载时调用本方法。
"""
problems: list[str] = []
if self.is_terminal and self.ended_at is None:
problems.append(f"终态 {self.status.value} 必须带有 ended_at")
if not self.is_terminal and self.ended_at is not None:
problems.append(f"非终态 {self.status.value} 不应带有 ended_at")
if self.status is TaskStatus.RUNNING and self.started_at is None:
problems.append("RUNNING 状态必须带有 started_at")
if self.status is TaskStatus.SUCCEEDED and self.failure is not None:
problems.append("SUCCEEDED 不应带有 failure")
if self.status is TaskStatus.RETRY_WAIT and self.failure is None:
problems.append("RETRY_WAIT 必须带有上次失败原因，否则无法解释为何在等待重试")
if self.overridden and self.requested_action is not None:
problems.append("已被人工改写的任务不应残留 requested_action")
if self.requested_action is not None and self.is_terminal:
problems.append(f"终态 {self.status.value} 不应残留 requested_action")
return tuple(problems)
./src\chongming_task\pipeline.py
"""DAG 管道：把多个任务按依赖关系串起来，由一个调度器推进。
状态模型早就为管道准备好了（``PENDING → BLOCKED → QUEUED``、``SKIPPED``、
``CANCELLED``），这里补的是**谁来驱动**：
* :class:`Pipeline` —— 静态定义：节点、依赖、汇合语义、参数；
* :class:`PipelineRun` —— 一次运行：创建节点任务，监听每个节点落终态，放行满足条件的
下游，跳过注定跑不了的下游，按结算策略收场。
## 两个不同的轴，别混
| 概念 | 回答的问题 |
|---|---|
| :class:`~chongming_task.policies.JoinMode`（逐节点） | **这一个节点何时该跑** |
| :class:`~chongming_task.policies.AggregationPolicy`（整条管道） | **管道何时结束、是否连坐** |
## 一个刻意的限制
管道的**结构在内存里**，不落库：任务的函数与参数从未持久化（理由同 :meth:`Controller.reconcile`
的说明）。因此进程重启后，一条跑到一半的管道无法续跑——它的节点会被 ``reconcile``
如实标记，而不是被假装恢复。
"""
from __future__ import annotations
import threading
import uuid
from collections.abc import Iterable, Mapping
from dataclasses import dataclass, field
from enum import StrEnum
from typing import TYPE_CHECKING, Any
from .errors import ConfigurationError
from .policies import AggregationMode, AggregationPolicy, JoinMode, RetryPolicy, Timeout
from .states import TERMINAL, TaskStatus
if TYPE_CHECKING: # pragma: no cover - 仅为类型检查
from .controller import Controller
from .run import TaskRun
__all__ = [
"Node",
"Pipeline",
"PipelineOutcome",
"PipelineRun",
"PipelineStatus",
"render_pipeline",
]
@dataclass(frozen=True, slots=True)
class Node:
"""管道里的一个节点：一个任务 + 它的上游 + 汇合语义。"""
node_id: str
"""节点在本管道内的唯一标识，也是 ``task_id`` 的后半段。"""
task: str
"""在 :class:`~chongming_task.controller.TaskRegistry` 里注册过的任务名。"""
args: tuple[Any, ...] = ()
kwargs: Mapping[str, Any] = field(default_factory=dict)
depends_on: tuple[str, ...] = ()
join: JoinMode = JoinMode.ALL_SUCCEEDED
label: str | None = None
bucket_key: str | None = None
tags: frozenset[str] = frozenset()
retry: RetryPolicy | None = None
timeout: Timeout | None = None
class Pipeline:
"""一条管道的**静态定义**。可重复提交，每次运行互不影响。"""
def __init__(
self,
name: str,
*,
aggregation: AggregationPolicy | None = None,
) -> None:
if not name:
raise ConfigurationError("管道必须有名字")
self.name = name
self.aggregation = aggregation if aggregation is not None else AggregationPolicy.fail_fast()
self._nodes: list[Node] = []
self._by_id: dict[str, Node] = {}
self._counts: dict[str, int] = {}
# ------------------------------------------------------------------
# 定义
# ------------------------------------------------------------------
def node(
self,
task: str,
*args: Any,
depends_on: Iterable[Node | str] = (),
join: JoinMode = JoinMode.ALL_SUCCEEDED,
node_id: str | None = None,
label: str | None = None,
bucket_key: str | None = None,
tags: Iterable[str] = (),
retry: RetryPolicy | None = None,
timeout: Timeout | None = None,
**kwargs: Any,
) -> Node:
"""追加一个节点并返回它，供下游用 ``depends_on=[它]`` 引用。
第一个参数是**已注册的任务名**，其后是传给任务的位置参数——节点 id 由
``node_id`` 关键字指定，因此不会被误读成第二个位置参数（写成 ``name=`` 时
很容易误以为 ``p.node("mark", "after")`` 里的 ``"after"`` 是节点名，其实它是
传给 ``mark`` 的参数）。
``node_id`` 省略时取任务名；同一任务名出现多次时自动加 ``#2`` / ``#3``
后缀——既不强迫调用方给每个节点起名，也不会撞车。
"""
if not task:
raise ConfigurationError("节点必须给出已注册的任务名")
resolved_id = node_id if node_id is not None else self._next_id(task)
if resolved_id in self._by_id:
raise ConfigurationError(f"节点 id 重复：{resolved_id}")
node = Node(
node_id=resolved_id,
task=task,
args=tuple(args),
kwargs=dict(kwargs),
depends_on=tuple(_as_node_id(item) for item in depends_on),
join=join,
label=label,
bucket_key=bucket_key,
tags=frozenset(tags),
retry=retry,
timeout=timeout,
)
self._nodes.append(node)
self._by_id[resolved_id] = node
return node
def _next_id(self, task: str) -> str:
self._counts[task] = self._counts.get(task, 0) + 1
index = self._counts[task]
return task if index == 1 else f"{task}#{index}"
@property
def nodes(self) -> tuple[Node, ...]:
"""按定义顺序返回全部节点。"""
return tuple(self._nodes)
def __len__(self) -> int:
return len(self._nodes)
def node_by_id(self, node_id: str) -> Node:
"""按 id 取节点；不存在即抛错。"""
try:
return self._by_id[node_id]
except KeyError:
raise ConfigurationError(f"管道 {self.name!r} 里没有节点 {node_id!r}") from None
# ------------------------------------------------------------------
# 校验
# ------------------------------------------------------------------
def validate(self) -> None:
"""定义期的静态校验。**在创建任何任务之前**调用，避免留下半条管道。"""
if not self._nodes:
raise ConfigurationError(f"管道 {self.name!r} 没有任何节点")
if self.aggregation.mode is AggregationMode.QUORUM:
# 诚实拒绝而不是发明语义：quorum 是"同一件事跑 N 份、成功 k 份即可"，
# 它属于**一组同质节点**，而不是一条节点各异、依赖交错的 DAG。
# 硬套到 DAG 上只能得到一个没人能预测含义的数字。
raise ConfigurationError(
f"管道 {self.name!r} 不支持 quorum 结算：quorum 描述的是"
"「同一件事的 N 份副本成功 k 份」，它属于扇出分组而不是整条 DAG。"
"请改用 fail_fast 或 all_settled。"
)
for node in self._nodes:
for dependency in node.depends_on:
if dependency not in self._by_id:
raise ConfigurationError(
f"节点 {node.node_id!r} 依赖了不存在的节点 {dependency!r}"
)
if node.node_id in node.depends_on:
raise ConfigurationError(f"节点 {node.node_id!r} 依赖了自己")
self._check_acyclic()
def _check_acyclic(self) -> None:
"""DFS 找环。
必须显式检测：有环时调度器只会安静地什么也不做（所有节点都在等对方），
表现为"管道卡住了"而不是"定义错了"——那种故障最难查。
"""
state: dict[str, int] = {} # 0=未访问 1=访问中 2=已完成
path: list[str] = []
def visit(node_id: str) -> None:
if state.get(node_id) == 2:
return
if state.get(node_id) == 1:
cycle = " -> ".join([*path[path.index(node_id) :], node_id])
raise ConfigurationError(f"管道 {self.name!r} 存在环形依赖：{cycle}")
state[node_id] = 1
path.append(node_id)
for dependency in self._by_id[node_id].depends_on:
visit(dependency)
path.pop()
state[node_id] = 2
for node in self._nodes:
visit(node.node_id)
def topological_order(self) -> tuple[Node, ...]:
"""拓扑序。仅用于展示与调试——真正的推进由运行时状态决定。"""
order: list[Node] = []
seen: set[str] = set()
def visit(node_id: str) -> None:
if node_id in seen:
return
seen.add(node_id)
for dependency in self._by_id[node_id].depends_on:
visit(dependency)
order.append(self._by_id[node_id])
for node in self._nodes:
visit(node.node_id)
return tuple(order)
class PipelineStatus(StrEnum):
"""整条管道的状态。"""
RUNNING = "running"
SUCCEEDED = "succeeded"
FAILED = "failed"
CANCELLED = "cancelled"
@dataclass(frozen=True, slots=True)
class PipelineOutcome:
"""一次管道运行的结论。"""
run_id: str
status: PipelineStatus
nodes: Mapping[str, TaskStatus]
failure: str | None = None
@property
def ok(self) -> bool:
"""是否成功。"""
return self.status is PipelineStatus.SUCCEEDED
#: 节点落终态后算"不成功"的状态。
_BAD_TERMINALS = frozenset(
{
TaskStatus.FAILED,
TaskStatus.LOST,
TaskStatus.CANCELLED,
TaskStatus.STOPPED,
}
)
#: 会触发 fail_fast 连坐的结局。
#:
#: **``CANCELLED`` 刻意不在其中**：它是连坐的**结果**，不是原因。把它也算作触发条件
#: 会同时造成两个错误——用户主动取消整条管道会被记成"失败"；被连坐取消的节点再触发
#: 一轮连坐，因果搅成一团。
_FAIL_FAST_TRIGGERS = frozenset({TaskStatus.FAILED, TaskStatus.LOST, TaskStatus.STOPPED})
class PipelineRun:
"""一次管道运行。由 :meth:`Controller.submit_pipeline` 创建。"""
def __init__(
self,
controller: Controller,
pipeline: Pipeline,
run_id: str,
task_ids: Mapping[str, str],
) -> None:
self._controller = controller
self.pipeline = pipeline
self.run_id = run_id
self._task_ids = dict(task_ids)
self._runs: dict[str, TaskRun] = {}
self._released: set[str] = set()
self._settled = threading.Event()
self._lock = threading.RLock()
self._failure: str | None = None
self._cancelled = False
# ------------------------------------------------------------------
# 只读视图
# ------------------------------------------------------------------
def __repr__(self) -> str:
return f"<PipelineRun {self.pipeline.name} {self.run_id} {self.status.value}>"
@property
def task_ids(self) -> Mapping[str, str]:
"""节点 id → task_id。"""
return dict(self._task_ids)
def task_id_of(self, node_id: str) -> str:
"""某个节点对应的 task_id。"""
return self._task_ids[node_id]
def node_status(self, node_id: str) -> TaskStatus:
"""某个节点当前的执行状态。"""
run = self._runs.get(node_id)
if run is not None:
return run.status
snapshot = self._controller.get(self._task_ids[node_id])
return snapshot.status if snapshot is not None else TaskStatus.PENDING
@property
def status(self) -> PipelineStatus:
"""整条管道的状态，由各节点状态与结算策略推导。
**先看是否还有节点没落终态。** 顺序反了会出真问题：fail_fast 触发时 ``_failure``
已经被记下，但被连坐的节点还在收尾；此时若直接返回 ``FAILED``，等待者会被提前
唤醒，看到一个"已经失败、但还有任务在跑"的矛盾状态。
"""
with self._lock:
statuses = [self.node_status(node.node_id) for node in self.pipeline.nodes]
if any(status not in TERMINAL for status in statuses):
return PipelineStatus.RUNNING
if self._failure is not None:
return PipelineStatus.FAILED
if self._cancelled:
return PipelineStatus.CANCELLED
if any(status in _BAD_TERMINALS for status in statuses):
return PipelineStatus.FAILED
if all(status is TaskStatus.SKIPPED for status in statuses):
# 一个都没跑：既不是成功也不是失败，但管道确实结束了。
return PipelineStatus.CANCELLED
return PipelineStatus.SUCCEEDED
def outcome(self) -> PipelineOutcome:
"""当前结论。"""
return PipelineOutcome(
run_id=self.run_id,
status=self.status,
nodes={node.node_id: self.node_status(node.node_id) for node in self.pipeline.nodes},
failure=self._failure,
)
def wait(self, timeout: float | None = None) -> PipelineOutcome:
"""等到整条管道收场（或超时），返回结论。
超时**不改变任何任务状态**——"我没等到"不等于"它失败了"。
"""
self._settled.wait(timeout)
return self.outcome()
@property
def settled(self) -> bool:
"""是否已有结论。"""
return self._settled.is_set()
# ------------------------------------------------------------------
# 控制
# ------------------------------------------------------------------
def cancel(self, *, reason: str | None = None) -> int:
"""取消整条管道：所有未落终态的节点都落 ``CANCELLED``。返回处理了几个。"""
with self._lock:
self._cancelled = True
count = 0
for node in self.pipeline.nodes:
run = self._runs.get(node.node_id)
if run is None or run.snapshot.is_terminal:
continue
run.cancel(reason=reason or f"管道 {self.pipeline.name} 被取消")
count += 1
self._settle()
return count
# ------------------------------------------------------------------
# 内部：由 Controller 驱动
# ------------------------------------------------------------------
def _register(self, node_id: str, run: TaskRun) -> None:
self._runs[node_id] = run
if not self.pipeline.node_by_id(node_id).depends_on:
with self._lock:
self._released.add(node_id)
def _on_task_finished(self, task_id: str) -> None:
"""某个节点落终态：先看要不要连坐，再放行/跳过下游。"""
node_id = self._node_id_of(task_id)
if node_id is None:
return
run = self._runs.get(node_id)
if run is not None and run.status in _FAIL_FAST_TRIGGERS:
self._maybe_fail_fast(node_id, run.status)
self._evaluate()
def _node_id_of(self, task_id: str) -> str | None:
for node_id, candidate in self._task_ids.items():
if candidate == task_id:
return node_id
return None
def _maybe_fail_fast(self, node_id: str, status: TaskStatus) -> None:
"""首个真失败即取消其余**已经放行**的节点。
只取消"已经放行"的（排队中 / 正在跑）：**等待上游的节点交给** :meth:`_evaluate`
去 ``skip``。两者的语义不同——``cancelled`` 是"本来要跑，被连坐了"，
``skipped`` 是"上游没成功，它本来就不该跑"。把后者也标成 cancelled，事后就
分不清哪些是真被中断的工作。
"""
if self.pipeline.aggregation.mode is not AggregationMode.FAIL_FAST:
return
if status not in _FAIL_FAST_TRIGGERS:
return
with self._lock:
if self._failure is not None or self._cancelled:
return
released = set(self._released)
self._failure = f"节点 {node_id} 落 {status.value}，按 fail_fast 取消其余任务"
reason = self._failure
for node in self.pipeline.nodes:
if node.node_id not in released:
continue
run = self._runs.get(node.node_id)
if run is not None and not run.snapshot.is_terminal:
run.cancel(reason=reason)
def _evaluate(self) -> None:
"""反复放行可跑的下游、跳过注定跑不了的下游，直到没有变化。"""
while True:
statuses = {
node.node_id: self.node_status(node.node_id) for node in self.pipeline.nodes
}
with self._lock:
decisions: list[tuple[Node, bool, str]] = []
for node in self.pipeline.nodes:
if node.node_id in self._released:
continue
verdict = _decide(node, statuses)
if verdict is None:
continue
self._released.add(node.node_id)
decisions.append((node, verdict[0], verdict[1]))
if not decisions:
break
for node, runnable, detail in decisions:
run = self._runs[node.node_id]
if runnable:
# 先注入上游返回值再放行——任务一旦开始执行就注入不了了。
run.set_inputs(self._upstream_values(node))
self._controller._release(run)
else:
run.skip(reason=detail)
self._settle()
def _upstream_values(self, node: Node) -> dict[str, Any]:
"""收集某个节点全部**成功**上游的返回值。
只有成功的上游才有值：失败或跳过的上游没有返回值，塞一个 ``None`` 进去会让
下游以为"上游返回了 None"，而不是"上游没有值"。下游用
``ctx.input(节点id, 默认值)`` 取即可。
"""
values: dict[str, Any] = {}
for dependency in node.depends_on:
upstream = self._runs.get(dependency)
if upstream is None:
continue
result = upstream.result
if result is not None and result.ok:
values[dependency] = result.value
return values
def _settle(self) -> None:
"""如果已经有结论，唤醒等待者。"""
if self.status is not PipelineStatus.RUNNING:
self._settled.set()
def _decide(node: Node, statuses: Mapping[str, TaskStatus]) -> tuple[bool, str] | None:
"""判断一个节点现在该放行、该跳过，还是继续等。
返回 ``None`` 表示"还不到时候"。返回 ``(True, "")`` 放行，``(False, 原因)`` 跳过。
"""
if not node.depends_on:
return (True, "")
upstream = [statuses[dependency] for dependency in node.depends_on]
succeeded = sum(1 for status in upstream if status is TaskStatus.SUCCEEDED)
settled = all(status in TERMINAL for status in upstream)
if node.join is JoinMode.ALL_SUCCEEDED:
if succeeded == len(upstream):
return (True, "")
if settled:
failed = [
dependency
for dependency in node.depends_on
if statuses[dependency] is not TaskStatus.SUCCEEDED
]
return (False, f"上游未成功：{', '.join(failed)}")
return None
if node.join is JoinMode.ANY_SUCCEEDED:
if succeeded:
return (True, "")
if settled:
return (False, "全部上游都未成功")
return None
# ALL_SETTLED：等所有上游落终态，不管成败。
if settled:
return (True, "")
return None
def _as_node_id(item: Node | str) -> str:
return item.node_id if isinstance(item, Node) else item
def render_pipeline(run: PipelineRun) -> str:
"""把一次管道运行渲染成表格（CLI 与示例共用）。"""
rows: list[tuple[str, ...]] = [("NODE", "TASK", "STATUS", "DEPENDS_ON", "JOIN")]
for node in run.pipeline.topological_order():
upstream = ",".join(node.depends_on) if node.depends_on else "-"
rows.append(
(
node.node_id,
node.task,
run.node_status(node.node_id).value,
upstream,
node.join.value,
)
)
widths = [max(len(row[index]) for row in rows) for index in range(len(rows[0]))]
lines = [
" ".join(cell.ljust(width) for cell, width in zip(row, widths, strict=True)).rstrip()
for row in rows
]
lines.insert(1, " ".join("-" * width for width in widths))
return "\n".join(lines)
def start_pipeline(
controller: Controller,
pipeline: Pipeline,
*,
pipeline_run_id: str | None = None,
) -> PipelineRun:
"""创建一条管道的全部节点任务，并返回可等待的运行句柄。
先做**全部**校验再创建任务：否则一个写错的 DAG 会在存储里留下一堆半成品任务，
而它们看起来就像正常提交的一样。
"""
pipeline.validate()
for node in pipeline.nodes:
# 未注册即抛 ConfigurationError——控制面只接受任务名，绝不接受代码。
controller.registry.get(node.task)
run_id = pipeline_run_id if pipeline_run_id is not None else uuid.uuid4().hex[:12]
task_ids = {node.node_id: f"{run_id}:{node.node_id}" for node in pipeline.nodes}
run = PipelineRun(controller, pipeline, run_id, task_ids)
for node in pipeline.topological_order():
task_run = controller.submit(
node.task,
*node.args,
task_id=task_ids[node.node_id],
label=node.label if node.label is not None else f"{pipeline.name}/{node.node_id}",
tags=node.tags,
bucket_key=node.bucket_key,
retry=node.retry,
timeout=node.timeout,
pipeline_run_id=run_id,
deferred=bool(node.depends_on),
**dict(node.kwargs),
)
run._register(node.node_id, task_run)
controller._watch_pipeline(run)
return run
./src\chongming_task\policies.py
"""重试、超时、并发与聚合策略。
所有策略对象都不可变，可安全共享、复用并作为配置跨进程传递。
一个刻意的选择：``RetryPolicy.retry_on`` 默认为 ``(Exception,)``，
而 :class:`chongming_task.errors.TaskCancelled` 继承自 ``BaseException``，
因此**取消永远不会被当作可重试的失败**——这是协作式取消能否生效的关键。
"""
from __future__ import annotations
import hashlib
import random
from collections.abc import Mapping
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Protocol, runtime_checkable
__all__ = [
"AggregationMode",
"AggregationPolicy",
"Backoff",
"BucketPolicy",
"ConcurrencyPolicy",
"ExponentialBackoff",
"FixedBackoff",
"Jitter",
"JoinMode",
"RetryPolicy",
"Timeout",
]
class Jitter(StrEnum):
"""退避抖动策略。默认加抖动，否则大量同时失败的任务会同步重试形成风暴。"""
NONE = "none"
FULL = "full"
"""在 ``[0, delay)`` 内均匀取值。"""
EQUAL = "equal"
"""在 ``[delay/2, delay)`` 内均匀取值，保留下限。"""
@runtime_checkable
class Backoff(Protocol):
"""退避算法。``attempt`` 从 1 开始。"""
def delay_for(self, attempt: int, *, rng: random.Random | None = None) -> float:
"""第 ``attempt`` 次尝试失败后应等待的秒数。"""
...
@dataclass(frozen=True, slots=True)
class FixedBackoff:
"""固定间隔退避。"""
delay: float = 1.0
def __post_init__(self) -> None:
if self.delay < 0:
raise ValueError(f"delay 不能为负：{self.delay}")
def delay_for(self, attempt: int, *, rng: random.Random | None = None) -> float:
del rng # 固定退避无需随机源
return self.delay
@dataclass(frozen=True, slots=True)
class ExponentialBackoff:
"""指数退避：``min(base * factor ** (attempt - 1), max_delay)``，再按需施加抖动。"""
base: float = 1.0
factor: float = 2.0
max_delay: float = 300.0
jitter: Jitter = Jitter.FULL
def __post_init__(self) -> None:
if self.base < 0:
raise ValueError(f"base 不能为负：{self.base}")
if self.factor < 1:
raise ValueError(f"factor 至少为 1，收到 {self.factor}")
if self.max_delay < 0:
raise ValueError(f"max_delay 不能为负：{self.max_delay}")
def delay_for(self, attempt: int, *, rng: random.Random | None = None) -> float:
if self.base == 0:
return 0.0
raw = min(self.base * self.factor ** max(attempt - 1, 0), self.max_delay)
if self.jitter is Jitter.NONE:
return raw
source = rng if rng is not None else random
draw = source.random()
if self.jitter is Jitter.FULL:
return raw * draw
return raw * (0.5 + 0.5 * draw)
@dataclass(frozen=True, slots=True)
class RetryPolicy:
"""重试策略。
``max_attempts`` 是**总尝试次数**（含首次），因此 ``1`` 表示不重试。
"""
max_attempts: int = 1
backoff: Backoff = field(default_factory=ExponentialBackoff)
retry_on: tuple[type[BaseException], ...] = (Exception,)
give_up_on: tuple[type[BaseException], ...] = ()
def __post_init__(self) -> None:
if self.max_attempts < 1:
raise ValueError(f"max_attempts 至少为 1，收到 {self.max_attempts}")
for exc_type in (*self.retry_on, *self.give_up_on):
if not (isinstance(exc_type, type) and issubclass(exc_type, BaseException)):
raise TypeError(f"重试条件必须是异常类型，收到 {exc_type!r}")
def should_retry(self, attempt: int, exc: BaseException | None = None) -> bool:
"""第 ``attempt`` 次尝试失败后是否应重试。
``exc`` 为 ``None`` 表示失败原因不是异常（如非零退出码），此时只要还有
尝试余额就重试。
"""
if attempt >= self.max_attempts:
return False
if exc is None:
return True
if self.give_up_on and isinstance(exc, self.give_up_on):
return False
return isinstance(exc, self.retry_on)
def delay_for(self, attempt: int, *, rng: random.Random | None = None) -> float:
"""第 ``attempt`` 次尝试失败后应等待的秒数。"""
return self.backoff.delay_for(attempt, rng=rng)
@dataclass(frozen=True, slots=True)
class Timeout:
"""超时策略。三个作用域相互独立，可任意组合。
``idle`` 基于心跳判定，用于发现"没有报错但也不动"的死锁——这是其余两者
无法覆盖的情形。
"""
start: float | None = None
"""从提交到真正开始执行的等待上限。"""
per_attempt: float | None = None
"""单次尝试的总时长上限。"""
idle: float | None = None
"""两次心跳之间的最大间隔。"""
def __post_init__(self) -> None:
for name in ("start", "per_attempt", "idle"):
value = getattr(self, name)
if value is not None and value <= 0:
raise ValueError(f"Timeout.{name} 必须为正数，收到 {value}")
@property
def is_enabled(self) -> bool:
"""是否设置了任何一档超时。"""
return any(v is not None for v in (self.start, self.per_attempt, self.idle))
@dataclass(frozen=True, slots=True)
class ConcurrencyPolicy:
"""并发与资源约束：调度器（``Controller``）真正执行的那一份配置。
三个字段各管一件事，互不重叠：
* ``max_parallel`` —— **同时执行**的任务总数上限。超出部分排队（状态 ``queued``），
不是拒绝，也不是多起线程。
* ``limits`` —— 按**任务标签**的具名信号量，用来保护外部稀缺资源：一个带 ``gpu``
标签的任务，只有当该标签下当前活跃的任务数还没到 ``limits["gpu"]`` 时才会被
启动。没写进 ``limits`` 的标签不受限；一个任务带多个受限期标签时，每个都要有
余量才放行。
* ``priority`` —— **默认优先级**：:meth:`~chongming_task.controller.Controller.submit`
没给 ``priority`` 时用它；数值越大越先从就绪队列里被取走，同优先级仍按提交顺序。
两个容易误解的地方，写在这里省得猜：
* 名额与优先级都是**本进程**的账，不做跨进程协调。它保护的是本节点那一份资源
预算（例如本节点到某个库的连接数）；多节点各自配自己的份额。优先级也只在
本进程的就绪队列里比较——不做跨节点抢占。
* 统计的是**活跃**任务（已占执行槽的），**包含被暂停的**——暂停不释放执行槽，
也就不该释放资源名额。
超限的任务**排队等待**而不是失败：调度器每次有槽位释放都会重新扫描就绪队列，
因此资源一空出来它就会被取走，不需要调用方重试。
"""
max_parallel: int | None = None
limits: Mapping[str, int] = field(default_factory=dict, hash=False)
priority: int = 0
"""默认优先级；数值越大越先被调度。"""
def __post_init__(self) -> None:
if self.max_parallel is not None and self.max_parallel < 1:
raise ValueError(f"max_parallel 至少为 1，收到 {self.max_parallel}")
for name, limit in self.limits.items():
if limit < 1:
raise ValueError(f"资源 {name!r} 的并发上限至少为 1，收到 {limit}")
class JoinMode(StrEnum):
"""一个节点在**什么条件下可以开始**——即它怎么"汇合"多个上游。
注意它与 :class:`AggregationPolicy` 是**两个不同的轴**，不要混为一谈：
* ``JoinMode`` 回答"这一个节点何时该跑"（逐节点的前置条件）；
* ``AggregationPolicy`` 回答"整条管道何时算结束、要不要连坐取消"（全局结算）。
"""
ALL_SUCCEEDED = "all_succeeded"
"""全部上游成功才跑（默认）。最常见的屏障语义。"""
ALL_SETTLED = "all_settled"
"""全部上游落终态就跑，**不管成败**。用于清理、通知、汇总这类收尾节点。"""
ANY_SUCCEEDED = "any_succeeded"
"""任一上游成功就跑（先到先得）。用于多源冗余取数的场景。"""
class AggregationMode(StrEnum):
"""整条管道的结算方式。"""
FAIL_FAST = "fail_fast"
"""首个失败即取消其余未完成任务。"""
ALL_SETTLED = "all_settled"
"""等所有任务落终态，再统一判定。"""
QUORUM = "quorum"
"""成功数达到阈值即视为通过。"""
@dataclass(frozen=True, slots=True)
class AggregationPolicy:
"""管道层如何由子任务状态推导自身状态。
工厂方法刻意命名为 ``fail_fast`` / ``all_settled`` / ``of_quorum``：
若把 quorum 工厂也叫作 ``quorum``，它会与同名字段在类属性上冲突，
dataclass 的默认值会覆盖掉该方法，导致调用时报 ``NoneType is not callable``。
"""
mode: AggregationMode = AggregationMode.FAIL_FAST
quorum: int | None = None
def __post_init__(self) -> None:
if self.mode is AggregationMode.QUORUM:
if self.quorum is None or self.quorum < 1:
raise ValueError("quorum 模式必须给出 >= 1 的 quorum")
elif self.quorum is not None:
raise ValueError("仅 quorum 模式可以设置 quorum")
@classmethod
def fail_fast(cls) -> AggregationPolicy:
return cls(AggregationMode.FAIL_FAST)
@classmethod
def all_settled(cls) -> AggregationPolicy:
return cls(AggregationMode.ALL_SETTLED)
@classmethod
def of_quorum(cls, quorum: int) -> AggregationPolicy:
return cls(AggregationMode.QUORUM, quorum=quorum)
@dataclass(frozen=True, slots=True)
class BucketPolicy:
"""分桶策略：把业务键**稳定地**映射到固定数量的桶。
桶既是本地并发隔离的单位（每个桶各自限流，互不挤占），也是将来分布式分配
的单位（一台机器拥有一组桶）。两者都要求同一个键在任何进程、任何机器上都
落到同一个桶。
为什么桶数必须固定
------------------
``bucket = hash(key) % num_buckets``。改动 ``num_buckets`` 会让**几乎所有**
键改换桶位，局部性与隔离假设同时失效。因此桶数是创建时确定的常量；真要改
就得做迁移——这也是任务会同时记下 ``bucket_key`` 的原因：迁移需要原始的键。
为什么不用内置 ``hash()``
-------------------------
CPython 对 ``str`` 的 ``hash()`` 带**进程级随机盐**（``PYTHONHASHSEED``）。
同一个键在不同进程里会落到不同的桶，而且是**静默**失效——"同桶任务应当跑在
同一台机器上"这条约定会莫名其妙地不成立，却没有任何报错。这里用 ``blake2b``，
跨进程、跨机器、跨版本都稳定。
``num_buckets`` 默认取 64：粒度过粗则隔离不足，过细则单桶任务太少、调度开销
占比升高。将来要做分布式时，它应当明显大于机器数量（经验上 4~16 倍），
以便某台机器退出时能把它的桶较均匀地分给其余机器。
"""
num_buckets: int = 64
def __post_init__(self) -> None:
if self.num_buckets < 1:
raise ValueError(f"num_buckets 至少为 1，收到 {self.num_buckets}")
def bucket_for(self, key: str) -> int:
"""把业务键映射到桶号。同一个键永远得到同一个桶。"""
if not key:
raise ValueError("bucket_key 不能为空字符串")
digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
return int.from_bytes(digest, "big") % self.num_buckets
def is_valid(self, bucket: int) -> bool:
"""``bucket`` 是否落在当前策略的范围内。
跨策略校验很有用：读取用旧桶数写入的任务时，应当发现它越界，而不是
悄悄按错误的桶处理。
"""
return 0 <= bucket < self.num_buckets
@property
def buckets(self) -> range:
"""全部桶号。"""
return range(self.num_buckets)
./src\chongming_task\run.py
"""单任务受控执行。
这是运行时零件（状态机 + 取消令牌 + 心跳 + 超时 + 执行器）的**第一个可用入口**。
在调度器（DAG、并发上限、持久化恢复）落地之前，它是唯一能把它们串起来跑的东西；
落地之后它仍然是"跑一个任务"的最短路径，用于脚本、测试与简单场景。
它负责的全部事情：
* 驱动状态机走过 ``PENDING → QUEUED → STARTING → RUNNING → 终态``；
* 按 :class:`~chongming_task.policies.RetryPolicy` 自动重试并退避；
* 按 :class:`~chongming_task.policies.Timeout` 判定超时，并在超时后请求停止；
* 把外部控制（停止 / 暂停 / 恢复）转成状态机迁移。
它**不**负责：DAG、并发上限、持久化与重启恢复、CLI。那些属于后续阶段。
"""
from __future__ import annotations
import contextlib
import threading
import time
import traceback
import uuid
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
from .clock import Clock, utcnow
from .core.state_machine import ChangeObserver, StateMachine, Transition
from .errors import InvalidTransition
from .models import (
Failure,
FailureKind,
Progress,
TaskSnapshot,
TimeoutScope,
describe_serialization_problem,
)
from .policies import RetryPolicy, Timeout
from .runtime.executors.base import (
ExecutionHandle,
ExecutionOutcome,
ExecutionRequest,
Executor,
classify_exception,
)
from .runtime.executors.thread import ThreadExecutor
from .runtime.timeouts import is_expired, next_deadline
from .states import TaskStatus
__all__ = ["TaskRun", "TaskRunOutcome", "run_task"]
#: 监督循环的轮询间隔。它同时决定了停止与超时的响应粒度。
_DEFAULT_POLL_SECONDS = 0.05
#: 默认的优雅停止窗口。
_DEFAULT_GRACEFUL_STOP_SECONDS = 5.0
def _cancelled_failure(reason: str | None) -> Failure:
"""编排层级联取消所对应的失败信息。"""
return Failure(
kind=FailureKind.CANCELLED,
message=reason or "编排层级联取消",
retryable=False,
)
@dataclass(frozen=True, slots=True)
class TaskRunOutcome:
"""一次受控执行的最终结果。"""
snapshot: TaskSnapshot
value: Any = None
"""任务返回值；未成功时为 ``None``。"""
@property
def status(self) -> TaskStatus:
"""最终状态。"""
return self.snapshot.status
@property
def ok(self) -> bool:
"""是否正常完成。"""
return self.snapshot.status is TaskStatus.SUCCEEDED
@property
def failure(self) -> Failure | None:
"""失败信息；正常完成时为 ``None``。"""
return self.snapshot.failure
class TaskRun:
"""一次受监督的任务运行。
典型用法::
run = TaskRun(extract, ("s3://bucket/x",))
run.start()
run.pause() # 任务在下一个 ctx.checkpoint() 处挂起
run.resume()
run.stop() # 请求停止；优雅窗口内自行收尾
outcome = run.wait()
线程安全：``start`` / ``wait`` / ``stop`` / ``pause`` / ``resume`` 可从不同
线程调用。所有状态机写入都在同一把锁下进行，因此"状态机只有一个写者"这条
不变量仍然成立——后续调度器必须遵守同一条纪律。
"""
def __init__(
self,
target: Callable[..., Any],
args: tuple[Any, ...] = (),
*,
kwargs: dict[str, Any] | None = None,
task_id: str | None = None,
name: str | None = None,
tags: Iterable[str] = (),
bucket: int | None = None,
bucket_key: str | None = None,
pipeline_run_id: str | None = None,
priority: int = 0,
executor: Executor | None = None,
retry: RetryPolicy | None = None,
timeout: Timeout | None = None,
artifact_dir: Path | None = None,
clock: Clock | None = None,
graceful_stop: float = _DEFAULT_GRACEFUL_STOP_SECONDS,
poll: float = _DEFAULT_POLL_SECONDS,
on_finished: Callable[[TaskRunOutcome], None] | None = None,
on_change: ChangeObserver | None = None,
) -> None:
self._target = target
self._args = args
self._kwargs = dict(kwargs or {})
self._executor: Executor = executor if executor is not None else ThreadExecutor()
self._retry = retry if retry is not None else RetryPolicy()
self._timeout = timeout if timeout is not None else Timeout()
self._clock: Clock = clock if clock is not None else utcnow
self._graceful_stop = graceful_stop
self._poll = poll
self._artifact_dir = artifact_dir
self._on_finished = on_finished
self._priority = priority
resolved_id = task_id if task_id is not None else uuid.uuid4().hex[:12]
resolved_name = name if name is not None else getattr(target, "__name__", "task")
self._machine = StateMachine.create(
resolved_id,
resolved_name,
max_attempts=self._retry.max_attempts,
tags=frozenset(tags),
bucket=bucket,
bucket_key=bucket_key,
pipeline_run_id=pipeline_run_id,
clock=self._clock,
on_change=on_change,
)
self._lock = threading.RLock()
self._handle: ExecutionHandle | None = None
self._thread: threading.Thread | None = None
self._result: TaskRunOutcome | None = None
self._stop_requested = False
self._stop_reason: str | None = None
self._cancel_requested = False
self._cancel_reason: str | None = None
self._inputs: Mapping[str, Any] = {}
self._wake = threading.Event()
# ------------------------------------------------------------------
# 只读视图
# ------------------------------------------------------------------
@property
def snapshot(self) -> TaskSnapshot:
"""当前快照，可在任务运行期间安全读取。"""
with self._lock:
return self._machine.snapshot
@property
def status(self) -> TaskStatus:
"""当前状态。"""
with self._lock:
return self._machine.status
@property
def transitions(self) -> tuple[Transition, ...]:
"""完整的状态迁移历史，可审计。"""
with self._lock:
return self._machine.history
@property
def started(self) -> bool:
"""是否已经启动过。
排队中（尚未获得执行槽）的任务为 ``False``；此时调用 :meth:`wait` 会报错，
因为根本没有可等待的执行。
"""
return self._thread is not None
@property
def priority(self) -> int:
"""排队时的优先级；数值越大越先被取走（同优先级按提交顺序）。
只在被 :class:`~chongming_task.controller.Controller` 排队调度时起作用——
直接 :meth:`start` 一个任务没有队列，也就用不到它。
"""
return self._priority
def refresh_activity(self) -> TaskSnapshot:
"""把执行载体的**真实心跳**同步进快照，并返回它。
为什么必须同步：心跳由执行句柄的 :class:`~chongming_task.runtime.heartbeat.Heartbeat`
持有，而快照是另一份数据。若不同步，下一次状态迁移的全量 ``save`` 会把
``touch`` 刚写进去的心跳**覆盖回 ``None``**——同一行有两条写入路径，这是最
典型的坑，而且是静默的（看板上心跳列永远是空的）。
注意取的是**载体记录的心跳时刻**，不是"现在"：后者会让卡死的任务也显得心跳
新鲜，``idle`` 超时就永不触发。
"""
with self._lock:
handle = self._handle
if handle is None: # pragma: no cover - 尚未启动
return self.snapshot
with self._lock:
return self._machine.heartbeat(at=handle.heartbeat.last_beat)
# ------------------------------------------------------------------
# 生命周期
# ------------------------------------------------------------------
def enqueue(self) -> None:
"""标记为"已就绪、等待执行槽"（``queued``）。
由调度方（如 :class:`~chongming_task.controller.Controller`）在**排队**时
调用，而不是在真正启动时。这样"提交了但排不上队"的任务状态是 ``queued``
而非 ``pending``——后者会让人以为它还没被受理。
"""
with self._lock:
if self._machine.status is TaskStatus.PENDING:
self._machine.mark_queued()
def block(self, *, reason: str | None = None) -> None:
"""标记为"等待上游"（``blocked``）。
与 ``queued`` 的区别是它**不会**被执行器取走：上游还没满足，现在跑它没有
意义。管道调度器负责在上游就绪后调用 :meth:`unblock`。
"""
with self._lock:
if self._machine.status is TaskStatus.PENDING:
self._machine.mark_blocked(reason=reason)
def unblock(self) -> None:
"""上游已满足，回到"等待执行槽"（``blocked`` → ``queued``）。
只对 ``blocked`` 生效：对已经在跑或已落终态的任务调用它是无操作，而不是报错
——调度器的重新评估可能重入，宽进严出比抛异常更合适。
"""
with self._lock:
if self._machine.status is TaskStatus.BLOCKED:
self._machine.mark_queued()
def start(self) -> None:
"""启动受监督执行。立即返回，不等待任务结束。"""
with self._lock:
if self._thread is not None:
raise RuntimeError("该 TaskRun 已经启动过；请新建一个实例")
# 调度方可能已经用 enqueue() 标过 queued；只有仍是 pending 才需要迁移。
if self._machine.status is TaskStatus.PENDING:
self._machine.mark_queued()
self._thread = threading.Thread(
target=self._run,
name=f"run-{self._machine.snapshot.task_id}",
daemon=True,
)
self._thread.start()
def wait(self, timeout: float | None = None) -> TaskRunOutcome:
"""等待结束并返回结果。
``timeout`` 到期只抛 :class:`TimeoutError`，**不改变任务状态**——把"我没
等到"当成"任务失败"是最常见的误用。
"""
thread = self._thread
if thread is None:
raise RuntimeError("请先调用 start()")
thread.join(timeout)
if thread.is_alive():
raise TimeoutError(f"等待超时（{timeout} 秒）；任务仍在运行，状态未被改变")
result = self._result
if result is None: # pragma: no cover - 监督循环必然留下结果
raise RuntimeError("监督循环结束却没有留下结果")
return result
def stop(self, *, reason: str | None = None) -> None:
"""请求停止。不阻塞，也不直接改状态。
这里只置位信号并让令牌抛错，**状态迁移仍由监督循环（唯一写者）完成**——
否则控制线程与监督线程会并发写状态机。
"""
with self._lock:
if self._stop_requested:
return
self._stop_requested = True
self._stop_reason = reason
handle = self._handle
self._wake.set()
if handle is not None:
handle.token.request_stop(reason)
def set_inputs(self, values: Mapping[str, Any]) -> None:
"""注入上游节点的返回值（管道数据传递）。
只能在任务**开始执行之前**调用：调度器在上游满足后才放行下游，因此这是自然
的时机。任务一旦跑起来，它已经拿到上下文里的那份快照，再改也影响不到它——
与其让调用方以为改了有用，不如明确这一点。
需要显式加锁是因为状态机快照与这里的 ``_inputs`` 会被两个线程读：调度线程写、
监督线程读。
"""
with self._lock:
if self._machine.status not in (
TaskStatus.PENDING,
TaskStatus.BLOCKED,
TaskStatus.QUEUED,
):
raise RuntimeError("任务已经开始执行，不能再注入上游数据")
self._inputs = dict(values)
@property
def result(self) -> TaskRunOutcome | None:
"""已完成时的结果；尚未结束或从未启动时为 ``None``。
与 :meth:`wait` 的区别：``wait`` 会阻塞并要求任务已启动，而本属性只是读一眼
当前状态——管道调度器用它取上游的返回值。
"""
with self._lock:
return self._result
def _record_progress(self, progress: Progress) -> None:
"""进度回写：只更新快照，不碰状态与版本号。
存储落地由 :meth:`Controller._flush_activity` 经 ``touch`` 完成——进度是高
频写入，不能挤进状态迁移那条带乐观并发的路径。
"""
with self._lock:
if self._machine.snapshot.is_terminal:
return
self._machine.update_progress(progress)
def cancel(self, *, reason: str | None = None) -> None:
"""**编排层级联取消**（上游失败、整条管道被取消）。
与 :meth:`stop` 的区别**只在结局**，但那个区别很重要：
* ``stop`` 是"用户要求的停止"，任务最终落 ``STOPPED``；
* ``cancel`` 是编排层的取消，任务最终落 ``CANCELLED``。
责任方不同，告警与统计也必须能分开——"值班同学手动停的"和"上游挂了被连坐的"
是两件完全不同的事，混成一个状态就再也分不出来了。
因此这里**不置** ``_stop_requested``：监督循环据此区分"我们请求的"与"别人
取消的"（见 ``_settle``）。
"""
with self._lock:
if self._machine.snapshot.is_terminal:
return
self._cancel_requested = True
self._cancel_reason = reason
status = self._machine.status
handle = self._handle
if status in (TaskStatus.PENDING, TaskStatus.BLOCKED):
# 从未交给执行器、也不会再交出去：直接落终态，不必走中间态缓冲。
self._machine.mark_cancelled(reason=reason, failure=_cancelled_failure(reason))
return
self._wake.set()
if handle is not None:
handle.token.request_stop(reason)
def skip(self, *, reason: str | None = None) -> None:
"""按策略跳过（``pending`` / ``blocked`` / ``queued`` → ``skipped``）。
与 :meth:`cancel` 的区别在**意图**：``skip`` 是"它本来就不该跑"（上游没成功、
条件不成立），``cancel`` 是"本来要跑，但不跑了"。两者都不是失败，但把原因
混在一起，事后就分不清"被跳过"和"被连坐"。
已经占用执行资源的任务无法跳过（状态机没有那条边）：到那一步只能 ``cancel``。
"""
with self._lock:
if self._machine.status in (
TaskStatus.PENDING,
TaskStatus.BLOCKED,
TaskStatus.QUEUED,
):
self._machine.mark_skipped(reason=reason)
def pause(self, *, timeout: float = _DEFAULT_GRACEFUL_STOP_SECONDS) -> bool:
"""请求暂停并等待确认。
返回 ``False`` 表示任务未到达检查点、暂停**没有生效**：此时状态回滚为
``RUNNING``，绝不假装已暂停。
注意这是**就地暂停**——任务阻塞在 ``ctx.checkpoint()`` 里，并**继续占用
执行槽**。我们无法把任意的 Python 调用栈换出内存再换回来，所以暂停不释放
槽位；暂停大量任务可能耗尽并发额度。
"""
with self._lock:
handle = self._handle
if handle is None or self._machine.status is not TaskStatus.RUNNING:
return False
self._machine.request_pause(actor="user")
if not handle.pause(timeout=timeout):
with self._lock:
self._machine.abort_pause(reason="任务未在窗口内到达检查点")
return False
with self._lock:
try:
self._machine.mark_paused()
except InvalidTransition:
# 任务恰好在"确认暂停"与"迁移状态"之间结束了，如实放弃这次暂停。
return False
return True
def resume(self) -> None:
"""撤销暂停。
就地暂停的任务回到 ``RUNNING``；从未开始执行就被暂停的任务（``QUEUED``
或 ``RETRY_WAIT`` 期间暂停）回到 ``QUEUED``。
"""
with self._lock:
handle = self._handle
if handle is None or self._machine.status is not TaskStatus.PAUSED:
return
handle.resume()
self._machine.resume(actor="user")
self._wake.set()
# ------------------------------------------------------------------
# 监督循环：唯一写者
# ------------------------------------------------------------------
def _run(self) -> None:
try:
self._supervise()
except BaseException as exc:
# 监督循环静默死掉是最糟的失败模式：wait() 会挂住并抛出一个与真实原因
# 毫无关系的错误。这里把异常落成一次失败，并保留完整 traceback。
self._result = self._record_supervisor_crash(exc)
finally:
self._notify_finished()
def _notify_finished(self) -> None:
"""通知外部（如 Controller 的调度器）本次运行已结束、可以让出并发额度。"""
callback = self._on_finished
result = self._result
if callback is None or result is None:
return
with contextlib.suppress(Exception):
# 回调是外部代码；它抛异常不该影响任务本身的结论。
callback(result)
def _record_supervisor_crash(self, exc: BaseException) -> TaskRunOutcome:
failure = Failure(
kind=FailureKind.EXCEPTION,
message=f"监督循环内部错误：{type(exc).__name__}: {exc}",
traceback="".join(traceback.format_exception(exc)),
retryable=False,
)
with self._lock:
if not self._machine.snapshot.is_terminal:
with contextlib.suppress(InvalidTransition):
self._machine.mark_failed(failure)
return TaskRunOutcome(snapshot=self._machine.snapshot)
def _supervise(self) -> None:
attempt = 0
while True:
with self._lock:
if self._machine.snapshot.is_terminal:
# 在"已出队、尚未启动"的窗口里被取消或跳过：状态已经落定，
# 绝不能再走 mark_starting——那会撞上非法状态转移。
self._result = TaskRunOutcome(snapshot=self._machine.snapshot)
return
not_started = self._machine.status in (
TaskStatus.PENDING,
TaskStatus.QUEUED,
TaskStatus.BLOCKED,
TaskStatus.PAUSED,
)
# 编排层级联取消优先于用户停止：取消是"这条管道不跑了"，更根本。
if self._cancel_requested and not_started:
self._machine.mark_cancelled(
reason=self._cancel_reason,
failure=_cancelled_failure(self._cancel_reason),
)
self._result = TaskRunOutcome(snapshot=self._machine.snapshot)
return
# 还没开始跑就已被要求停止：直接收场，不要白白消耗一次尝试。
if self._stop_requested and not_started:
self._machine.request_stop(actor="user", reason=self._stop_reason)
self._result = TaskRunOutcome(snapshot=self._machine.snapshot)
return
attempt += 1
launch_failure = self._begin_attempt(attempt)
outcome = launch_failure if launch_failure is not None else self._await_attempt()
settled = self._settle(attempt, outcome)
if settled is not None:
self._result = settled
return
def _begin_attempt(self, attempt: int) -> ExecutionOutcome | None:
"""起一次尝试。返回 ``None`` 表示已拉起，否则返回起始就失败的结论。"""
with self._lock:
if self._machine.status is TaskStatus.PENDING:
self._machine.mark_queued()
self._machine.mark_starting(executor=self._executor.name)
snapshot = self._machine.snapshot
inputs = self._inputs
request = ExecutionRequest(
task_id=snapshot.task_id,
name=snapshot.name,
target=self._target,
attempt=attempt,
args=self._args,
kwargs=self._kwargs,
artifact_dir=self._artifact_dir,
inputs=inputs,
on_progress=self._record_progress,
clock=self._clock,
)
try:
handle = self._executor.launch(request)
except BaseException as exc:
# 执行器没能把任务拉起来：配置错（目标不可 pickle、任务没注册）、资源耗尽……
# 必须在这里接住。异常若逃进监督线程，任务会以"监督循环内部错误"收场——
# 那等于把一个**任务侧**的问题说成框架的 bug，指向完全错误的排查方向。
return ExecutionOutcome(failure=classify_exception(exc))
with self._lock:
self._handle = handle
# pid 来自句柄：同进程执行器为 None（任务就跑在本进程里），进程执行器给出
# 子进程的进程号——事后 `ps` 对得上号的只有后者。
self._machine.mark_running(pid=handle.pid)
return None
def _mark_stopping_if_needed(self) -> None:
"""把"已请求停止"落成 ``STOPPING``。
必须在**交出结果之前**调用。停止请求发出后，任务可能在下一次
``checkpoint()`` 立刻抛错结束——那样监督循环会直接从 ``handle.wait``
拿到结果，从未走过 ``STOPPING``，于是结果会被错误地记成"级联取消"而不是
"用户停止"。这个竞态在第一次运行时就真实发生了。
"""
with self._lock:
if self._machine.status in (
TaskStatus.STARTING,
TaskStatus.RUNNING,
TaskStatus.PAUSING,
):
self._machine.request_stop(actor="user", reason=self._stop_reason)
def _await_attempt(self) -> ExecutionOutcome:
"""等待本次尝试结束，期间处理停止请求与超时。"""
handle = self._must_handle()
started_at = self._clock()
while True:
# 先取到局部变量：否则 mypy 会把 self._stop_requested 收窄为 False，
# 并把循环末尾的同名检查判成不可达——而那个检查正是"任务抢在停止分支
# 之前就结束"这个竞态的兜底。
stop_requested = self._stop_requested
if stop_requested:
self._mark_stopping_if_needed()
return handle.stop(graceful_timeout=self._graceful_stop, reason=self._stop_reason)
# 暂停（含暂停中）期间不判超时：被暂停的任务不是"卡住"的，
# 而它此时本来就不会推进心跳。
with self._lock:
pause_in_flight = self._machine.status in (
TaskStatus.PAUSING,
TaskStatus.PAUSED,
)
if not pause_in_flight:
scope = self._expired_scope(handle, started_at)
if scope is not None:
limit = self._limit_for(scope)
stopped = handle.stop(
graceful_timeout=self._graceful_stop,
reason=f"{scope.value} 超时",
)
return ExecutionOutcome(
failure=Failure(
kind=FailureKind.for_timeout(scope),
message=(f"任务超时（{scope.value}）：已超过 {limit} 秒，已请求停止"),
retryable=True,
),
abandoned=stopped.abandoned,
)
outcome = handle.wait(timeout=self._poll)
if outcome is not None:
# 停止请求可能已生效但还没被上面的分支看到：任务抢先结束了。
if self._stop_requested:
self._mark_stopping_if_needed()
return outcome
def _expired_scope(self, handle: ExecutionHandle, started_at: datetime) -> TimeoutScope | None:
deadline = next_deadline(
self._timeout,
started_at=started_at,
last_beat_at=handle.heartbeat.last_beat,
)
if deadline is None or not is_expired(deadline, now=self._clock()):
return None
return deadline.scope
def _limit_for(self, scope: TimeoutScope) -> float | None:
if scope is TimeoutScope.ATTEMPT:
return self._timeout.per_attempt
if scope is TimeoutScope.IDLE:
return self._timeout.idle
return self._timeout.start
def _settle(self, attempt: int, outcome: ExecutionOutcome) -> TaskRunOutcome | None:
"""把执行结果落成状态。返回 ``None`` 表示"还要再试一次"。"""
if outcome.ok:
problem = describe_serialization_problem(outcome.value)
if problem is not None:
# 返回值会被持久化，因此必须可 JSON 序列化。**必须在写库之前拦下来**：
# 否则失败会发生在写穿路径上，而那条路径的失败只被记录、不会让任务失败
# ——结果是进程内显示 succeeded、存储里停在 running，重启恢复再把它判成
# LOST。把一个真正成功过的任务报成"结果未知"是最糟的一种错。
#
# 不重试：同一个返回值重跑多少次都还是不可序列化。
failure = Failure(
kind=FailureKind.RESULT_INVALID,
message=(
f"返回值不满足持久化契约：{problem}。"
"返回值会写入存储，请返回可 JSON 序列化的值；大数据请传指针"
"（对象存储 key、表名 + 主键、文件路径）。"
),
retryable=False,
)
with self._lock:
self._machine.mark_failed(failure)
return TaskRunOutcome(snapshot=self._machine.snapshot)
with self._lock:
self._machine.mark_succeeded(outcome={"value": outcome.value})
return TaskRunOutcome(snapshot=self._machine.snapshot, value=outcome.value)
failure = outcome.failure or Failure(kind=FailureKind.LOST, message="执行器未给出结果")
if outcome.abandoned:
# 结果未知：任务可能仍在后台运行并产生副作用，绝不能记成"已停止"。
# 失败原因必须一并记录，否则事后无从判断该重跑还是该人工核对。
with self._lock:
self._machine.mark_lost(reason=failure.message, failure=failure)
return TaskRunOutcome(snapshot=self._machine.snapshot)
if failure.kind is FailureKind.CANCELLED:
with self._lock:
if self._stop_requested:
# 我们自己请求的停止：必须走过 STOPPING，最终落 STOPPED。
self._mark_stopping_if_needed()
self._machine.mark_stopped(reason=self._stop_reason, failure=failure)
else:
# 别人取消的（编排层级联），落 CANCELLED。
self._machine.mark_cancelled(reason=failure.message, failure=failure)
return TaskRunOutcome(snapshot=self._machine.snapshot)
if failure.retryable and self._retry.should_retry(attempt):
with self._lock:
self._machine.mark_retry_wait(failure)
if not self._backoff(attempt):
with self._lock:
self._mark_stopping_if_needed()
self._machine.mark_stopped(reason=self._stop_reason, failure=failure)
return TaskRunOutcome(snapshot=self._machine.snapshot)
return None
with self._lock:
self._machine.mark_failed(failure)
return TaskRunOutcome(snapshot=self._machine.snapshot)
def _backoff(self, attempt: int) -> bool:
"""退避等待。返回 ``False`` 表示等待期间被要求停止。
用可唤醒的事件等待，而不是 ``time.sleep``：否则一次 300 秒的退避会让
``stop()`` 也等上 300 秒才生效。
"""
deadline = time.monotonic() + self._retry.delay_for(attempt)
while True:
if self._stop_requested:
return False
remaining = deadline - time.monotonic()
if remaining <= 0:
return True
self._wake.wait(min(remaining, self._poll))
def _must_handle(self) -> ExecutionHandle:
handle = self._handle
if handle is None: # pragma: no cover - 监督循环内必然已启动
raise RuntimeError("执行尚未启动")
return handle
def run_task(
target: Callable[..., Any],
/,
*args: Any,
task_id: str | None = None,
name: str | None = None,
executor: Executor | None = None,
retry: RetryPolicy | None = None,
timeout: Timeout | None = None,
artifact_dir: Path | None = None,
clock: Clock | None = None,
graceful_stop: float = _DEFAULT_GRACEFUL_STOP_SECONDS,
**kwargs: Any,
) -> TaskRunOutcome:
"""跑一个任务并等它结束，返回最终结果。
这是 :class:`TaskRun` 的阻塞式便捷包装。需要运行期控制（边跑边暂停、停止）
时请直接用 :class:`TaskRun`。
``target`` 的第一个位置参数由执行器注入执行期上下文::
def extract(ctx, source: str) -> list[Row]: ...
"""
run = TaskRun(
target,
args,
kwargs=kwargs,
task_id=task_id,
name=name,
executor=executor,
retry=retry,
timeout=timeout,
artifact_dir=artifact_dir,
clock=clock,
graceful_stop=graceful_stop,
)
run.start()
return run.wait()
./src\chongming_task\states.py
"""任务生命周期状态与合法转移表。
设计要点
--------
1. ``TaskStatus`` 只表达**生命周期阶段**。重试次数、进度、心跳健康、失败原因属于
正交维度，分别由 :mod:`chongming_task.models` 中的 ``attempt`` / ``progress`` /
``last_heartbeat_at`` / ``failure`` 承载。把它们的组合塞进同一个枚举会迅速导致
状态爆炸，并且无法表达"已超时但正在重试"这类情形。
2. 转移表以**数据**表达而非 ``if/else``，因此可被属性测试穷举验证、可内省、
可由工具生成文档。
3. 终态在表中显式声明为空集，使"终态不可离开"成为可断言的事实，而不是散落在
各处的隐式约定。复活只能经由显式的 ``retry()`` / ``reset()``。
"""
from __future__ import annotations
from collections.abc import Mapping
from enum import StrEnum
from types import MappingProxyType
from typing import Final
__all__ = [
"NON_TERMINAL",
"RESOURCE_HOLDING",
"REVIVABLE",
"TERMINAL",
"TRANSITIONS",
"TaskStatus",
"allowed_from",
"can_transition",
"is_terminal",
]
class TaskStatus(StrEnum):
"""任务生命周期状态。终态集合见 :data:`TERMINAL`。"""
# ---- 非终态 ----
PENDING = "pending"
"""已创建，依赖尚未解析。"""
BLOCKED = "blocked"
"""依赖未满足，等待上游。"""
QUEUED = "queued"
"""就绪，等待执行槽；暂停后恢复也回到此状态。"""
STARTING = "starting"
"""执行器启动中（spawn 进程 / 建立连接 / 加载资源）。"""
RUNNING = "running"
PAUSING = "pausing"
"""已请求暂停，等待任务到达静默点。"""
PAUSED = "paused"
"""已静止，不再消耗执行资源。"""
RETRY_WAIT = "retry_wait"
"""本次尝试失败，正在退避等待。"""
STOPPING = "stopping"
"""已请求停止，处于优雅停止窗口内。"""
# ---- 终态 ----
SUCCEEDED = "succeeded"
"""正常完成。"""
FAILED = "failed"
"""出错且重试耗尽，或遇到不可重试的错误。"""
STOPPED = "stopped"
"""人工主动停止。语义上不是错误，与 CANCELLED 区分责任方。"""
CANCELLED = "cancelled"
"""编排层级联取消（上游失败 / 整条管道被取消）。"""
SKIPPED = "skipped"
"""条件为假，或依策略跳过。"""
LOST = "lost"
"""执行载体失联（进程消失、心跳超时），结果未知。"""
# 兼容早期版本中语义模糊的 FINISHED：它无法区分"成功结束"与"任意终止"。
FINISHED = "succeeded"
#: 合法状态转移表：``当前状态 -> 允许的后续状态``。
#:
#: 终态显式声明为空 ``frozenset``，因此 :data:`TERMINAL` 可以由本表推导，
#: 无需两处维护（避免表与集合不一致这一最常见缺陷）。
TRANSITIONS: Final[Mapping[TaskStatus, frozenset[TaskStatus]]] = MappingProxyType(
{
# **统一规则**：控制动作是否需要中间态，取决于任务是否**已占用执行资源**。
# 未占用（PENDING/BLOCKED/QUEUED/PAUSED/RETRY_WAIT）→ 立即生效；
# 已占用（STARTING/RUNNING/PAUSING）→ 经中间态缓冲。
# 让未启动的任务显示"停止中/暂停中"是虚假信息，会让运维误以为需要等待。
TaskStatus.PENDING: frozenset(
{
TaskStatus.BLOCKED,
TaskStatus.QUEUED,
TaskStatus.SKIPPED,
TaskStatus.CANCELLED,
TaskStatus.STOPPED,
}
),
TaskStatus.BLOCKED: frozenset(
{
TaskStatus.QUEUED,
TaskStatus.SKIPPED,
TaskStatus.CANCELLED,
TaskStatus.STOPPED,
}
),
TaskStatus.QUEUED: frozenset(
{
TaskStatus.STARTING,
TaskStatus.PAUSED,
TaskStatus.SKIPPED,
TaskStatus.STOPPED,
TaskStatus.CANCELLED,
}
),
TaskStatus.STARTING: frozenset(
{
TaskStatus.RUNNING,
TaskStatus.RETRY_WAIT,
TaskStatus.PAUSING,
TaskStatus.STOPPING,
TaskStatus.FAILED,
TaskStatus.CANCELLED,
TaskStatus.LOST,
}
),
# CANCELLED 是编排层的**级联**取消，运行中的任务同样可能被级联取消
# （例如同批次的兄弟任务失败）。这与"用户主动 stop"是不同的路径：
# 后者必须经过 STOPPING，最终落 STOPPED。
TaskStatus.RUNNING: frozenset(
{
TaskStatus.SUCCEEDED,
TaskStatus.FAILED,
TaskStatus.RETRY_WAIT,
TaskStatus.PAUSING,
TaskStatus.STOPPING,
TaskStatus.CANCELLED,
TaskStatus.LOST,
}
),
# RUNNING 表示暂停请求被撤销/回滚；STOPPING 表示暂停期间收到停止请求。
TaskStatus.PAUSING: frozenset(
{
TaskStatus.PAUSED,
TaskStatus.RUNNING,
TaskStatus.STOPPING,
TaskStatus.FAILED,
TaskStatus.CANCELLED,
TaskStatus.LOST,
}
),
# 恢复的目标取决于**暂停发生在哪一层**：
# 就地暂停（STARTING/RUNNING -> PAUSING -> PAUSED）→ RUNNING，
# 任务始终占着执行槽，调用栈还在内存里；
# 未启动即暂停（QUEUED/RETRY_WAIT -> PAUSED）→ QUEUED，重新排队。
# 无 STOPPING 出边：暂停中的任务没有占用新的执行资源，stop() 直接落 STOPPED。
TaskStatus.PAUSED: frozenset(
{
TaskStatus.QUEUED,
TaskStatus.RUNNING,
TaskStatus.STOPPED,
TaskStatus.CANCELLED,
}
),
# 无 PAUSING/STOPPING 出边：退避等待中的任务没有占用执行资源，
# 暂停与停止都立即生效。
TaskStatus.RETRY_WAIT: frozenset(
{
TaskStatus.STARTING,
TaskStatus.QUEUED,
TaskStatus.PAUSED,
TaskStatus.STOPPED,
TaskStatus.FAILED,
TaskStatus.CANCELLED,
}
),
# SUCCEEDED 是"恰好跑完":停止请求发出后任务自然结束，属于合法结果。
TaskStatus.STOPPING: frozenset(
{
TaskStatus.STOPPED,
TaskStatus.SUCCEEDED,
TaskStatus.FAILED,
TaskStatus.LOST,
}
),
# ---- 终态：无出边 ----
TaskStatus.SUCCEEDED: frozenset(),
TaskStatus.FAILED: frozenset(),
TaskStatus.STOPPED: frozenset(),
TaskStatus.CANCELLED: frozenset(),
TaskStatus.SKIPPED: frozenset(),
TaskStatus.LOST: frozenset(),
}
)
#: 终态，由转移表推导（出边为空者）。
TERMINAL: Final[frozenset[TaskStatus]] = frozenset(
status for status, successors in TRANSITIONS.items() if not successors
)
#: 非终态。注意 ``frozenset(TaskStatus)`` 只包含规范成员，不含 ``FINISHED`` 之类的别名。
NON_TERMINAL: Final[frozenset[TaskStatus]] = frozenset(TaskStatus) - TERMINAL
#: 允许经 ``retry()`` 重新进入执行流程的终态。
#:
#: ``SUCCEEDED`` 与 ``SKIPPED`` 不在其中：前者重跑没有意义，后者是策略决定的结果。
REVIVABLE: Final[frozenset[TaskStatus]] = frozenset(
{TaskStatus.FAILED, TaskStatus.STOPPED, TaskStatus.CANCELLED, TaskStatus.LOST}
)
#: **占用执行资源**的非终态：这些状态下存在一个活的执行载体（线程/进程/远端作业）。
#:
#: 它是重启恢复的判据之一：这类任务若失去监督者，不是"等待中"而是**已经丢了**——
#: 载体随进程一起消失，谁也不会再推进它，只能落 ``LOST``。相对的，``queued`` /
#: ``paused`` / ``retry_wait`` 不占资源，失去监督者只是"没人接手"，不是"丢了"。
#:
#: 注意与 ``state_machine._PAUSE_IS_DEFERRED`` 的区别：那个集合更小（只有
#: ``starting`` / ``running``），语义是"暂停要等到检查点"，而不是"占着资源"——
#: 已经 ``pausing`` 的任务不必再等一次检查点。
RESOURCE_HOLDING: Final[frozenset[TaskStatus]] = frozenset(
{
TaskStatus.STARTING,
TaskStatus.RUNNING,
TaskStatus.PAUSING,
TaskStatus.STOPPING,
}
)
def is_terminal(status: TaskStatus) -> bool:
"""``status`` 是否为终态（不可能再自行变化）。"""
return status in TERMINAL
def allowed_from(status: TaskStatus) -> frozenset[TaskStatus]:
"""从 ``status`` 出发允许到达的状态集合；终态返回空集合。"""
return TRANSITIONS[status]
def can_transition(source: TaskStatus, target: TaskStatus) -> bool:
"""``source`` 是否可以合法转移到 ``target``。
仅依据转移表判断，不涉及守卫条件（如并发槽、依赖是否满足）。
"""
return target in TRANSITIONS[source]
./src\chongming_task\__init__.py
"""精细化状态控制的 Python 管道任务控制工具。
公开 API 收敛在这里，方便直接嵌进其他应用::
from chongming_task import Controller, ThreadExecutor
controller = Controller(max_parallel=4)
controller.register("extract", extract) # 先登记，之后只传名字
controller.submit("extract", "s3://bucket/x", task_id="job-1")
只跑一个任务而不需要管理多个时::
from chongming_task import run_task
outcome = run_task(extract, "s3://bucket/x")
一切都以**名字**而不是代码提交，因此 :class:`Controller` 依赖
:class:`~chongming_task.controller.TaskRegistry` 做白名单——一个能启动任意可调用
对象的入口等于远程代码执行。
存储层
------
持久化契约是 :class:`~chongming_task.stores.protocol.TaskStore`，内置
:class:`~chongming_task.stores.memory.InMemoryTaskStore` 与
:class:`~chongming_task.stores.sqlite.SqliteTaskStore`（标准库 ``sqlite3``）。
跨机部署用 :class:`~chongming_task.stores.postgres.PostgresTaskStore`——它需要
可选的 psycopg（``pip install "chongming-task[postgres]"``），因此**惰性导入**：
``from chongming_task import PostgresTaskStore`` 能用，但它不在 ``__all__`` 里，
装没装驱动都不影响 ``import chongming_task`` 本身。
``Controller(store=...)`` 一旦给出存储，本实例就成为跨进程可见的一环：状态迁移
写穿、查询读穿、对不属于自己的任务下发指令。第三方后端可用
``chongming_task.testing.store_contract.StoreContract`` 验证自己。
"""
from __future__ import annotations
from typing import Any
from .clock import Clock, utcnow
from .controller import (
Controller,
ReconcileEntry,
ReconcileOutcome,
ReconcileReport,
Registration,
StatusCounts,
TaskRegistry,
render_table,
)
from .core.state_machine import StateMachine, Transition, Trigger
from .decorators import TaskSpec, get_task_spec, task
from .errors import (
AttemptsExhausted,
CircularDependency,
ConcurrencyConflict,
ConfigurationError,
DependencyFailed,
InvalidTransition,
StoreError,
TaskCancelled,
TaskError,
TaskExecutionError,
TaskTimeout,
)
from .models import (
AttemptRecord,
Failure,
FailureKind,
Progress,
RequestedAction,
TaskSnapshot,
TimeoutScope,
)
from .pipeline import (
Node,
Pipeline,
PipelineOutcome,
PipelineRun,
PipelineStatus,
render_pipeline,
)
from .policies import (
AggregationMode,
AggregationPolicy,
Backoff,
BucketPolicy,
ConcurrencyPolicy,
ExponentialBackoff,
FixedBackoff,
Jitter,
JoinMode,
RetryPolicy,
Timeout,
)
from .run import TaskRun, TaskRunOutcome, run_task
from .runtime.cancellation import CancellationToken
from .runtime.context import RuntimeContext, TaskContext
from .runtime.executors.base import ExecutionHandle, ExecutionOutcome, Executor
from .runtime.executors.inline import InlineExecutor
from .runtime.executors.process import ProcessExecutor
from .runtime.executors.thread import ThreadExecutor
from .states import (
NON_TERMINAL,
RESOURCE_HOLDING,
REVIVABLE,
TERMINAL,
TRANSITIONS,
TaskStatus,
)
from .stores.memory import InMemoryTaskStore
from .stores.protocol import Command, CommandKind, TaskQuery, TaskStore
from .stores.sqlite import SqliteTaskStore
__version__ = "0.1.0"
def __getattr__(name: str) -> Any:
"""惰性导出可选依赖的实现（``PostgresTaskStore``）。
核心包**零第三方运行时依赖**，所以需要 psycopg 的实现不能放进 ``__all__``：
``from chongming_task import *`` 不该迫使所有使用方安装数据库驱动。真正取到
这个名字时才导入；缺驱动时 psycopg 的 ImportError 会带上安装指引。
"""
if name == "PostgresTaskStore":
from .stores.postgres import PostgresTaskStore
return PostgresTaskStore
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
__all__ = [
"NON_TERMINAL",
"RESOURCE_HOLDING",
"REVIVABLE",
"TERMINAL",
"TRANSITIONS",
"AggregationMode",
"AggregationPolicy",
"AttemptRecord",
"AttemptsExhausted",
"Backoff",
"BucketPolicy",
"CancellationToken",
"CircularDependency",
"Clock",
"Command",
"CommandKind",
"ConcurrencyConflict",
"ConcurrencyPolicy",
"ConfigurationError",
"Controller",
"DependencyFailed",
"ExecutionHandle",
"ExecutionOutcome",
"Executor",
"ExponentialBackoff",
"Failure",
"FailureKind",
"FixedBackoff",
"InMemoryTaskStore",
"InlineExecutor",
"InvalidTransition",
"Jitter",
"JoinMode",
"Node",
"Pipeline",
"PipelineOutcome",
"PipelineRun",
"PipelineStatus",
"ProcessExecutor",
"Progress",
"ReconcileEntry",
"ReconcileOutcome",
"ReconcileReport",
"Registration",
"RequestedAction",
"RetryPolicy",
"RuntimeContext",
"SqliteTaskStore",
"StateMachine",
"StatusCounts",
"StoreError",
"TaskCancelled",
"TaskContext",
"TaskError",
"TaskExecutionError",
"TaskQuery",
"TaskRegistry",
"TaskRun",
"TaskRunOutcome",
"TaskSnapshot",
"TaskSpec",
"TaskStatus",
"TaskStore",
"TaskTimeout",
"ThreadExecutor",
"Timeout",
"TimeoutScope",
"Transition",
"Trigger",
"get_task_spec",
"render_pipeline",
"render_table",
"run_task",
"task",
"utcnow",
]
./src\chongming_task\cli\main.py
"""命令行控制台。
CLI 是**另一个进程**：它看不到任务进程的内存，只能通过共享存储查询与控制。这正是
持久层存在的理由——没有它，``chongming-task ps`` 什么也看不到。
因此除 ``run``（它自己就是 worker，前台执行到结束）之外，所有子命令都要求给出存储
位置，通过 ``--db`` 或环境变量 ``CHONGMING_TASK_DB``。
退出码：
* ``0`` 命令成功；
* ``1`` 命令本身失败（参数错、任务不存在、存储打不开等）；
* ``3`` 命令成功，但被查询/被执行的任务**没有成功**（便于 shell 里判分支）。
"""
from __future__ import annotations
import argparse
import importlib
import json
import os
import sys
from collections.abc import Sequence
from datetime import datetime
from typing import Any
from ..clock import utcnow
from ..controller import Controller, render_table
from ..decorators import get_task_spec
from ..errors import ConfigurationError, TaskError
from ..models import TaskSnapshot
from ..policies import BucketPolicy, RetryPolicy, Timeout
from ..runtime.executors.base import Executor
from ..runtime.executors.process import ProcessExecutor
from ..states import TERMINAL, TaskStatus
from ..stores.codec import snapshot_to_row
from ..stores.protocol import CommandKind, TaskQuery, TaskStore
from ..stores.sqlite import SqliteTaskStore
__all__ = ["main"]
EXIT_OK = 0
EXIT_ERROR = 1
EXIT_TASK_NOT_SUCCESSFUL = 3
_ENV_DB = "CHONGMING_TASK_DB"
# ----------------------------------------------------------------------
# 参数解析
# ----------------------------------------------------------------------
def build_parser() -> argparse.ArgumentParser:
"""构造参数解析器。单独暴露是为了让测试能直接查帮助与解析结果。"""
parser = argparse.ArgumentParser(
prog="chongming-task",
description="精细化状态控制的管道任务控制台。",
epilog=(
"存储位置可用 --db 或环境变量 CHONGMING_TASK_DB 指定；"
"没有共享存储就看不到别的进程的任务。"
),
)
parser.add_argument(
"--db",
metavar="PATH",
help=(
"任务存储：SQLite 文件路径，或 postgresql:// 连接串（跨机部署，需要 "
f'"chongming-task[postgres]"). 未给出时读环境变量 {_ENV_DB}。'
),
)
parser.add_argument("--json", action="store_true", help="以 JSON 输出，便于脚本消费。")
parser.add_argument("--owner", default=None, help="本客户端的标识，写进它下发的指令里。")
parser.add_argument(
"--num-buckets",
type=int,
default=None,
metavar="N",
help=(
"集群的桶数。**必须与部署一致**——它决定 --owned-buckets 的合法性以及 buckets "
"视图的行数。不给出时按默认值（64）。"
),
)
parser.add_argument(
"--owned-buckets",
default=None,
metavar="SPEC",
help=(
"本节点负责的桶，如 '0-15,20,30'。给出后：reconcile 只清理这些桶里的僵尸任务，"
"ps --mine 只看这些桶。不给出表示不限定归属（单机语义）。"
),
)
sub = parser.add_subparsers(dest="command", required=True, metavar="COMMAND")
# --- 定义 ---
tasks = sub.add_parser("tasks", help="列出某个任务模块里注册的任务定义。")
tasks.add_argument("--tasks", required=True, metavar="MODULE", help="形如 myapp.tasks 的模块。")
tasks.set_defaults(handler=_cmd_tasks)
# --- 提交 ---
run = sub.add_parser("run", help="提交并**前台执行**一个任务，直到它落终态。")
run.add_argument("task", metavar="TASK", help="已注册的任务名。")
run.add_argument("--tasks", required=True, metavar="MODULE", help="形如 myapp.tasks 的模块。")
run.add_argument("--task-id", default=None, help="指定 task_id；默认随机生成。")
run.add_argument("--label", default=None, help="展示名；默认用注册名。")
run.add_argument(
"--bucket-key",
default=None,
help="分桶的业务键（租户、数据源……）。同一个键在任何进程都落到同一个桶。",
)
run.add_argument(
"--arg",
action="append",
default=[],
metavar="K=V",
help="任务参数，可重复。值若是合法 JSON 会按 JSON 解析，否则按字符串。",
)
run.add_argument("--max-attempts", type=int, default=1, help="总尝试次数（含首次）。")
run.add_argument(
"--timeout", type=float, default=None, metavar="SECONDS", help="单次尝试超时。"
)
run.add_argument(
"--executor",
choices=("thread", "process"),
default="thread",
help=(
"任务在哪跑。默认 thread；process 起独立子进程，因此**能被真的杀掉**"
"（优雅窗口用尽后 terminate/kill），代价是参数与返回值必须可 pickle、"
"目标函数必须在可导入的模块里。"
),
)
run.add_argument(
"--graceful-stop", type=float, default=5.0, metavar="SECONDS", help="优雅停止窗口。"
)
run.set_defaults(handler=_cmd_run)
# --- 查询 ---
ps = sub.add_parser("ps", help="列出任务。默认只显示未落终态的。")
ps.add_argument("--all", action="store_true", help="包含已落终态的任务。")
ps.add_argument(
"--status", action="append", default=[], metavar="STATUS", help="按状态过滤，可重复。"
)
ps.add_argument("--bucket", type=int, default=None, metavar="N", help="只看某个桶。")
ps.add_argument(
"--pipeline",
default=None,
metavar="RUN_ID",
help="只看某次管道运行的任务（RUN_ID 即 pipeline_run_id）。",
)
ps.add_argument(
"--mine",
action="store_true",
help="只看本节点负责的桶（需要 --owned-buckets）。",
)
ps.add_argument("--limit", type=int, default=None, metavar="N", help="最多显示多少条。")
ps.set_defaults(handler=_cmd_ps)
status = sub.add_parser("status", help="单个任务的详情与迁移历史。")
status.add_argument("task_id", metavar="TASK_ID")
status.set_defaults(handler=_cmd_status)
history = sub.add_parser("history", help="单个任务的迁移历史（可审计轨迹）。")
history.add_argument("task_id", metavar="TASK_ID")
history.set_defaults(handler=_cmd_history)
stats = sub.add_parser("stats", help="各状态计数。")
stats.set_defaults(handler=_cmd_stats)
reconcile = sub.add_parser(
"reconcile",
help="清理上次进程留下的僵尸任务（重启恢复）。启动时先跑一次。",
)
reconcile.add_argument(
"--stale-after",
type=float,
default=None,
metavar="SECONDS",
help="从未续过租的任务的宽限期；默认取租约时长的两倍。",
)
reconcile.add_argument("--dry-run", action="store_true", help="只报告，不改任何状态。")
reconcile.add_argument("--reason", default=None, help="写进审计轨迹的原因。")
reconcile.set_defaults(handler=_cmd_reconcile)
buckets = sub.add_parser("buckets", help="每个桶的任务数；标出本节点负责的桶。")
buckets.set_defaults(handler=_cmd_buckets)
# --- 控制 ---
stop = sub.add_parser("stop", help="请求停止一个任务。")
stop.add_argument("task_id", metavar="TASK_ID")
stop.add_argument("--reason", default=None, help="停止原因，会写进审计轨迹。")
stop.set_defaults(handler=_cmd_stop)
pause = sub.add_parser("pause", help="请求暂停一个任务。")
pause.add_argument("task_id", metavar="TASK_ID")
pause.add_argument(
"--timeout", type=float, default=5.0, metavar="SECONDS", help="等待确认的时长。"
)
pause.set_defaults(handler=_cmd_pause)
resume = sub.add_parser("resume", help="恢复一个被暂停的任务。")
resume.add_argument("task_id", metavar="TASK_ID")
resume.set_defaults(handler=_cmd_resume)
return parser
# ----------------------------------------------------------------------
# 公共设施
# ----------------------------------------------------------------------
def _emit(args: argparse.Namespace, *, text: str, payload: Any = None) -> None:
"""按 ``--json`` 决定输出文本还是 JSON。"""
if args.json:
sys.stdout.write(json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n")
else:
sys.stdout.write(text + "\n")
def _open_store(args: argparse.Namespace) -> TaskStore:
"""打开任务存储。
``--db`` 接受两种形态，按前缀区分：
* ``postgresql://`` / ``postgres://`` —— Postgres 连接串（跨机部署用；
需要 ``pip install "chongming-task[postgres]"``）；
* 其余 —— SQLite 文件路径。
两者走同一份 :class:`TaskStore` 契约，因此上面所有子命令对它们都成立。
"""
target = args.db or os.environ.get(_ENV_DB)
if not target:
raise ConfigurationError(
f"未指定任务存储。请用 --db PATH（SQLite 文件或 postgresql:// 连接串），"
f"或设置环境变量 {_ENV_DB}。\n"
"CLI 是另一个进程，没有共享存储就看不到任何任务。"
)
if target.startswith(("postgresql://", "postgres://")):
try:
from ..stores.postgres import PostgresTaskStore
except ImportError as exc:
raise ConfigurationError(
f"--db 给的是 Postgres 连接串，但可选依赖未安装：{exc}"
) from exc
return PostgresTaskStore(target)
return SqliteTaskStore(target)
def _bucket_policy(args: argparse.Namespace) -> BucketPolicy:
"""按 ``--num-buckets`` 构造分桶策略。
桶数是**部署范围的约定**：各节点必须一致，否则同一个键会落到不同的桶，而且不会有
任何报错。CLI 默认按 64 假设，部署用了别的值就必须显式告知。
"""
count = getattr(args, "num_buckets", None)
return BucketPolicy() if count is None else BucketPolicy(num_buckets=count)
def _controller(args: argparse.Namespace, store: Any) -> Controller:
"""按命令行选项构造控制器。
**所有需要控制器的子命令都必须走这里**：曾经 ``run`` 自己 ``Controller(...)``，
于是 ``--owned-buckets`` / ``--num-buckets`` 被静默忽略——参数接受了却不起作用，
比报错更难查。
"""
return Controller(
store=store,
owner=args.owner or _default_owner(),
bucket_policy=_bucket_policy(args),
owned_buckets=args.owned,
executor=_executor_for(args),
graceful_stop=getattr(args, "graceful_stop", 5.0),
)
def _executor_for(args: argparse.Namespace) -> Executor | None:
"""按 ``--executor`` 选执行器；``None`` 表示"交给 Controller 的默认值"。
默认值只写在一个地方（``Controller``），这里不再抄一遍——两处默认值迟早分叉。
"""
if getattr(args, "executor", None) == "process":
return ProcessExecutor()
return None
def _resolve_owned_buckets(args: argparse.Namespace) -> frozenset[int] | None:
"""解析并**校验** ``--owned-buckets``，结果存回 ``args.owned``。
在 ``main()`` 里对所有子命令统一做一次：有些命令（``ps`` / ``stats`` / ``buckets``）
直接查存储、不构造 ``Controller``，如果只在 ``Controller`` 里校验，越界的桶号在这些
命令上会被**静默接受**——筛选条件匹配不到任何东西，看起来就像"这个桶没有任务"。
"""
spec = getattr(args, "owned_buckets", None)
if spec is None:
args.owned = None
return None
resolved = _parse_bucket_spec(spec)
limit = _bucket_policy(args).num_buckets
for bucket in sorted(resolved):
if bucket < 0 or bucket >= limit:
raise ConfigurationError(
f"桶号越界：{bucket}；桶数是 {limit}（合法范围 0..{limit - 1}）。"
f"若部署的桶数不是 {limit}，请用 --num-buckets 指明。"
)
args.owned = resolved
return resolved
def _parse_bucket_spec(text: str) -> frozenset[int]:
"""解析 ``'0-15,20,30'`` 形式的桶集合。
支持区间是为了让"这个节点负责 0 到 15 号桶"这类配置写起来不难受——手写 16 个数字
既啰嗦又容易漏。
"""
buckets: set[int] = set()
for chunk in text.split(","):
piece = chunk.strip()
if not piece:
continue
if "-" in piece:
low_text, _, high_text = piece.partition("-")
try:
low, high = int(low_text), int(high_text)
except ValueError:
raise ConfigurationError(f"桶区间写法不正确：{piece!r}（应为 0-15）") from None
if low > high:
raise ConfigurationError(f"桶区间起止颠倒：{piece!r}")
buckets.update(range(low, high + 1))
continue
try:
buckets.add(int(piece))
except ValueError:
raise ConfigurationError(f"桶号必须是整数：{piece!r}") from None
if not buckets:
raise ConfigurationError(f"--owned-buckets 没有解析出任何桶：{text!r}")
return frozenset(buckets)
def _default_owner() -> str:
return f"cli:{os.getpid()}"
def _parse_statuses(raw: Sequence[str]) -> frozenset[TaskStatus] | None:
if not raw:
return None
parsed: set[TaskStatus] = set()
for item in raw:
try:
parsed.add(TaskStatus(item))
except ValueError:
known = "、".join(status.value for status in TaskStatus)
raise ConfigurationError(f"未知状态 {item!r}；可用：{known}") from None
return frozenset(parsed)
def _parse_args_pairs(pairs: Sequence[str]) -> dict[str, Any]:
"""把 ``--arg k=v`` 解析成关键字参数。
值先尝试按 JSON 解析（于是 ``n=3`` 得到整数、``flag=true`` 得到布尔），失败则按
字符串——这样不必为"传个数字"发明额外语法。
"""
result: dict[str, Any] = {}
for pair in pairs:
if "=" not in pair:
raise ConfigurationError(f"--arg 需要 K=V 形式，收到 {pair!r}")
key, _, raw = pair.partition("=")
if not key:
raise ConfigurationError(f"--arg 的键不能为空：{pair!r}")
try:
result[key] = json.loads(raw)
except json.JSONDecodeError:
result[key] = raw
return result
def _load_tasks_module(name: str) -> Any:
"""导入任务模块。
这是"任务定义从哪里来"的答案：CLI 只传名字，因此必须有一个模块把可调用对象登记
进注册表——一个能启动任意可调用对象的命令行入口等于任意代码执行。
两种写法都接受：
* 定义 ``register(controller)`` 自行登记（早期约定，最灵活）；
* 只用 ``@task`` 装饰函数，交给 :meth:`Controller.register_module` 扫描
——这是后来加的写法，CLI 也必须支持，否则 ``@task`` 在命令行下用不了。
"""
try:
module = importlib.import_module(name)
except ImportError as exc:
missing = getattr(exc, "name", None)
if missing and (name == missing or name.startswith(f"{missing}.")):
# 模块自身没找到：最常见的一种，也是最不该让人猜的一种。
raise ConfigurationError(
f"无法导入任务模块 {name!r}：{exc}\n"
"任务模块必须在 sys.path 上：要么是**已安装的包**，要么把它所在的目录加进 "
"PYTHONPATH。\n"
"注意：`chongming-task` 是入口脚本，**不会**把当前目录放进 sys.path——"
"放在当前目录的模块请用：\n"
f" PYTHONPATH=. chongming-task ... --tasks {name}\n"
f"或 uv run python -m chongming_task.cli.main ... --tasks {name}"
) from exc
# 模块找到了，但模块**内部**导入失败——那是它自己的依赖问题，别把矛头指向路径。
raise ConfigurationError(f"任务模块 {name!r} 找到了，但它导入失败：{exc}") from exc
if not callable(getattr(module, "register", None)) and not _has_decorated_tasks(module):
raise ConfigurationError(
f"任务模块 {name!r} 里既没有 register(controller) 函数，也没有用 @task 装饰过的函数。"
)
return module
def _has_decorated_tasks(module: Any) -> bool:
"""模块里是否存在 ``@task`` 装饰过的函数。"""
return any(get_task_spec(member) is not None for member in vars(module).values())
def _register_tasks(module: Any, controller: Controller) -> None:
"""把模块里的任务登记进控制器：优先 ``register()``，否则扫描 ``@task``。"""
register = getattr(module, "register", None)
if callable(register):
register(controller)
return
controller.register_module(module)
def _format_elapsed(snapshot: TaskSnapshot, now: datetime) -> str:
if snapshot.started_at is None:
return "-"
end = snapshot.ended_at or now
return f"{(end - snapshot.started_at).total_seconds():.1f}s"
def _format_detail(snapshot: TaskSnapshot, *, now: datetime | None = None) -> str:
"""``status`` 的文本视图。
刻意把"正交维度"分开展示：状态是一件事，重试、进度、健康、失败原因是另外几件。
"""
moment = now if now is not None else utcnow()
failure = snapshot.failure
rows: list[tuple[str, str]] = [
("task_id", snapshot.task_id),
("name", snapshot.name),
("status", snapshot.status.value),
("bucket", "-" if snapshot.bucket is None else str(snapshot.bucket)),
("bucket_key", snapshot.bucket_key or "-"),
("attempt", f"{snapshot.attempt}/{snapshot.max_attempts}"),
("status_version", str(snapshot.status_version)),
("executor", snapshot.executor or "-"),
("created_at", _iso(snapshot.created_at)),
("started_at", _iso(snapshot.started_at)),
("ended_at", _iso(snapshot.ended_at)),
("last_heartbeat_at", _iso(snapshot.last_heartbeat_at)),
("lease_owner", snapshot.lease_owner or "-"),
("lease_until", _lease_text(snapshot, moment)),
("elapsed", _format_elapsed(snapshot, moment)),
("progress", "-" if snapshot.progress is None else _progress_text(snapshot)),
("failure", "-" if failure is None else f"{failure.kind.value}: {failure.message}"),
("retryable", "-" if failure is None else str(failure.retryable)),
("overridden", str(snapshot.overridden)),
]
if snapshot.tags:
rows.insert(4, ("tags", ", ".join(sorted(snapshot.tags))))
width = max(len(key) for key, _ in rows)
return "\n".join(f"{key:>{width}} : {value}" for key, value in rows)
def _lease_text(snapshot: TaskSnapshot, now: datetime) -> str:
"""监督者租约的**可判读**形式。
为什么不能只显示 ``last_heartbeat_at``：一个"看起来在跑"的任务可能早就没有监督者了
（进程崩了、被杀了），而心跳停在半路与"任务卡住"长得一模一样。租约过期才是**没有人在管
它**的判据。这正是本命令最容易被误读的地方，所以直接给出结论而不是原始时间。
"""
if snapshot.lease_until is None:
return "-（从未续租）" if not snapshot.is_terminal else "-"
remaining = (snapshot.lease_until - now).total_seconds()
stamp = _iso(snapshot.lease_until)
if snapshot.is_terminal:
# 终态任务留着最后一任监督者，纯属审计信息。
return f"{stamp}（任务已结束）"
if remaining > 0:
return f"{stamp}（剩余 {remaining:.0f}s）"
return f"{stamp} ⚠ 已过期 {-remaining:.0f}s —— 没有监督者在管它，reconcile 会判它 lost"
def _progress_text(snapshot: TaskSnapshot) -> str:
progress = snapshot.progress
if progress is None: # pragma: no cover - 调用方已判空
return "-"
if progress.total is None:
return str(progress.done)
return f"{progress.done}/{progress.total}"
def _iso(value: datetime | None) -> str:
return "-" if value is None else value.isoformat(timespec="seconds")
def _render_history(records: Sequence[Any]) -> str:
"""迁移历史的表格视图。"""
if not records:
return "(没有迁移记录)"
rows: list[tuple[str, ...]] = [("VER", "FROM", "TO", "AT", "TRIGGER", "ACTOR", "REASON")]
for record in records:
rows.append(
(
str(record.version),
record.from_status.value,
record.to_status.value,
_iso(record.at),
record.trigger,
record.actor,
record.reason or "-",
)
)
widths = [max(len(row[index]) for row in rows) for index in range(len(rows[0]))]
lines = [
" ".join(cell.ljust(width) for cell, width in zip(row, widths, strict=True)).rstrip()
for row in rows
]
lines.insert(1, " ".join("-" * width for width in widths))
return "\n".join(lines)
def _snapshot_json(snapshot: TaskSnapshot) -> dict[str, Any]:
"""快照的 JSON 视图。
直接用存储层的行编码：输出与存储同构、全部是原语，因此脚本消费它不需要了解
``datetime`` / 枚举的序列化规则。
"""
return snapshot_to_row(snapshot)
# ----------------------------------------------------------------------
# 子命令
# ----------------------------------------------------------------------
def _cmd_tasks(args: argparse.Namespace) -> int:
module = _load_tasks_module(args.tasks)
probe = Controller()
_register_tasks(module, probe)
registrations = probe.registry.all()
payload = [
{
"name": item.name,
"description": item.description,
"tags": sorted(item.tags),
"max_attempts": item.retry.max_attempts,
}
for item in registrations
]
lines = ["NAME TAGS MAX_ATTEMPTS DESCRIPTION"]
lines.append("-" * 60)
for item in registrations:
tags = ",".join(sorted(item.tags)) or "-"
lines.append(f"{item.name} {tags} {item.retry.max_attempts} {item.description}")
if not registrations:
lines.append("(该模块没有注册任何任务)")
_emit(args, text="\n".join(lines), payload=payload)
return EXIT_OK
def _cmd_run(args: argparse.Namespace) -> int:
module = _load_tasks_module(args.tasks)
store = _open_store(args)
controller = _controller(args, store)
try:
_register_tasks(module, controller)
run = controller.submit(
args.task,
task_id=args.task_id,
label=args.label,
bucket_key=args.bucket_key,
retry=RetryPolicy(max_attempts=args.max_attempts),
timeout=Timeout(per_attempt=args.timeout),
**_parse_args_pairs(args.arg),
)
task_id = run.snapshot.task_id
if not args.json:
sys.stderr.write(f"已提交 {args.task} -> {task_id}（前台执行，Ctrl-C 可中断）\n")
outcome = controller.wait(task_id, timeout=None)
finally:
controller.close(close_store=True)
snapshot = outcome.snapshot
_emit(
args,
text=_format_detail(snapshot),
payload=_snapshot_json(snapshot),
)
# 写穿失败意味着**存储里的状态与上面的结论不一致**。这必须喊出来：否则运维在
# `status` 里看到的是"还在 running"，而这里的进程说 succeeded——两边都不算错，
# 却对不上，而且没有任何线索指向原因。
_warn_about_persistence(args, controller)
return EXIT_OK if outcome.ok else EXIT_TASK_NOT_SUCCESSFUL
def _warn_about_persistence(args: argparse.Namespace, controller: Controller) -> None:
"""把写穿失败与永久分叉报到 stderr。
这些信息**只存在于写进程的内存里**，`status` / `ps` 在另一个进程里看不到它们——
所以必须由 `run` 自己说出来。
"""
if args.json:
return
failures = controller.persist_failures
if failures:
sys.stderr.write(
f"\n警告：{len(failures)} 次写穿失败——**存储里的状态落后于上面的结论**。\n"
" 用 `status` 观察到的可能是旧状态；最近几次失败：\n"
)
for item in failures[-5:]:
sys.stderr.write(f" {item}\n")
for task_id, reason in controller.desynced_tasks.items():
sys.stderr.write(
f"\n警告：{task_id} 的存储状态已**永久分叉**（已停止写穿以免覆盖对方）：\n"
f" {reason}\n"
" 该任务的存储状态不会再更新，请人工核对。\n"
)
def _cmd_ps(args: argparse.Namespace) -> int:
store = _open_store(args)
try:
statuses = _parse_statuses(args.status)
if statuses is None and not args.all:
# 「在跑的东西」通常指未落终态——包括排队与暂停。
statuses = frozenset(status for status in TaskStatus if status not in TERMINAL)
wanted_buckets: frozenset[int] | None = None
if args.mine:
wanted_buckets = args.owned
if wanted_buckets is None:
raise ConfigurationError("--mine 需要同时给出 --owned-buckets")
elif args.bucket is not None:
wanted_buckets = frozenset({args.bucket})
snapshots = store.query(
TaskQuery(
statuses=statuses,
buckets=wanted_buckets,
pipeline_run_id=args.pipeline,
limit=args.limit,
)
)
finally:
store.close()
if not snapshots:
_emit(args, text="(没有匹配的任务)", payload=[])
return EXIT_OK
_emit(
args,
text=render_table(snapshots),
payload=[_snapshot_json(item) for item in snapshots],
)
return EXIT_OK
def _cmd_status(args: argparse.Namespace) -> int:
store = _open_store(args)
try:
snapshot = store.get(args.task_id)
if snapshot is None:
raise ConfigurationError(f"没有这个任务：{args.task_id}")
records = store.history(args.task_id)
finally:
store.close()
detail = _format_detail(snapshot)
text = f"{detail}\n\n迁移历史：\n{_render_history(records)}"
payload = {
"task": _snapshot_json(snapshot),
"history": [
{
"version": record.version,
"from": record.from_status.value,
"to": record.to_status.value,
"at": record.at.isoformat(),
"trigger": record.trigger,
"actor": record.actor,
"reason": record.reason,
}
for record in records
],
}
_emit(args, text=text, payload=payload)
return EXIT_OK
def _cmd_history(args: argparse.Namespace) -> int:
store = _open_store(args)
try:
records = store.history(args.task_id)
finally:
store.close()
_emit(
args,
text=_render_history(records),
payload=[
{
"version": record.version,
"from": record.from_status.value,
"to": record.to_status.value,
"at": record.at.isoformat(),
"trigger": record.trigger,
"actor": record.actor,
"reason": record.reason,
}
for record in records
],
)
return EXIT_OK
def _cmd_stats(args: argparse.Namespace) -> int:
store = _open_store(args)
try:
counts = store.counts()
finally:
store.close()
lines = [
f"{status.value:<12} {count}"
for status, count in sorted(counts.items(), key=lambda kv: kv[0].value)
]
total = sum(counts.values())
unfinished = sum(count for status, count in counts.items() if status not in TERMINAL)
lines.append("-" * 20)
lines.append(f"{'total':<12} {total}")
lines.append(f"{'unfinished':<12} {unfinished}")
_emit(
args,
text="\n".join(lines),
payload={
"by_status": {status.value: count for status, count in counts.items()},
"total": total,
"unfinished": unfinished,
},
)
return EXIT_OK
def _cmd_reconcile(args: argparse.Namespace) -> int:
store = _open_store(args)
try:
controller = _controller(args, store)
report = controller.reconcile(
stale_after=args.stale_after,
reason=args.reason,
dry_run=args.dry_run,
)
finally:
store.close()
payload = {
"dry_run": report.dry_run,
"entries": [
{
"task_id": item.task_id,
"previous_status": item.previous_status.value,
"outcome": item.outcome.value,
"detail": item.detail,
}
for item in report.entries
],
}
if not report:
_emit(args, text="没有需要处理的任务。", payload=payload)
return EXIT_OK
rows: list[tuple[str, ...]] = [("OUTCOME", "TASK_ID", "WAS", "DETAIL")]
for item in report.entries:
rows.append((item.outcome.value, item.task_id, item.previous_status.value, item.detail))
widths = [max(len(row[index]) for row in rows) for index in range(len(rows[0]))]
lines = [
" ".join(cell.ljust(width) for cell, width in zip(row, widths, strict=True)).rstrip()
for row in rows
]
lines.insert(1, " ".join("-" * width for width in widths))
lines.append("")
if report.dry_run:
lines.append(
f"预演：有 {len(report.lost)} 个任务**会被**标记为 lost，"
f"{len(report.adoptable)} 个等待重新提交。未写入任何状态。"
)
else:
lines.append(
f"已把 {len(report.lost)} 个僵尸任务标记为 lost，"
f"{len(report.adoptable)} 个任务无人接手（状态未改动，需要重新提交）。"
)
if report.skipped:
lines.append(f"{len(report.skipped)} 个被跳过（有人正在写）。")
lines.append("注意：本命令不会、也无法重新执行任务——存储里记的是状态，不是参数。")
_emit(args, text="\n".join(lines), payload=payload)
return EXIT_OK
def _cmd_buckets(args: argparse.Namespace) -> int:
store = _open_store(args)
try:
snapshots = store.query()
finally:
store.close()
counts: dict[int | None, int] = {}
for snapshot in snapshots:
counts[snapshot.bucket] = counts.get(snapshot.bucket, 0) + 1
num_buckets = _bucket_policy(args).num_buckets
mine = args.owned
lines = [f"{'':<1} {'BUCKET':>6} {'TASKS':>5}"]
lines.append("-" * 16)
for bucket in range(num_buckets):
# 星号标出本节点负责的桶——"我这一份"是自治架构里最常看的视图。
mark = "*" if mine is not None and bucket in mine else " "
lines.append(f"{mark} {bucket:>6} {counts.get(bucket, 0):>5}")
if None in counts:
lines.append(f"{' '} {'(none)':>6} {counts[None]:>5}")
_emit(
args,
text="\n".join(lines),
payload={
"num_buckets": num_buckets,
"by_bucket": {str(bucket): count for bucket, count in counts.items()},
"owned_buckets": None if mine is None else sorted(mine),
},
)
return EXIT_OK
def _route_control(
args: argparse.Namespace, kind: CommandKind, *, reason: str | None = None
) -> bool:
"""把一个控制动作交给 Controller 路由（本地直下，远程写成指令）。"""
store = _open_store(args)
try:
controller = _controller(args, store)
if kind is CommandKind.STOP:
return controller.stop(args.task_id, reason=reason)
if kind is CommandKind.PAUSE:
confirmed = controller.pause(args.task_id, timeout=getattr(args, "timeout", 5.0))
if not confirmed:
# 两种可能：任务不属于本进程（无法确认），或任务没到达检查点。
sys.stderr.write(
"暂停**未确认**。若任务由别的进程持有，这是正常的——"
"跨进程没有回执通道，请用 status 观察结果。\n"
)
return confirmed or controller.get(args.task_id) is not None
return controller.resume(args.task_id)
finally:
store.close()
def _cmd_stop(args: argparse.Namespace) -> int:
handled = _route_control(args, CommandKind.STOP, reason=args.reason)
if not handled:
raise ConfigurationError(f"没有找到任务 {args.task_id}；无法停止。")
_emit(
args,
text=f"已请求停止 {args.task_id}" + (f"（原因：{args.reason}）" if args.reason else ""),
payload={"task_id": args.task_id, "requested": "stop", "reason": args.reason},
)
return EXIT_OK
def _cmd_pause(args: argparse.Namespace) -> int:
handled = _route_control(args, CommandKind.PAUSE)
if not handled:
raise ConfigurationError(f"没有找到任务 {args.task_id}；无法暂停。")
_emit(
args,
text=f"已请求暂停 {args.task_id}（用 status 观察是否真的挂起）",
payload={"task_id": args.task_id, "requested": "pause"},
)
return EXIT_OK
def _cmd_resume(args: argparse.Namespace) -> int:
handled = _route_control(args, CommandKind.RESUME)
if not handled:
raise ConfigurationError(f"没有找到任务 {args.task_id}；无法恢复。")
_emit(
args,
text=f"已请求恢复 {args.task_id}",
payload={"task_id": args.task_id, "requested": "resume"},
)
return EXIT_OK
# ----------------------------------------------------------------------
# 入口
# ----------------------------------------------------------------------
def main(argv: Sequence[str] | None = None) -> int:
"""命令行入口。返回退出码（见模块文档）。"""
parser = build_parser()
args = parser.parse_args(list(argv) if argv is not None else None)
try:
# 无论哪个子命令，都先解析并校验一次桶归属：有些命令直接查存储、不构造
# Controller，把校验放在 Controller 里会让越界桶号在这些命令上被静默接受。
_resolve_owned_buckets(args)
return int(args.handler(args))
except (ConfigurationError, TaskError) as exc:
sys.stderr.write(f"错误：{exc}\n")
return EXIT_ERROR
except KeyboardInterrupt:
sys.stderr.write("已中断\n")
return EXIT_ERROR
except BrokenPipeError: # pragma: no cover - 管道被下游关闭
return EXIT_OK
if __name__ == "__main__": # pragma: no cover
raise SystemExit(main())
./src\chongming_task\cli\__init__.py
"""命令行控制台。
``chongming-task`` 是**另一个进程**：它只能通过共享存储看到任务。除 ``run``（自身即
worker，前台执行到结束）之外，所有子命令都需要 ``--db`` 或 ``CHONGMING_TASK_DB``。
用法与退出码见 :mod:`chongming_task.cli.main`。
"""
from __future__ import annotations
from .main import build_parser, main
__all__ = ["build_parser", "main"]
./src\chongming_task\core\state_machine.py
"""任务状态机——任务状态的**唯一写入者**。
不变量
------
1. 任何状态变更都经转移表校验，非法转移抛 :class:`~chongming_task.errors.InvalidTransition`，
而不是静默忽略或返回 ``False``——"停止指令没生效"是最难排查的一类问题。
2. 每次成功变更都递增 ``status_version`` 并追加一条 :class:`Transition`，
因此状态历史只追加、可审计，并可作为存储层乐观并发的依据。
3. 终态不可离开；复活只能走显式的 :meth:`StateMachine.retry` / :meth:`StateMachine.reset`。
4. 时间全部取自注入的 ``clock``，使重试/超时相关测试无需真实等待。
不负责的事
----------
* **线程安全**。调用方（调度器）必须保证同一任务只有一个写入者。
* 依赖解析、并发槽、执行器启动。这些由调度器决定，状态机只记录结果。
"""
from __future__ import annotations
from collections.abc import Callable, Mapping
from dataclasses import dataclass, replace
from datetime import datetime
from enum import StrEnum
from typing import Any
from ..clock import Clock, utcnow
from ..errors import AttemptsExhausted, InvalidTransition
from ..models import Failure, Progress, RequestedAction, TaskSnapshot
from ..states import REVIVABLE, TERMINAL, TRANSITIONS, TaskStatus
__all__ = ["ChangeObserver", "Clock", "StateMachine", "Transition", "Trigger", "utcnow"]
class Trigger(StrEnum):
"""状态变更的标准触发源。
用枚举而非自由字符串，避免审计日志里出现 ``user:stop`` / ``stop`` /
``user_stop`` 之类的拼写变体——一旦出现，"按触发源统计"就彻底失效了。
"""
DEPENDENCIES_UNMET = "orchestrator:dependencies_unmet"
DEPENDENCIES_RESOLVED = "orchestrator:dependencies_resolved"
DISPATCH = "scheduler:dispatch"
LAUNCHED = "executor:launched"
COMPLETED = "executor:completed"
FAILED = "executor:failed"
RETRY_SCHEDULED = "retry:scheduled"
RETRY_EXHAUSTED = "retry:exhausted"
RETRY_REQUESTED = "user:retry"
RESET_REQUESTED = "user:reset"
STOP_REQUESTED = "user:stop"
STOP_CONFIRMED = "runtime:stop_confirmed"
STOP_FORCED = "runtime:stop_forced"
PAUSE_REQUESTED = "user:pause"
PAUSE_REACHED = "runtime:pause_reached"
PAUSE_ABORTED = "runtime:pause_aborted"
RESUME_REQUESTED = "user:resume"
CANCEL_REQUESTED = "orchestrator:cancel"
SKIP_REQUESTED = "orchestrator:skip"
HEARTBEAT_LOST = "runtime:heartbeat_lost"
ADMIN_OVERRIDE = "admin:override"
#: **统一规则**：控制动作是否需要中间态，取决于任务是否**已占用执行资源**。
#: 未占用（PENDING/BLOCKED/QUEUED/PAUSED/RETRY_WAIT）→ 立即生效；
#: 已占用（STARTING/RUNNING/PAUSING）→ 经中间态缓冲。
#: 让未启动的任务显示"停止中/暂停中"是虚假信息，会让运维误以为需要等待。
_STOP_IS_DEFERRED: frozenset[TaskStatus] = frozenset(
{TaskStatus.STARTING, TaskStatus.RUNNING, TaskStatus.PAUSING}
)
_PAUSE_IS_DEFERRED: frozenset[TaskStatus] = frozenset({TaskStatus.STARTING, TaskStatus.RUNNING})
#: 状态变更观察者：``(变更后的快照, 本次迁移记录)``。
#:
#: **契约：本回调不得抛异常。** 它在状态已经落定之后被调用；若它抛出去，调用方会
#: 看到一个"状态已变但操作失败"的中间态。需要处理失败的实现（例如写存储失败）应当
#: 自己捕获并记录下来。
ChangeObserver = Callable[[TaskSnapshot, "Transition"], None]
@dataclass(frozen=True, slots=True)
class Transition:
"""一次已生效的状态变更记录。只追加，不修改。"""
task_id: str
from_status: TaskStatus
to_status: TaskStatus
at: datetime
trigger: str
actor: str = "system"
reason: str | None = None
version: int = 0
"""变更后的 ``status_version``，用于与存储层比对。"""
@property
def is_override(self) -> bool:
"""是否为人工强制改写。此类记录在审计视图中必须单独可见。"""
return self.trigger == Trigger.ADMIN_OVERRIDE
class StateMachine:
"""持有单个任务的状态，并作为其唯一写入者。"""
__slots__ = ("_clock", "_history", "_on_change", "_snapshot")
def __init__(
self,
snapshot: TaskSnapshot,
*,
clock: Clock | None = None,
on_change: ChangeObserver | None = None,
) -> None:
self._snapshot = snapshot
self._clock: Clock = clock if clock is not None else utcnow
self._history: list[Transition] = []
self._on_change = on_change
@classmethod
def create(
cls,
task_id: str,
name: str,
*,
max_attempts: int = 1,
pipeline_run_id: str | None = None,
tags: frozenset[str] = frozenset(),
bucket: int | None = None,
bucket_key: str | None = None,
clock: Clock | None = None,
on_change: ChangeObserver | None = None,
) -> StateMachine:
"""创建处于 ``PENDING`` 的新任务。"""
now = (clock if clock is not None else utcnow)()
snapshot = TaskSnapshot(
task_id=task_id,
name=name,
status=TaskStatus.PENDING,
created_at=now,
max_attempts=max_attempts,
pipeline_run_id=pipeline_run_id,
tags=tags,
bucket=bucket,
bucket_key=bucket_key,
)
return cls(snapshot, clock=clock, on_change=on_change)
# ------------------------------------------------------------------
# 只读视图
# ------------------------------------------------------------------
@property
def snapshot(self) -> TaskSnapshot:
"""当前快照。不可变，可安全外传。"""
return self._snapshot
@property
def status(self) -> TaskStatus:
"""当前状态。"""
return self._snapshot.status
@property
def history(self) -> tuple[Transition, ...]:
"""本实例产生过的全部转移记录。"""
return tuple(self._history)
def can(self, target: TaskStatus) -> bool:
"""能否合法转移到 ``target``（仅依据转移表）。"""
return target in TRANSITIONS[self._snapshot.status]
@property
def can_stop(self) -> bool:
"""能否停止。供控制面预检，使 UI 能提前置灰按钮。"""
return self.can(TaskStatus.STOPPING) or self.can(TaskStatus.STOPPED)
@property
def can_pause(self) -> bool:
"""能否暂停。未启动的任务立即暂停，运行中的任务经 ``PAUSING`` 缓冲。"""
return self.can(TaskStatus.PAUSING) or self.can(TaskStatus.PAUSED)
@property
def can_resume(self) -> bool:
"""能否恢复。目标状态由 :meth:`resume` 根据暂停层级判定。"""
return self.status is TaskStatus.PAUSED
@property
def can_retry(self) -> bool:
"""能否经 :meth:`retry` 复活。"""
return self.status in REVIVABLE
# ------------------------------------------------------------------
# 写入核心
# ------------------------------------------------------------------
def _apply(
self,
target: TaskStatus,
trigger: str,
*,
actor: str = "system",
reason: str | None = None,
enforce_table: bool = True,
**changes: Any,
) -> Transition:
"""写入路径的唯一实现。
``enforce_table=False`` 仅供 :meth:`retry` / :meth:`reset` /
:meth:`force_transition` 使用——这三者会自行完成更严格的守卫与审计，
而不受转移表约束。
"""
current = self._snapshot
if enforce_table and target not in TRANSITIONS[current.status]:
raise InvalidTransition(current.status, target, action=str(trigger))
now = self._clock()
if target in TERMINAL:
changes.setdefault("ended_at", now)
changes.setdefault("requested_action", None)
new = replace(
current,
status=target,
status_version=current.status_version + 1,
**changes,
)
record = Transition(
task_id=new.task_id,
from_status=current.status,
to_status=target,
at=now,
trigger=str(trigger),
actor=actor,
reason=reason,
version=new.status_version,
)
self._snapshot = new
self._history.append(record)
if self._on_change is not None:
# 契约：观察者不得抛异常（见 ChangeObserver）。挂在 _apply 上而不是各
# 业务方法里，是为了保证**任何**一次迁移都不会漏掉——漏一次就意味着
# 存储与内存悄悄分叉。
self._on_change(new, record)
return record
# ------------------------------------------------------------------
# 调度推进（由调度器/执行器调用）
# ------------------------------------------------------------------
def mark_blocked(self, *, reason: str | None = None) -> Transition:
"""依赖未满足，等待上游。"""
return self._apply(TaskStatus.BLOCKED, Trigger.DEPENDENCIES_UNMET, reason=reason)
def mark_queued(self, *, reason: str | None = None) -> Transition:
"""依赖已满足，进入就绪队列等待执行槽。"""
return self._apply(
TaskStatus.QUEUED,
Trigger.DEPENDENCIES_RESOLVED,
reason=reason,
queued_at=self._clock(),
)
def mark_starting(self, *, executor: str, reason: str | None = None) -> Transition:
"""开始一次新的尝试，同时递增 ``attempt``。
重试预算在此处统一校验，而不是分散到各调用点——预算只在"真正要跑"的
那一刻才应该被消费。
"""
used = self._snapshot.attempt
if used + 1 > self._snapshot.max_attempts:
raise AttemptsExhausted(self._snapshot.task_id, used, self._snapshot.max_attempts)
return self._apply(
TaskStatus.STARTING,
Trigger.DISPATCH,
reason=reason,
attempt=used + 1,
executor=executor,
started_at=self._clock(),
)
def mark_running(self, *, pid: int | None = None) -> Transition:
"""执行器已拉起，工作真正开始。
``started_at`` 在 :meth:`mark_starting` 时即已写入——它标记**本次尝试**的
起点；这里只补执行载体信息（``pid``：进程执行器下就是子进程的进程号）。
这样任何到达 ``RUNNING`` 的路径（包括 ``STARTING → PAUSING → RUNNING`` 的
暂停撤销）都必然带有开始时间。
带**守卫**：没有 ``started_at`` 就不允许进入 ``RUNNING``。转移表无法表达
"必须在本次尝试中已经启动过"——``PAUSED → RUNNING`` 这条边对就地暂停是
合法的，对"排队期间被暂停、从未启动"的任务却毫无意义。守卫负责补上这个
区别，否则会产出 ``RUNNING`` 却没有开始时间的自相矛盾快照。
"""
if self._snapshot.started_at is None:
raise InvalidTransition(
self._snapshot.status,
TaskStatus.RUNNING,
action="guard:尚未开始本次尝试",
)
return self._apply(TaskStatus.RUNNING, Trigger.LAUNCHED, pid=pid)
def mark_succeeded(self, *, outcome: Mapping[str, Any] | None = None) -> Transition:
"""正常完成。"""
return self._apply(TaskStatus.SUCCEEDED, Trigger.COMPLETED, outcome=outcome)
def mark_retry_wait(self, failure: Failure) -> Transition:
"""本次尝试失败但仍有重试余额，进入退避等待。
必须携带 ``failure``：否则界面无法回答"它为什么在等"。
顺带清掉 ``pid``：那一次尝试的载体已经没了，留着一个死进程号只会让人去 `ps`
一个不存在的进程——比没有 pid 更容易误导。
"""
return self._apply(
TaskStatus.RETRY_WAIT, Trigger.RETRY_SCHEDULED, failure=failure, pid=None
)
def mark_retry_exhausted(self, failure: Failure) -> Transition:
"""重试余额耗尽，落 ``FAILED``。"""
return self._apply(TaskStatus.FAILED, Trigger.RETRY_EXHAUSTED, failure=failure)
def mark_failed(self, failure: Failure) -> Transition:
"""遇到不可重试的错误，直接失败。"""
return self._apply(TaskStatus.FAILED, Trigger.FAILED, failure=failure)
def mark_skipped(self, *, reason: str | None = None) -> Transition:
"""条件为假或依策略跳过。"""
return self._apply(TaskStatus.SKIPPED, Trigger.SKIP_REQUESTED, reason=reason)
def mark_cancelled(
self, *, reason: str | None = None, failure: Failure | None = None
) -> Transition:
"""编排层级联取消。"""
return self._apply(
TaskStatus.CANCELLED,
Trigger.CANCEL_REQUESTED,
reason=reason,
failure=failure,
)
def mark_lost(self, *, reason: str | None = None, failure: Failure | None = None) -> Transition:
"""心跳丢失、载体被放弃，结果未知。
应当带上 ``failure``：``LOST`` 的全部意义就是"不知道结果如何"，
若不同时记录为什么不知道，事后根本无从判断该重跑还是该人工核对。
"""
return self._apply(
TaskStatus.LOST,
Trigger.HEARTBEAT_LOST,
reason=reason,
failure=failure,
)
# ------------------------------------------------------------------
# 控制面（由 Controller / 用户调用）
# ------------------------------------------------------------------
def request_stop(self, *, actor: str = "user", reason: str | None = None) -> Transition:
"""请求停止。
任务若正占用执行资源（STARTING/RUNNING/PAUSING），进入 ``STOPPING`` 等待
静默；否则立即落 ``STOPPED``。对尚未真正跑起来的任务显示"停止中"是虚假
信息，会让运维误以为需要等待。
"""
if self._snapshot.status in _STOP_IS_DEFERRED:
return self._apply(
TaskStatus.STOPPING,
Trigger.STOP_REQUESTED,
actor=actor,
reason=reason,
requested_action=RequestedAction.STOP,
)
return self._apply(TaskStatus.STOPPED, Trigger.STOP_REQUESTED, actor=actor, reason=reason)
def mark_stopped(
self,
*,
actor: str = "system",
forced: bool = False,
reason: str | None = None,
failure: Failure | None = None,
) -> Transition:
"""运行期确认任务已停止。``forced`` 表示走的是强杀/放弃等待路径。"""
return self._apply(
TaskStatus.STOPPED,
Trigger.STOP_FORCED if forced else Trigger.STOP_CONFIRMED,
actor=actor,
reason=reason,
failure=failure,
)
def request_pause(self, *, actor: str = "user", reason: str | None = None) -> Transition:
"""请求暂停。
与 :meth:`request_stop` 同一原则：任务是否已占用执行资源，决定要不要经过
``PAUSING`` 缓冲。尚未跑起来的任务（QUEUED / RETRY_WAIT）本就静止，直接落
``PAUSED``——对它们显示"暂停中"是虚假信息。
已在 ``PAUSING`` 的任务不会被"加急"成 ``PAUSED``：它还没到达检查点，
谎称已暂停会让用户以为资源已释放。幂等处理由控制面负责。
"""
status = self._snapshot.status
if status is TaskStatus.PAUSING:
raise InvalidTransition(status, TaskStatus.PAUSED, action=str(Trigger.PAUSE_REQUESTED))
if status in _PAUSE_IS_DEFERRED:
return self._apply(
TaskStatus.PAUSING,
Trigger.PAUSE_REQUESTED,
actor=actor,
reason=reason,
requested_action=RequestedAction.PAUSE,
)
if self.can(TaskStatus.PAUSED):
return self._apply(
TaskStatus.PAUSED, Trigger.PAUSE_REQUESTED, actor=actor, reason=reason
)
raise InvalidTransition(status, TaskStatus.PAUSING, action=str(Trigger.PAUSE_REQUESTED))
def mark_paused(self, *, reason: str | None = None) -> Transition:
"""任务已到达静默点并挂起。
请求已经完成，因此清空 ``requested_action``——它承载的是"已受理但尚未
完成的请求"，留着一个已完成的请求会让 UI 永远显示"暂停中"。
"""
return self._apply(
TaskStatus.PAUSED,
Trigger.PAUSE_REACHED,
reason=reason,
requested_action=None,
)
def abort_pause(self, *, reason: str | None = None) -> Transition:
"""暂停超时仍未静默，回滚为 ``RUNNING``。
如实告知"没能暂停"，而不是假装已暂停——后者会让用户以为资源已释放。
带**守卫**：只允许从 ``PAUSING`` 回滚。``PAUSED → RUNNING`` 这条边是为
*就地恢复* 准备的（见 :meth:`resume`），它要求本次尝试已经启动过；把
"回滚"也放进来会产出一个没有 ``started_at`` 的 ``RUNNING`` 快照。
"""
if self._snapshot.status is not TaskStatus.PAUSING:
raise InvalidTransition(
self._snapshot.status,
TaskStatus.RUNNING,
action=str(Trigger.PAUSE_ABORTED),
)
return self._apply(
TaskStatus.RUNNING,
Trigger.PAUSE_ABORTED,
reason=reason,
requested_action=None,
)
def resume(self, *, actor: str = "user", reason: str | None = None) -> Transition:
"""恢复被暂停的任务。
目标状态取决于**暂停发生在哪一层**（见 :meth:`_resume_target`）：
* 就地暂停——任务阻塞在 ``checkpoint()`` 里，始终占着执行槽，它的调用栈
还在内存中，因此恢复就是回到 ``RUNNING``。我们没有办法把任意的 Python
调用栈换出内存再换回来，所以**暂停不释放并发槽位**。
* 尚未开始就被暂停——从未启动过，恢复应当回到 ``QUEUED`` 重新排队。
"""
target = self._resume_target()
changes: dict[str, Any] = {"requested_action": None}
if target is TaskStatus.QUEUED:
changes["queued_at"] = self._clock()
return self._apply(
target,
Trigger.RESUME_REQUESTED,
actor=actor,
reason=reason,
**changes,
)
def _resume_target(self) -> TaskStatus:
"""判定恢复应该回到哪个状态。
先看迁移历史：暂停若是从 ``PAUSING`` 进入的，就是就地暂停。
历史缺失时（例如快照直接从存储加载、尚未经历任何迁移）退化为按
``attempt`` 判断——已经开过尝试的任务只可能是就地暂停的。
"""
for record in reversed(self._history):
if record.to_status is TaskStatus.PAUSED:
return (
TaskStatus.RUNNING
if record.from_status is TaskStatus.PAUSING
else TaskStatus.QUEUED
)
return TaskStatus.RUNNING if self._snapshot.attempt >= 1 else TaskStatus.QUEUED
# ------------------------------------------------------------------
# 复活与人工干预（绕过转移表，因此自带更严格的守卫与审计）
# ------------------------------------------------------------------
def retry(
self,
*,
actor: str = "user",
reason: str | None = None,
force: bool = False,
) -> Transition:
"""从终态复活，重新进入就绪队列。
不修改历史：``attempt`` 只在 :meth:`mark_starting` 中递增，既有尝试记录
保持不变。重试预算耗尽时抛 :class:`AttemptsExhausted`；``force=True``
会抬高 ``max_attempts`` 以放行——这是"我知道它已经失败够多次了，再试一次"。
"""
current = self._snapshot
if current.status not in REVIVABLE:
raise InvalidTransition(
current.status,
TaskStatus.QUEUED,
action=str(Trigger.RETRY_REQUESTED),
)
changes: dict[str, Any] = {
"failure": None,
"ended_at": None,
"outcome": None,
"requested_action": None,
"overridden": False,
"queued_at": self._clock(),
# 执行载体与进度都属于已经结束的那次尝试，不应带入新尝试：
# pid 是**那一次**的进程号，重试会起一个全新的载体。
# （逐次尝试的完整明细在迁移历史里：starting / retry_wait 自带版本号、
# 时间戳与原因。AttemptRecord 目前没有写入方。）
"started_at": None,
"progress": None,
"pid": None,
}
if current.attempt >= current.max_attempts:
if not force:
raise AttemptsExhausted(current.task_id, current.attempt, current.max_attempts)
changes["max_attempts"] = current.attempt + 1
return self._apply(
TaskStatus.QUEUED,
Trigger.RETRY_REQUESTED,
actor=actor,
reason=reason,
enforce_table=False,
**changes,
)
def reset(self, *, actor: str = "user", reason: str | None = None) -> Transition:
"""重置为 ``PENDING``，丢弃全部尝试历史。
危险操作：重试计数与失败原因都会消失。调用方应要求二次确认。
"""
return self._apply(
TaskStatus.PENDING,
Trigger.RESET_REQUESTED,
actor=actor,
reason=reason,
enforce_table=False,
attempt=0,
failure=None,
outcome=None,
progress=None,
requested_action=None,
queued_at=None,
started_at=None,
ended_at=None,
last_heartbeat_at=None,
executor=None,
pid=None,
overridden=False,
)
def force_transition(
self,
target: TaskStatus,
*,
actor: str,
reason: str,
failure: Failure | None = None,
) -> Transition:
"""绕过转移表人工改写状态。
现实运维需要"这个任务卡住了，我确认后把它标成失败"的能力；堵死它只会让人
绕过系统去改数据库。因此本方法存在，但要求非空 ``actor`` 与 ``reason``，
并把 ``overridden=True`` 写进快照，使人工干预在报表中单独可见。
"""
if not actor or not actor.strip():
raise ValueError("force_transition 必须记录 actor")
if not reason or not reason.strip():
raise ValueError("force_transition 必须记录 reason")
changes: dict[str, Any] = {"overridden": True, "requested_action": None}
if failure is not None:
changes["failure"] = failure
if target in TERMINAL:
# _apply 会补 ended_at。SUCCEEDED 语义上不可能带失败原因，顺手抹平，
# 否则会留下一个恒被 validate() 判为不一致的快照。
if target is TaskStatus.SUCCEEDED:
changes["failure"] = None
else:
# 非终态带着结束时间是自相矛盾的。
changes["ended_at"] = None
if target is TaskStatus.RUNNING and self._snapshot.started_at is None:
changes["started_at"] = self._clock()
if target is TaskStatus.RETRY_WAIT and failure is None:
raise ValueError(
"force_transition 到 RETRY_WAIT 必须给出 failure，否则无法解释为何在等待重试"
)
return self._apply(
target,
Trigger.ADMIN_OVERRIDE,
actor=actor,
reason=reason,
enforce_table=False,
**changes,
)
# ------------------------------------------------------------------
# 高频旁路更新
# ------------------------------------------------------------------
def heartbeat(self, *, at: datetime | None = None) -> TaskSnapshot:
"""刷新心跳时间。
``at`` 允许调用方给出**真实的心跳时刻**（来自执行载体），而不是"现在"。
这个区别很关键：若一律填"现在"，一个卡死的任务也会显得心跳新鲜，``idle``
超时就永远不会触发——而那正是它存在的唯一理由。
刻意**不**递增 ``status_version``：心跳是高频写入，若参与乐观并发，会和
状态写入持续冲突。存储层应将心跳与进度走各自的写入路径。
"""
self._snapshot = replace(self._snapshot, last_heartbeat_at=at or self._clock())
return self._snapshot
def update_progress(self, progress: Progress) -> TaskSnapshot:
"""更新进度。同样不改变状态，也不递增 ``status_version``。"""
self._snapshot = replace(self._snapshot, progress=progress)
return self._snapshot
./src\chongming_task\core\__init__.py
"""运行时内核：状态机、依赖图、调度器。
本层刻意**不**依赖 ``asyncio``、不访问数据库、不认识日志库——领域逻辑必须能被
属性测试穷举验证，这要求它保持纯粹。
"""
./src\chongming_task\runtime\cancellation.py
"""协作式取消与暂停。
Python 既不能安全强杀线程，也没有跨平台的抢占式暂停手段（见设计文档 §11）。
因此停止与暂停都必须是**任务主动配合**的：任务在长耗时循环里调用
:meth:`CancellationToken.checkpoint`，由它响应信号。
这个选择带来一条硬约束，必须让使用者知道：
不调用 ``checkpoint()`` 的任务，既无法优雅停止，也无法暂停。
本模块只管信号，不管状态迁移——``RUNNING -> STOPPING`` 之类的写入仍由
:class:`~chongming_task.core.state_machine.StateMachine` 独占。
"""
from __future__ import annotations
import asyncio
import threading
import time
from ..errors import TaskCancelled
__all__ = ["CancellationToken"]
class CancellationToken:
"""可跨线程使用的取消/暂停信号。
线程安全：同一个令牌会被任务线程（调用 :meth:`checkpoint`）与调度线程
（调用 :meth:`request_stop` / :meth:`request_pause`）同时访问。
"""
__slots__ = ("_ack", "_cond", "_pause", "_stop", "_stop_reason")
def __init__(self) -> None:
self._cond = threading.Condition(threading.Lock())
self._stop = False
self._stop_reason: str | None = None
self._pause = False
self._ack = False
"""是否已有 checkpoint 真正停在了暂停点上。"""
# ------------------------------------------------------------------
# 只读
# ------------------------------------------------------------------
@property
def is_stop_requested(self) -> bool:
"""是否已请求停止。"""
with self._cond:
return self._stop
@property
def is_pause_requested(self) -> bool:
"""是否已请求暂停。"""
with self._cond:
return self._pause
@property
def is_pause_acknowledged(self) -> bool:
"""任务是否已确认暂停（某个 checkpoint 正在暂停点上阻塞）。"""
with self._cond:
return self._ack
# ------------------------------------------------------------------
# 请求侧（调度线程）
# ------------------------------------------------------------------
def request_stop(self, reason: str | None = None) -> None:
"""请求停止。
会唤醒所有正阻塞在暂停点上的任务，使它们抛 :class:`TaskCancelled`——
否则"暂停中的任务停不掉"。重复调用是幂等的，且**保留首个原因**，
因为最先给出的那个通常最接近真实起因。
"""
with self._cond:
if self._stop:
return
self._stop = True
self._stop_reason = reason
self._cond.notify_all()
def request_pause(self) -> None:
"""请求暂停。任务将在下一个 :meth:`checkpoint` 挂起。"""
with self._cond:
if self._pause:
return
self._pause = True
self._ack = False
def resume(self) -> None:
"""撤销暂停，唤醒阻塞中的任务。"""
with self._cond:
if not self._pause:
return
self._pause = False
self._ack = False
self._cond.notify_all()
def wait_paused(self, timeout: float | None = None) -> bool:
"""等待任务确认已暂停，返回是否在超时前确认。
调度器用它区分"已进入 ``PAUSED``"与"暂停没能生效"。后者必须回滚为
``RUNNING`` 并告警，而不是假装成功——假装会让用户以为资源已释放。
"""
deadline = None if timeout is None else time.monotonic() + timeout
with self._cond:
while not self._ack:
if self._stop:
return False
if deadline is None:
self._cond.wait()
continue
remaining = deadline - time.monotonic()
if remaining <= 0:
return False
self._cond.wait(remaining)
return True
# ------------------------------------------------------------------
# 任务侧
# ------------------------------------------------------------------
def _raise_if_stopped_locked(self) -> None:
"""须在持有 ``_cond`` 时调用。
抽成独立方法不只是为了去重：若把它内联进 :meth:`checkpoint`，mypy 会把
第一次检查后的 ``self._stop`` 收窄为 ``Literal[False]``，进而把循环之后
的复查判为不可达语句——而那个复查恰恰是"暂停中被停止"能生效的关键。
"""
if self._stop:
raise TaskCancelled(self._stop_reason)
def checkpoint(self) -> None:
"""协作点：响应停止与暂停。
这是任务与运行时之间的**唯一约定**。长耗时循环必须周期性调用它。
停止优先于暂停：即使任务正处于暂停挂起状态，停止请求也会立刻生效并
抛出 :class:`TaskCancelled`——否则"暂停中的任务停不掉"。
"""
with self._cond:
self._raise_if_stopped_locked()
while self._pause and not self._stop:
# ``_ack`` 必须在**循环体内**设置，不能只在进入时设一次。
#
# 否则会有一次丢唤醒：``resume()`` 把 ``_pause`` 置回 False 并通知后，
# 任务线程还没来得及重新取到锁，新的 ``request_pause()`` 就把
# ``_pause`` 又置成了 True —— 线程重新检查循环条件时发现"还要暂停"，
# 于是继续等待，而它已经越过了设置 ``_ack`` 的那一行。结果
# ``wait_paused()`` 永远等不到确认、暂停超时回滚。
#
# 表现出来就是"暂停 → 恢复 → 再暂停"里第二次暂停永久失效。
self._ack = True
self._cond.notify_all()
self._cond.wait()
self._raise_if_stopped_locked()
async def acheckpoint(self) -> None:
""":meth:`checkpoint` 的异步版本。
非阻塞快速路径（绝大多数调用）保持同步执行，避免在紧密循环里为每次
检查都付一次线程调度成本；只有真的需要挂起时才把阻塞等待交给线程。
"""
with self._cond:
self._raise_if_stopped_locked()
if not self._pause:
return
await asyncio.to_thread(self.checkpoint)
def raise_if_stopped(self) -> None:
"""只检查停止信号，不阻塞、不响应暂停。
用于无法安全挂起的位置（例如 ``except`` 块的清理逻辑中），此时阻塞会
导致清理永远无法完成。
"""
with self._cond:
if self._stop:
raise TaskCancelled(self._stop_reason)
./src\chongming_task\runtime\context.py
"""执行期上下文：任务与运行时之间的**唯一协作入口**。
任务作者看到的就是这个对象（公开 API 里叫 ``TaskContext``，同一个类）：
```python
def extract(ctx, source: str) -> list[Row]:
ctx.log(f"开始读取 {source}")
rows = fetch(source)
for index, row in enumerate(rows):
ctx.checkpoint() # 响应停止/暂停 + 刷新心跳
ctx.progress(index + 1, total=len(rows))
return rows
三件事收敛在这里，而且都必须在这里：
- 停止 / 暂停 / 心跳——与执行载体强相关，只有执行器知道载体是什么；
- 进度——高频写入，走旁路而不是状态机（见
touch）； - 上游数据——管道节点之间传值，只有调度器知道谁是谁的上游。
为什么上下文必须是注入的
见 :class:~chongming_task.runtime.executors.base.ExecutionRequest 的说明：上下文里
有 :class:~chongming_task.runtime.cancellation.CancellationToken，它每次尝试才
诞生（由执行器在启动那一刻创建），所以任务根本无从捕获——它只能从第一个位置参数
拿到。这不是风格问题：换成闭包捕获就得让令牌活到任务定义之外，指针一旦指向"上一次
尝试"的令牌，停止指令就会发给一个已经结束的执行。
"""
from future import annotations
import logging from collections import deque from collections.abc import Callable, Mapping from dataclasses import dataclass, field from datetime import UTC, datetime from pathlib import Path from typing import Any
from ..models import Progress from .cancellation import CancellationToken from .heartbeat import Heartbeat
all = ["RuntimeContext", "TaskContext"]
#: 上下文在内存里保留的日志条数上限。 _LOG_BUFFER = 200
#: 任务日志的 logger 名称前缀。 LOGGER_PREFIX = "chongming_task.task"
@dataclass(slots=True) class RuntimeContext: """任务执行期可用的运行时能力。"""
task_id: str
name: str
attempt: int
token: CancellationToken
heartbeat: Heartbeat
artifact_dir: Path | None = None
"""本次尝试的产物目录。按 attempt 隔离，重试不会污染上一次的产物。"""
inputs: Mapping[str, Any] = field(default_factory=dict)
"""上游节点的返回值，按节点 id 索引。非管道任务为空。
只有**成功**的上游才会出现在这里：失败或跳过的上游没有值可传，与其塞一个
``None`` 让人误以为"上游返回了 None"，不如干脆不出现。用 :meth:`input` 取则
可以给出默认值。
"""
on_progress: Callable[[Progress], None] | None = None
"""进度回写通道，由执行器注入。为 ``None`` 时 :meth:`progress` 静默丢弃。"""
clock: Callable[[], datetime] | None = None
"""时钟。与任务本身共用同一个，避免时间戳来自两个不同的时间源。"""
logs: deque[str] = field(default_factory=lambda: deque(maxlen=_LOG_BUFFER))
"""本任务最近的日志（内存中，最多 200 条）。
它**不是**日志存储方案：进程退出即消失。它的用途是让"任务失败时最后几条日志是
什么"在事后可查——把完整日志落盘/入库是应用自己的事（挂个 ``logging`` handler
即可）。
"""
# ------------------------------------------------------------------
# 协作点
# ------------------------------------------------------------------
def checkpoint(self) -> None:
"""协作点：响应停止与暂停，并刷新心跳。
把心跳绑在这里是有意的：``checkpoint()`` 就是"我还在推进"的信号。
不调用它的任务既不响应停止，也会被 ``idle`` 超时判定为疑似卡死。
收到停止请求时抛 :class:`~chongming_task.errors.TaskCancelled`。
"""
self.token.checkpoint()
self.heartbeat.beat()
async def acheckpoint(self) -> None:
""":meth:`checkpoint` 的异步版本，供原生 ``async`` 任务使用。"""
await self.token.acheckpoint()
self.heartbeat.beat()
@property
def is_stop_requested(self) -> bool:
"""是否已请求停止。"""
return self.token.is_stop_requested
@property
def is_pause_requested(self) -> bool:
"""是否已请求暂停。"""
return self.token.is_pause_requested
# ------------------------------------------------------------------
# 进度
# ------------------------------------------------------------------
def progress(
self,
done: int,
*,
total: int | None = None,
unit: str | None = None,
message: str | None = None,
) -> None:
"""上报进度。
**不参与状态机**，也不递增 ``status_version``：进度可能每秒几十次，若参与
乐观并发会与状态写入持续冲突，把版本号变成噪声。
它同时也刷新心跳吗？**不。** 心跳是"我还在推进"的整体信号，进度是"推进到哪
了"。两者都调用一次是常见写法，但把心跳藏进进度里会让一个只报进度、不做实际
工作的任务看起来永远健康——那正是 ``idle`` 超时要抓的东西。
"""
sink = self.on_progress
if sink is None:
return
# 这两个字段会被写进存储。类型错误必须在**调用点**炸掉：让它流到写穿路径上，
# 失败只会被记录下来、任务继续跑，表现是"看板的进度永远是旧的"——离原因很远。
if unit is not None and not isinstance(unit, str):
raise TypeError(f"unit 必须是 str 或 None，收到 {type(unit).__name__}")
if message is not None and not isinstance(message, str):
raise TypeError(f"message 必须是 str 或 None，收到 {type(message).__name__}")
moment = self.clock() if self.clock is not None else datetime.now(UTC)
sink(
Progress(
done=done,
total=total,
unit=unit,
message=message,
updated_at=moment,
)
)
# ------------------------------------------------------------------
# 日志
# ------------------------------------------------------------------
@property
def logger(self) -> logging.Logger:
"""本任务的 logger。
名字是 ``chongming_task.task.<task_id>``，因此可以在应用侧按任务维度单独配
handler / 级别，而不必给每个任务各写一套。
"""
return logging.getLogger(f"{LOGGER_PREFIX}.{self.task_id}")
def log(
self,
message: str,
*,
level: int = logging.INFO,
exc_info: bool = False,
) -> None:
"""记一条日志。
走标准库 ``logging``，并把 ``task_id`` / ``task_name`` / ``attempt`` 作为
``extra`` 附上，因此 handler 可以据此过滤或加字段。同时保留最近
:data:`_LOG_BUFFER` 条在 :attr:`logs` 里，供事后查看。
"""
self.logs.append(message)
self.logger.log(
level,
message,
exc_info=exc_info,
extra={
"task_id": self.task_id,
"task_name": self.name,
"attempt": self.attempt,
},
)
def warning(self, message: str) -> None:
""":meth:`log` 的 ``WARNING`` 快捷方式。"""
self.log(message, level=logging.WARNING)
def error(self, message: str) -> None:
""":meth:`log` 的 ``ERROR`` 快捷方式。"""
self.log(message, level=logging.ERROR)
# ------------------------------------------------------------------
# 上游数据
# ------------------------------------------------------------------
@property
def upstream(self) -> Mapping[str, Any]:
"""上游节点的返回值，按节点 id 索引。只读。"""
return self.inputs
def input(self, node_id: str, default: Any = None) -> Any:
"""取某个上游节点的返回值；没有就给 ``default``。"""
return self.inputs.get(node_id, default)
#: 公开 API 里的名字。同一个类，不是包装——任务作者只面对一个 ctx。
TaskContext = RuntimeContext
`./src\chongming_task\runtime\heartbeat.py`
```python
"""心跳。
心跳由 ``checkpoint()`` 驱动，**不额外起线程**。理由：一个不调用 checkpoint 的
任务本来就无法被暂停或优雅停止，再给它伪造心跳只会掩盖问题。
代价必须明确：``idle`` 超时衡量的是 **checkpoint 间隔**，而不是"CPU 是否在干活"。
一段长时间不 checkpoint 的重计算会被判定为疑似卡死。这是我们希望暴露的信号，
但也要求任务作者让 checkpoint 频率与 ``idle`` 超时相匹配。
"""
from __future__ import annotations
from datetime import datetime
from ..clock import Clock, utcnow
__all__ = ["Heartbeat"]
class Heartbeat:
"""最近一次活动时间。
线程安全性与令牌一致：由持有该任务执行权的单个线程写入。
"""
__slots__ = ("_clock", "_last")
def __init__(self, clock: Clock | None = None) -> None:
self._clock: Clock = clock if clock is not None else utcnow
self._last = self._clock()
def beat(self) -> None:
"""记录一次活动。"""
self._last = self._clock()
@property
def last_beat(self) -> datetime:
"""最近一次活动的时刻。"""
return self._last
def idle_seconds(self) -> float:
"""距最近一次活动已过去的秒数。"""
return (self._clock() - self._last).total_seconds()
def is_stale(self, idle_timeout: float) -> bool:
"""静默时长是否已达到 ``idle_timeout``。"""
if idle_timeout <= 0:
raise ValueError(f"idle_timeout 必须为正数，收到 {idle_timeout}")
return self.idle_seconds() >= idle_timeout
./src\chongming_task\runtime\timeouts.py
"""超时判定。
三个作用域相互独立，取**最早**到期者：
* ``start`` —— 从提交到真正开始执行，一旦开始执行即失效；
* ``per_attempt`` —— 单次尝试的总时长；
* ``idle`` —— 距最近一次心跳的静默时长（语义见 :mod:`~chongming_task.runtime.heartbeat`）。
本模块只做**纯计算**，不做等待、不设计时循环——那属于调度器。这样判定逻辑可以
用 ``FakeClock`` 精确测试，而调度器只负责在到期时唤醒。
"""
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime, timedelta
from ..models import FailureKind, TimeoutScope
from ..policies import Timeout
__all__ = ["Deadline", "is_expired", "next_deadline", "remaining_seconds"]
@dataclass(frozen=True, slots=True)
class Deadline:
"""一个已确定的到期点。"""
scope: TimeoutScope
at: datetime
"""到期时刻。"""
@property
def kind(self) -> FailureKind:
"""由作用域得到对应的失败归类，避免调用处手写映射而拼错。"""
return FailureKind.for_timeout(self.scope)
def next_deadline(
timeout: Timeout,
*,
submitted_at: datetime | None = None,
started_at: datetime | None = None,
last_beat_at: datetime | None = None,
) -> Deadline | None:
"""返回下一个到期点；没有任何生效的超时则返回 ``None``。
传 ``None`` 表示对应阶段尚未开始，该作用域自动失效——例如 ``started_at``
一旦给出，``start`` 超时就不再参与竞争。
"""
if not timeout.is_enabled:
# 一档超时都没配：这条路径在监督循环里每轮都会走，早退掉省下建表与排序。
return None
candidates: list[Deadline] = []
if timeout.start is not None and submitted_at is not None and started_at is None:
candidates.append(
Deadline(TimeoutScope.START, submitted_at + timedelta(seconds=timeout.start))
)
if timeout.per_attempt is not None and started_at is not None:
candidates.append(
Deadline(TimeoutScope.ATTEMPT, started_at + timedelta(seconds=timeout.per_attempt))
)
# idle 必须要求已开始执行：任务还没跑起来就谈不上"卡住"，
# 否则排队等待的时长会被误判成静默超时。
if timeout.idle is not None and last_beat_at is not None and started_at is not None:
candidates.append(
Deadline(TimeoutScope.IDLE, last_beat_at + timedelta(seconds=timeout.idle))
)
if not candidates:
return None
return min(candidates, key=lambda deadline: deadline.at)
def is_expired(deadline: Deadline, *, now: datetime) -> bool:
"""``now`` 是否已到达或越过 ``deadline``。"""
return now >= deadline.at
def remaining_seconds(deadline: Deadline, *, now: datetime) -> float:
"""距到期还剩多少秒；已过期则为负值。
调度器用它决定下一次唤醒的等待时长——**不要把负数直接交给 sleep**。
"""
return (deadline.at - now).total_seconds()
./src\chongming_task\runtime\__init__.py
"""运行时：取消令牌、心跳、超时判定、执行器。
本层可以依赖 :mod:`chongming_task.core`，但**不反向被依赖**：领域内核不认识
线程、超时与执行器，因此它才能被属性测试穷举验证。
"""
./src\chongming_task\runtime\executors\base.py
"""执行器抽象。
执行器只负责四件事：把一段可调用对象跑起来、报告它是否还在跑、能被要求停止、
能被要求暂停。
它**不**做两件事，这是刻意划下的边界：
* 不决定状态迁移——那是 :class:`~chongming_task.core.state_machine.StateMachine`
的独占职责；
* 不决定并发上限——那是调度器的职责（``ConcurrencyPolicy``）。因此执行器内部
不再叠一层线程池，否则两层限制互相掩盖，很难推理。
"""
from __future__ import annotations
import abc
import asyncio
import inspect
import traceback
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Protocol, runtime_checkable
from ...clock import Clock
from ...errors import ConfigurationError, TaskCancelled, TaskTimeout
from ...models import Failure, FailureKind, Progress
from ..cancellation import CancellationToken
from ..context import RuntimeContext
from ..heartbeat import Heartbeat
__all__ = [
"ExecutionHandle",
"ExecutionOutcome",
"ExecutionRequest",
"Executor",
"LocalExecutionHandle",
"classify_exception",
"run_target",
]
@dataclass(frozen=True, slots=True)
class ExecutionRequest:
"""一次尝试的执行请求。
``target`` 的**第一个位置参数是执行期上下文**，由执行器注入::
def extract(ctx, source: str) -> list[Row]: ...
这个约定是刻意的：上下文（尤其是取消令牌）属于**某一次尝试**，在启动那一刻才
由执行器造出来，因此任务只能被喂给它，不可能自己攥着一个。若靠闭包捕获，任务
定义里就得提前握住某个令牌，而那多半是上一次尝试留下的——停止指令会发给一个
已经结束的执行，且完全静默。
"""
task_id: str
name: str
target: Callable[..., Any]
attempt: int = 1
args: tuple[Any, ...] = ()
kwargs: Mapping[str, Any] = field(default_factory=dict, hash=False)
artifact_dir: Path | None = None
inputs: Mapping[str, Any] = field(default_factory=dict, hash=False)
"""上游节点的返回值，按节点 id 索引。管道数据传递靠它。"""
on_progress: Callable[[Progress], None] | None = None
"""进度回写通道：进度不参与状态机，由调用方决定怎么落地。"""
clock: Clock | None = None
"""时钟。与自己的任务共用同一个，避免时间戳来自两个不同的时间源。"""
@dataclass(frozen=True, slots=True)
class ExecutionOutcome:
"""一次执行的结果。"""
value: Any = None
failure: Failure | None = None
abandoned: bool = False
"""执行载体未能停止、已被放弃。
为 ``True`` 时结果**未知**：任务可能仍在后台运行并产生副作用。调用方应据此
落 ``LOST`` 而不是 ``STOPPED``——把"没停住"记成"已停止"不只是不准确，
而是会让运维误判资源已释放。
"""
@property
def ok(self) -> bool:
"""是否正常完成。"""
return self.failure is None and not self.abandoned
def run_target(request: ExecutionRequest, context: RuntimeContext) -> Any:
"""调用任务体并取出返回值；目标是协程函数时，在本线程内把它跑完。
线程执行器与内联执行器共用这一个入口，语义才不会分叉。**必须真的等到协程结束**：
直接把协程对象当返回值的话，任务会"看起来成功"（结果是 ``<coroutine object …>``），
然后在写存储时以"返回值不可序列化"失败，或者在没有存储时静默通过——两种都难查。
协程在**本线程**新起一个事件循环来跑，不借用调用方已有的循环：因此不要求调用方
是异步的，也不会与调用方的循环互相干扰，``ctx.acheckpoint()`` 这类异步协作点
照常工作。调用方若本身已在事件循环里（只有内联执行器可能遇到），``asyncio.run``
会抛 ``RuntimeError``——如实报错，不猜。
"""
value = request.target(context, *request.args, **request.kwargs)
if not inspect.isawaitable(value):
return value
return asyncio.run(_drive_awaitable(value))
async def _drive_awaitable(value: Awaitable[Any]) -> Any:
"""把任意可等待对象包成协程——``asyncio.run`` 只接受协程。"""
return await value
def classify_exception(exc: BaseException) -> Failure:
"""把任务抛出的异常转成结构化失败信息。
``TaskCancelled`` 继承自 ``BaseException``，必须先于普通异常分支处理。
它被标记为**不可重试**：取消不是失败，重试它等于无视停止指令。
"""
detail = "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
if isinstance(exc, TaskCancelled):
return Failure(
kind=FailureKind.CANCELLED,
message=str(exc) or "任务已被取消",
traceback=detail,
retryable=False,
)
if isinstance(exc, TaskTimeout):
return Failure(
kind=FailureKind.for_timeout(exc.scope),
message=str(exc),
traceback=detail,
retryable=True,
)
if isinstance(exc, ConfigurationError):
# 配置类错误（目标不可 pickle、任务没注册……）重试多少次都是同一个结果，
# 再试只会白烧重试预算，还会把"改配置"这件事拖到重试耗尽之后才暴露。
return Failure(
kind=FailureKind.EXCEPTION,
message=str(exc),
traceback=detail,
retryable=False,
)
return Failure(
kind=FailureKind.EXCEPTION,
message=f"{type(exc).__name__}: {exc}",
traceback=detail,
retryable=True,
)
@runtime_checkable
class ExecutionHandle(Protocol):
"""一次执行的句柄。
实时控制（停止、暂停、恢复）都长在句柄上：只有句柄同时持有令牌、心跳与
执行载体。这样控制面就不必自己拼装这三者，也不会出现"令牌与载体对不上号"。
"""
@property
def token(self) -> CancellationToken:
"""本次执行的取消/暂停令牌。"""
...
@property
def heartbeat(self) -> Heartbeat:
"""本次执行的心跳。"""
...
@property
def pid(self) -> int | None:
"""执行载体的操作系统进程号；**同进程**执行器为 ``None``。
进程执行器下它是子进程的 pid，会写进任务快照——事后排障时"这个任务当时在哪个
进程里"是第一个要问的问题（看板、`ps`、`top` 都靠它对上号）。同进程执行器里
任务就跑在监督进程自己身上，再存一个 pid 只是重复。
"""
...
@property
def running(self) -> bool:
"""是否仍在运行。
被放弃之后依然返回 ``True``——它确实还在跑。
"""
...
def wait(self, timeout: float | None = None) -> ExecutionOutcome | None:
"""等待结束；``timeout`` 内未结束则返回 ``None``。"""
...
def stop(self, *, graceful_timeout: float, reason: str | None = None) -> ExecutionOutcome:
"""请求停止，最多等待 ``graceful_timeout`` 秒给出结论。
Python 无法安全地终止线程，所以超时之后**没有兜底动作**：只能放弃等待，并
如实报告 :attr:`ExecutionOutcome.abandoned`（结果未知，由上层落 ``LOST``）。
谎报成"已停止"会让运维以为资源已经释放。
"""
...
def pause(self, *, timeout: float) -> bool:
"""请求暂停并等待任务确认；返回是否在 ``timeout`` 内确认。
返回 ``False`` 表示暂停**没有生效**（任务未到达检查点）。调用方应据此
回滚为 ``RUNNING`` 并告警，而不是假装已暂停。
"""
...
def resume(self) -> None:
"""撤销暂停。"""
...
@runtime_checkable
class Executor(Protocol):
"""执行器。"""
name: str
"""标识，用于快照的 ``executor`` 字段与排障。"""
def launch(self, request: ExecutionRequest) -> ExecutionHandle:
"""启动一次执行并立即返回句柄（**不得**阻塞到任务结束）。"""
...
class LocalExecutionHandle(abc.ABC):
"""同进程执行器的句柄基类（线程与内联执行器实现它）。
集中实现三件与具体执行器无关、但语义必须完全一致的事：令牌与心跳的持有、
``pause`` / ``resume`` 的握手、``stop`` 的"请求 → 优雅窗口 → 兜底"三段式。
子类只需回答两个问题：还在跑吗（:attr:`running`），以及优雅窗口用尽之后
怎么办（:meth:`_on_stop_timeout`）。
"""
__slots__ = ("_forced_outcome", "_heartbeat", "_request", "_token")
def __init__(self, request: ExecutionRequest, *, clock: Clock | None = None) -> None:
self._request = request
self._token = self._make_token()
self._heartbeat = Heartbeat(clock)
self._forced_outcome: ExecutionOutcome | None = None
def _make_token(self) -> CancellationToken:
"""造本次执行的取消令牌。
子类可换掉实现——进程执行器用它把本地信号**转发**给子进程：控制面有几条路径
直接调 ``handle.token.request_stop()``（见 :meth:`~chongming_task.run.TaskRun.stop`），
绕过了句柄的 ``stop()``；令牌不转发的话，子进程永远收不到那次停止，
表现出来就是"停止按钮没反应"。
"""
return CancellationToken()
@property
def request(self) -> ExecutionRequest:
"""原始请求。"""
return self._request
@property
def token(self) -> CancellationToken:
"""本次执行的取消/暂停令牌。"""
return self._token
@property
def pid(self) -> int | None:
"""同进程执行器没有独立进程号——任务就跑在本进程里。"""
return None
@property
def heartbeat(self) -> Heartbeat:
"""本次执行的心跳。"""
return self._heartbeat
def build_context(self) -> RuntimeContext:
"""构造交给任务的执行期上下文。"""
return RuntimeContext(
task_id=self._request.task_id,
name=self._request.name,
attempt=self._request.attempt,
token=self._token,
heartbeat=self._heartbeat,
artifact_dir=self._request.artifact_dir,
inputs=self._request.inputs,
on_progress=self._request.on_progress,
clock=self._request.clock,
)
def pause(self, *, timeout: float) -> bool:
"""请求暂停并等待任务确认。"""
self._token.request_pause()
return self._token.wait_paused(timeout)
def resume(self) -> None:
"""撤销暂停。"""
self._token.resume()
def stop(self, *, graceful_timeout: float, reason: str | None = None) -> ExecutionOutcome:
"""请求停止，三段式收场。"""
self._token.request_stop(reason)
outcome = self.wait(graceful_timeout)
if outcome is not None:
return outcome
self._forced_outcome = self._on_stop_timeout(graceful_timeout)
return self._forced_outcome
def wait(self, timeout: float | None = None) -> ExecutionOutcome | None:
"""等待结束。
兜底结论一旦产生就是**粘性**的：后台载体日后返回的结果不会改写已记录的
``LOST``，否则审计记录会自相矛盾。
"""
if self._forced_outcome is not None:
return self._forced_outcome
return self._wait_impl(timeout)
@property
@abc.abstractmethod
def running(self) -> bool:
"""是否仍在运行。被放弃的任务依然返回 ``True``。"""
@abc.abstractmethod
def _wait_impl(self, timeout: float | None) -> ExecutionOutcome | None:
"""实际的等待实现；超时返回 ``None``。"""
@abc.abstractmethod
def _on_stop_timeout(self, graceful_timeout: float) -> ExecutionOutcome:
"""优雅窗口用尽后的兜底。
线程不可强杀，所以实现只能是"放弃等待"，并如实报告
:attr:`ExecutionOutcome.abandoned`。
"""
./src\chongming_task\runtime\executors\inline.py
"""内联执行器：在调用线程内直接跑完。
零调度开销，适合微秒级的纯计算步骤，也便于确定性测试。
**它刻意违反 ``Executor`` 的"launch 不得阻塞"约定**：``launch`` 会一直执行到
任务结束才返回。这就是它存在的意义，但也意味着**异步调度器不得使用它**——
否则单个任务会卡住整个调度循环。它的用途是同步运行路径与测试。
停止与暂停对它意义有限：整个执行发生在调用线程内，没有第二个线程可以投递
信号。任务仍可通过 ``ctx.is_stop_requested`` 自行检查，但外部无法在它运行期间
介入——因为"运行期间"调用方根本没拿到控制权。
``async def`` 任务同样能跑（在本线程新起一个事件循环跑完），但**调用方若已经在事件
循环里**就不能用它——``asyncio.run`` 不允许嵌套。异步应用请用线程执行器。
"""
from __future__ import annotations
from typing import Any
from ...clock import Clock
from .base import (
ExecutionOutcome,
ExecutionRequest,
LocalExecutionHandle,
classify_exception,
run_target,
)
__all__ = ["InlineExecutionHandle", "InlineExecutor"]
class InlineExecutor:
"""在调用线程内同步执行任务。"""
name = "inline"
def __init__(self, *, clock: Clock | None = None) -> None:
self._clock = clock
def launch(self, request: ExecutionRequest) -> InlineExecutionHandle:
"""同步执行到结束，返回一个已完成的句柄。
注意这与 ``Executor`` 协议的约定不同（见模块文档），是有意为之。
"""
handle = InlineExecutionHandle(request, clock=self._clock)
handle.run()
return handle
class InlineExecutionHandle(LocalExecutionHandle):
"""内联执行的句柄。它在 ``launch`` 返回时就已经结束。"""
__slots__ = ("_error", "_value")
def __init__(self, request: ExecutionRequest, *, clock: Clock | None = None) -> None:
super().__init__(request, clock=clock)
self._value: Any = None
self._error: BaseException | None = None
def run(self) -> None:
"""执行任务。由 :class:`InlineExecutor` 调用，不应直接使用。"""
context = self.build_context()
try:
self._value = run_target(self._request, context)
except BaseException as exc:
self._error = exc
@property
def running(self) -> bool:
"""内联执行在 ``launch`` 返回时必然已结束。"""
return False
def pause(self, *, timeout: float) -> bool:
"""内联执行无法被暂停：没有第二个线程能投递信号。
返回 ``False`` 是**如实回答**——调用方据此回滚为 ``RUNNING``，
而不是以为暂停成功了。
"""
del timeout
return False
def _wait_impl(self, timeout: float | None) -> ExecutionOutcome | None:
del timeout # 已完成，无需等待
if self._error is None:
return ExecutionOutcome(value=self._value)
return ExecutionOutcome(failure=classify_exception(self._error))
def _on_stop_timeout(self, graceful_timeout: float) -> ExecutionOutcome:
"""不会被调用：内联执行在 ``stop`` 之前就已结束。"""
return ExecutionOutcome(
failure=classify_exception(
RuntimeError(f"内联执行无法停止（已等待 {graceful_timeout} 秒）")
)
)
./src\chongming_task\runtime\executors\process.py
"""进程执行器：每个任务一个守护**子进程**。
它换来三件线程给不了的事：
* **硬性停止保证**：优雅窗口用尽后 ``terminate()`` → ``kill()``，任务是真的停了，
不是"放弃等待"（线程只能放弃，并如实记 ``LOST``）；
* **绕开 GIL**：CPU 密集的任务能真正并行；
* **崩溃隔离**：任务里加载的原生扩展段错误只会带走子进程，监督进程活着，
能如实记下一次失败并按策略重试。
代价也是三条，都是**结构性的**，不是实现偷懒：
1. **target / 参数 / 上游数据必须可 pickle**。``lambda``、闭包、局部函数统统不行
——它们在子进程里找不到。违反时在 ``launch()`` 当场抛
:class:`~chongming_task.errors.ConfigurationError`，并告诉调用方改用什么，
而不是等子进程里神秘失败。
2. **子进程是守护进程，不能再生子进程**（``multiprocessing`` 的硬规定，为的是让
解释器退出时不被卡住）。任务内部要并行，请自己开线程，或改用线程执行器。
3. **停止与暂停有轮询延迟**（当前 20ms）：信号得靠 ``multiprocessing`` 的事件过界，
子进程由一个轻量轮询线程把它映射到本地令牌上。线程执行器是即时的。
刻意选择 **spawn** 而不是 POSIX 上的 ``fork``：监督进程是**多线程**的，fork 出来的
子进程会继承一堆状态不一致的锁，那是死锁的经典来源。spawn 起的是全新解释器，代价是
启动慢一些（毫秒到百毫秒级，还要重新 import 目标模块）。
于是有两条使用要求，与 ``multiprocessing`` 本身一模一样：目标函数必须是**模块级**
的，且调用方的脚本要有 ``if __name__ == "__main__":`` 守卫——否则子进程会把主模块的
顶层代码再跑一遍。
## 父进程怎么知道子进程在干什么
一条**单向管道**（子 → 父）+ 两个 ``Event``（父 → 子）：
| 方向 | 载体 | 内容 |
|---|---|---|
| 子 → 父 | 管道 | ``started`` / ``beat`` / ``paused`` / ``progress`` / ``result`` / ``failure`` |
| 父 → 子 | ``Event`` | 停止、暂停（子进程的轮询线程映射成本地令牌信号） |
**父进程必须持续排空管道**，而且与有没有人在 ``wait()`` 无关：心跳就在那条管道上，
只在 ``wait()`` 里读的话，``idle`` 超时会在任务健康时也触发。所以句柄起一个读取线程，
它同时负责收集进度与最终结论。
管道在子进程死亡时给出 ``EOFError``——父进程靠它把"子进程没了"和"子进程给了结论"
区分开，这也正是"崩溃隔离"能如实上报的原因。
"""
from __future__ import annotations
import contextlib
import logging
import os
import pickle
import threading
import time
from dataclasses import dataclass, replace
from multiprocessing import get_context
from multiprocessing.connection import Connection
from multiprocessing.process import BaseProcess
from typing import Any
from ...clock import Clock, utcnow
from ...errors import ConfigurationError
from ...models import Failure, FailureKind, Progress
from ..cancellation import CancellationToken
from ..context import RuntimeContext
from ..heartbeat import Heartbeat
from .base import (
ExecutionOutcome,
ExecutionRequest,
LocalExecutionHandle,
classify_exception,
run_target,
)
__all__ = ["ProcessExecutionHandle", "ProcessExecutor"]
#: spawn 上下文。用 spawn 而不是 fork 的理由见模块文档。
_SPAWN = get_context("spawn")
#: 读取线程等子进程真正退出的时间（避免留下僵尸进程）。
_JOIN_TIMEOUT = 2.0
#: ``terminate()`` 之后留给子进程的时间；仍活着就升级为 ``kill()``。
#:
#: POSIX 上 ``terminate()`` 发 ``SIGTERM``（本执行器不装信号处理器，因此几乎立即生效，
#: 这点时间只是收尸余量）；Windows 上它就是 ``TerminateProcess``，毫无缓冲余地。
_TERMINATE_GRACE = 0.2
#: 子进程里镜像父进程信号的轮询周期。这就是"停止/暂停有 20ms 延迟"的来源：
#: 用一点延迟换掉"信号本身不可跨进程"这个硬约束。
_SIGNAL_POLL = 0.02
#: 心跳节流间隔：紧凑循环里的 ``checkpoint()`` 不该每次都写一次管道。
_BEAT_INTERVAL = 0.05
#: 读取线程的轮询周期。它同时决定"发现子进程没了"的延迟——管道不给 EOF 的那些情形
#: （见 :meth:`ProcessExecutionHandle._read_loop` 的说明）全靠它兜住。
_READ_POLL = 0.05
_log = logging.getLogger("chongming_task.executor.process")
class ProcessExecutor:
"""用子进程执行任务。"""
name = "process"
def __init__(self, *, clock: Clock | None = None) -> None:
"""``clock`` 只影响父进程这侧的心跳时间戳；子进程用系统时钟。"""
self._clock = clock
def launch(self, request: ExecutionRequest) -> ProcessExecutionHandle:
"""校验载荷可 pickle，然后启动一个守护子进程。"""
_ensure_picklable(request)
handle = ProcessExecutionHandle(request, clock=self._clock)
handle.start()
return handle
@dataclass(frozen=True, slots=True)
class _ChildJob:
"""发到子进程的载荷：只包含可 pickle 的部分。
``on_progress`` 与 ``clock`` 刻意置空——它们是**闭包**，跨不过进程边界。子进程
用自己的时钟，并把进度经管道发回来（父进程收到后再回调真正的 ``on_progress``）。
"""
request: ExecutionRequest
class _Channel:
"""子进程这侧的管道：把消息发给父进程。
管道不是线程安全的：任务的 ``ctx.checkpoint()``（心跳）与轮询线程（暂停期心跳）
会并发写同一根管道，不加锁内容就会交错损坏。锁只保护 ``send``，不保护任何业务
逻辑；发不出去（父进程已经走了）**绝不让任务失败**——可观测性坏了不是任务的错。
"""
__slots__ = ("_conn", "_last_beat", "_lock")
def __init__(self, conn: Connection) -> None:
self._conn = conn
self._lock = threading.Lock()
self._last_beat = 0.0
def send(self, kind: str, payload: Any = None) -> None:
with self._lock:
self._conn.send((kind, payload))
def beat(self) -> None:
"""回报一次心跳，按 :data:`_BEAT_INTERVAL` 节流。"""
now = time.monotonic()
if now - self._last_beat < _BEAT_INTERVAL:
return
self._last_beat = now
with contextlib.suppress(Exception):
# 父进程已经走了：心跳没人收，也绝不能让任务因此失败。
self.send("beat")
def progress(self, progress: Progress) -> None:
with contextlib.suppress(Exception):
# 同上：进度丢一条不是任务的问题。
self.send("progress", progress)
def close(self) -> None:
with contextlib.suppress(Exception):
self._conn.close()
class _MirrorToken(CancellationToken):
"""子进程这侧的令牌：把父进程的跨进程信号映射成本地信号。
``CancellationToken`` 内部是 ``threading.Condition``，**不可跨进程传递**，所以
子进程自建一个，由 :class:`_SignalMirror` 轮询父进程的 ``Event``，再调用
``request_stop`` / ``request_pause`` / ``resume`` 把状态映射进来。
映射（而不是直接解析事件）的好处是：暂停、停止、以及"暂停中的任务也要停得掉"
这些语义一条都不用重写——跨进程与同进程的行为由**同一个** ``checkpoint()`` 决定。
"""
__slots__ = ("_channel",)
def __init__(self, channel: _Channel) -> None:
super().__init__()
self._channel = channel
def checkpoint(self) -> None:
"""本地协作点 + 一次心跳回传。
心跳挂在这里是有意的：``checkpoint()`` 就是"我还在推进"的唯一信号，``idle``
超时靠它在跨进程时依然有效——而它是唯一能发现"卡住但没报错"的机制。
"""
super().checkpoint()
self._channel.beat()
class _SignalMirror(threading.Thread):
"""子进程里的轮询线程：把父进程的 ``Event`` 映射到本地令牌，并在暂停期代为心跳。
两个职责都必须在**任务线程之外**：
* 镜像信号：任务线程可能正阻塞在暂停点上，只有别的线程能叫醒它；
* 暂停期心跳：被暂停的任务不是卡死的，可它自己不会去调 ``checkpoint()``——心跳
得由这里代发，否则父进程的 ``idle`` 超时会把一个正常暂停的任务判成卡死。
"""
def __init__(
self,
token: _MirrorToken,
channel: _Channel,
stop_event: Any,
pause_event: Any,
) -> None:
super().__init__(name="task-signal-mirror", daemon=True)
self._token = token
self._channel = channel
self._stop_event = stop_event
self._pause_event = pause_event
self._done = threading.Event()
self._paused_reported = False
def shutdown(self) -> None:
self._done.set()
def run(self) -> None:
while not self._done.wait(_SIGNAL_POLL):
try:
self._mirror_once()
except Exception: # pragma: no cover - 父进程没了不该影响任务
return
def _mirror_once(self) -> None:
if self._stop_event.is_set():
self._token.request_stop()
# 停止之后不再需要心跳：任务要么在收尾，要么马上要被强杀。
return
if self._pause_event.is_set():
self._token.request_pause()
if self._token.is_pause_acknowledged:
if not self._paused_reported:
# 报一次确认——父进程等的就是这个：它据此把状态落成 PAUSED，
# 而不是以为暂停没生效（那会让两边对同一个任务的判断分叉）。
self._paused_reported = True
with contextlib.suppress(Exception):
self._channel.send("paused")
# 已经停在暂停点上：代为心跳，别让父进程以为它卡死了。
self._channel.beat()
return
self._paused_reported = False
self._token.resume()
class ProcessExecutionHandle(LocalExecutionHandle):
"""子进程执行的句柄。
它兑现了 :class:`~chongming_task.runtime.executors.base.ExecutionHandle` 里那句
"超时之后如何收场由具体执行器决定"：线程只能放弃并记 ``abandoned``，而这里
**真的把任务停掉**——所以结论是一个已知的 ``CANCELLED``（任务未完成），
而不是"结果未知"。
"""
def __init__(self, request: ExecutionRequest, *, clock: Clock | None = None) -> None:
# 这两个事件必须先于 super().__init__() 建好：基类在构造时就会通过
# _make_token() 取令牌，而转发令牌正需要它们。
self._stop_event = _SPAWN.Event()
self._pause_event = _SPAWN.Event()
self._pause_acked = threading.Event()
self._settled = threading.Event()
self._lock = threading.Lock()
self._verdict: ExecutionOutcome | None = None
self._child_started = False
self._terminating = threading.Event()
self._process: BaseProcess | None = None
self._reader: threading.Thread | None = None
self._conn: Connection | None = None
super().__init__(request, clock=clock)
# ------------------------------------------------------------------
# 控制面
# ------------------------------------------------------------------
def _make_token(self) -> CancellationToken:
"""父进程这侧的令牌要**转发**信号：有几条控制路径直接调 ``handle.token``。"""
return _RelayToken(self._stop_event, self._pause_event)
@property
def pid(self) -> int | None:
"""子进程的进程号；尚未启动时为 ``None``。"""
process = self._process
return None if process is None else process.pid
@property
def process(self) -> BaseProcess | None:
"""底层 ``multiprocessing.Process``，供诊断与测试使用。"""
return self._process
@property
def running(self) -> bool:
"""是否仍在运行。
判据是**载体**：子进程还活着，且我们还没拿到结论。被强杀之后立刻为 ``False``
——它确实已经死了；谎报成 ``True`` 会让上层以为资源还占着。
"""
process = self._process
if process is None:
return False
return process.is_alive() and not self._settled.is_set()
def pause(self, *, timeout: float) -> bool:
"""请求暂停，等子进程在下一个 ``checkpoint()`` 确认。
超时**必须回滚**：撤掉暂停请求，让子进程继续跑。不回滚的话，调用方已经按
"暂停没生效"把状态回滚成 ``RUNNING``，子进程却会在稍后的某个检查点自己停住
——那就成了一个没人知道、也没人会恢复的挂起。
"""
self._pause_acked.clear()
self._token.request_pause()
if self._pause_acked.wait(timeout):
return True
self._token.resume()
return False
def resume(self) -> None:
"""撤销暂停。幂等。"""
self._token.resume()
self._pause_acked.clear()
def _on_stop_timeout(self, graceful_timeout: float) -> ExecutionOutcome:
"""优雅窗口用尽：**真的把子进程杀掉**。
结论是 ``CANCELLED`` 而不是 ``LOST``：任务确实没有完成，而且我们知道这一点
（线程那边做不到，只能放弃并寄希望于它自己收工）。
"""
self._terminating.set()
self._terminate_and_reap()
return self._terminated_outcome(graceful_timeout)
def _terminated_outcome(self, graceful_timeout: float) -> ExecutionOutcome:
return ExecutionOutcome(
failure=Failure(
kind=FailureKind.CANCELLED,
message=(
f"任务在 {graceful_timeout} 秒内未响应停止请求，已强制终止子进程"
f"（pid={self.pid}）"
),
retryable=False,
)
)
# ------------------------------------------------------------------
# 生命周期
# ------------------------------------------------------------------
def start(self) -> None:
"""起子进程与读取线程。"""
# 单向管道：返回 (recv, send)。子进程拿 send，父进程留 recv。
# Pipe() 实际返回 PipeConnection，它实现了 Connection 协议的所有方法；
# typeshed 没有把它声明为 Connection 的子类型，Pylance 会误报赋值不兼容。
self._conn, child_conn = _SPAWN.Pipe(duplex=False) # type: ignore[assignment]
request = self._request
# 闭包过不去：进度回调与时钟都置空，子进程用自己的。
job = _ChildJob(request=replace(request, on_progress=None, clock=None))
self._process = _SPAWN.Process(
target=_child_entry,
args=(job, child_conn, self._stop_event, self._pause_event),
name=f"task-{request.task_id}",
daemon=True,
)
self._process.start()
# 父进程这侧的写端必须关掉：否则子进程死后管道不会给出 EOF，
# "载体没了"这件事就永远发现不了。
child_conn.close()
self._reader = threading.Thread(
target=self._read_loop, name=f"task-reader-{request.task_id}", daemon=True
)
self._reader.start()
def _wait_impl(self, timeout: float | None) -> ExecutionOutcome | None:
if not self._settled.wait(timeout):
return None
with self._lock:
return self._verdict
def _read_loop(self) -> None:
"""持续排空管道——**与有没有人在 wait 无关**。
只在 ``wait()`` 里读的话，心跳就只在调用方主动轮询时推进，而 ``idle`` 超时靠
的正是心跳：一个健康的长期任务会被判成卡死。所以这里是一个独立的线程。
**判活以 ``Process.is_alive()`` 为准，不靠管道断开。** 原因实测过：子进程若在
反序列化参数时就死掉（目标模块 import 失败、参数过不来），父进程那侧的写端
句柄**从未交出去**，管道永远不给 EOF——只等 EOF 的话这里会一直阻塞，子进程
早已死了而父进程毫无察觉。管道只负责搬消息。
"""
conn = self._conn
process = self._process
assert conn is not None and process is not None # start() 先建好它们
try:
while True:
if conn.poll(_READ_POLL):
try:
kind, payload = conn.recv()
except (EOFError, OSError):
# 管道断了：子进程正常退出时会走到这里——那时结论已在手上。
break
if not self._dispatch(kind, payload):
break
continue
if not process.is_alive():
# 载体没了。先把管道里剩下的读完：它可能还来得及把结论发出来。
self._drain(conn)
break
finally:
conn.close()
self._reap()
if not self._settled.is_set():
# 没给出结论就没了：崩溃、被外部杀掉、或 os._exit。
self._settle(self._on_child_vanished())
def _drain(self, conn: Connection) -> None:
"""子进程已经死了，把管道里残留的消息读完（没有残留就直接返回）。"""
while conn.poll(0):
try:
kind, payload = conn.recv()
except (EOFError, OSError):
return
if not self._dispatch(kind, payload):
return
def _dispatch(self, kind: str, payload: Any) -> bool:
"""处理一条消息；返回 ``False`` 表示"这是最后一条，读完收工"。"""
if kind == "beat":
# 记录的是**父进程收到**的时刻：子进程一旦停摆（哪怕它没崩），这里就不再
# 更新，idle 超时照常触发。
self._heartbeat.beat()
return True
if kind == "started":
self._child_started = True
return True
if kind == "paused":
self._pause_acked.set()
return True
if kind == "progress":
self._forward_progress(payload)
return True
if kind == "result":
self._settle(ExecutionOutcome(value=payload))
return False
if kind == "failure":
self._settle(ExecutionOutcome(failure=payload))
return False
_log.warning("进程执行器收到未知消息 %r，已忽略", kind)
return True
def _forward_progress(self, progress: Progress) -> None:
"""把子进程的进度交给父进程的 ``on_progress``。
进度是**可观测性**，它的失败绝不能让任务失败——与"存储故障不让任务失败"
是同一条原则。
"""
callback = self._request.on_progress
if callback is None:
return
try:
callback(progress)
except Exception as exc: # pragma: no cover - 调用方自己的回调
_log.warning("进度回调失败：%s: %s", type(exc).__name__, exc)
def _settle(self, outcome: ExecutionOutcome) -> None:
with self._lock:
if self._verdict is None:
self._verdict = outcome
self._settled.set()
def _reap(self) -> None:
"""等子进程真正退出，避免留下僵尸进程。"""
process = self._process
if process is not None:
process.join(_JOIN_TIMEOUT)
def _on_child_vanished(self) -> ExecutionOutcome:
"""子进程没给结论就没了。
"有没有报告过 started"决定了失败的性质，也决定了运维下一步该做什么：
* **没报告过** —— 子进程根本没跑起来（多半是载荷在子进程里无法反序列化，
或目标模块 import 失败）。任务没有执行过，重试是干净的。
* **报告过** —— 任务跑着跑着载体没了（段错误、``os._exit``、被外部杀掉）。
**结果未知**：副作用可能已经发生了一部分。
"""
if self._terminating.is_set():
# 是我们自己按停止请求杀的，结论以那次停止为准。
return self._terminated_outcome(_TERMINATE_GRACE)
code = self._process.exitcode if self._process is not None else None
if code is None or code == 0:
return ExecutionOutcome(
failure=Failure(
kind=FailureKind.LOST,
message=(
f"子进程（pid={self.pid}）在给出结论前退出（退出码 {code}）；"
"任务可能已在别处产生副作用，结果未知"
),
retryable=True,
)
)
if self._child_started:
message = (
f"子进程（pid={self.pid}）以退出码 {code} 结束且未给出结论："
"任务可能在崩溃前已产生部分副作用，结果未知"
)
else:
message = (
f"子进程（pid={self.pid}）以退出码 {code} 结束，任务从未开始："
"通常是 target 或参数在子进程里无法导入/反序列化"
"（子进程的报错已打印到本进程的 stderr）"
)
return ExecutionOutcome(
failure=Failure(kind=FailureKind.NONZERO_EXIT, message=message, exit_code=code)
)
def _terminate_and_reap(self) -> None:
"""``terminate()`` → 宽限 → ``kill()``。"""
process = self._process
if process is None: # pragma: no cover - 只会在未启动时发生
return
process.terminate()
process.join(_TERMINATE_GRACE)
if process.is_alive():
process.kill()
process.join(_TERMINATE_GRACE)
if self._reader is not None:
# 读取线程会因 EOF 自行退出；这里只是等它收工，好让结论落定。
self._reader.join(_JOIN_TIMEOUT)
class _RelayToken(CancellationToken):
"""父进程这侧的令牌：本地置位的同时把信号转发给子进程。
两条控制路径会**直接**调 ``handle.token.request_stop()``（见
:meth:`~chongming_task.run.TaskRun.stop`），绕过了句柄的 ``stop()``。令牌不转发的话
子进程永远收不到那次停止——表现出来就是"停止按钮没反应"。
"""
__slots__ = ("_pause_event", "_stop_event")
def __init__(self, stop_event: Any, pause_event: Any) -> None:
super().__init__()
self._stop_event = stop_event
self._pause_event = pause_event
def request_stop(self, reason: str | None = None) -> None:
super().request_stop(reason)
# 暂停中的任务也必须停得掉：先解除暂停，否则子进程会一直挂在暂停点上。
self._pause_event.clear()
self._stop_event.set()
def request_pause(self) -> None:
super().request_pause()
self._pause_event.set()
def resume(self) -> None:
super().resume()
self._pause_event.clear()
# ----------------------------------------------------------------------
# 子进程侧
# ----------------------------------------------------------------------
def _child_entry(
job: _ChildJob,
conn: Connection,
stop_event: Any,
pause_event: Any,
) -> None:
"""子进程入口。**必须是模块级函数**：spawn 靠"模块 + 限定名"定位它。
它自己拼装运行期上下文（父进程的令牌、心跳、闭包都过不来），跑完把结果经管道
发回去，然后正常退出——父进程那侧读到的 EOF 就是"载体结束了"。
"""
_ensure_child_logging()
request = job.request
channel = _Channel(conn)
token = _MirrorToken(channel)
mirror = _SignalMirror(token, channel, stop_event, pause_event)
mirror.start()
try:
context = RuntimeContext(
task_id=request.task_id,
name=request.name,
attempt=request.attempt,
token=token,
heartbeat=Heartbeat(),
artifact_dir=request.artifact_dir,
inputs=request.inputs,
# 进度经管道回传，由父进程回调真正的 on_progress。
on_progress=channel.progress,
clock=utcnow,
)
channel.send("started", os.getpid())
try:
value = run_target(request, context)
except BaseException as exc:
# TaskCancelled 是控制流信号、不继承 Exception，必须一起接住：让它逃出
# 子进程的话，父进程只会看到"子进程消失了"，结论就变成未知。
channel.send("failure", classify_exception(exc))
return
try:
channel.send("result", value)
except Exception as exc:
channel.send(
"failure",
Failure(
kind=FailureKind.RESULT_INVALID,
message=(
f"返回值无法跨进程传回（{type(exc).__name__}: {exc}）。"
"进程执行器要求返回值可 pickle，比线程执行器的"
"'可 JSON 序列化'更严——请返回可 pickle 的值，"
"或改用线程执行器。"
),
retryable=False,
),
)
finally:
mirror.shutdown()
channel.close()
def _ensure_child_logging() -> None:
"""子进程不继承父进程的 logging 配置（spawn 起的是全新解释器）。
不补一个 handler 的话，``ctx.log`` 这类 INFO 级日志会被 ``lastResort`` 直接丢掉
——任务在子进程里"什么日志都没有"，是最难查的一种现场。父进程的 stderr 由子进程
继承，所以挂到 stderr 上就是"日志还看得见"。已经有配置就不动它。
"""
root = logging.getLogger()
if root.handlers:
return
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter("%(levelname)s %(name)s: %(message)s"))
root.addHandler(handler)
root.setLevel(logging.INFO)
def _ensure_picklable(request: ExecutionRequest) -> None:
"""在起子进程**之前**验证载荷可 pickle。
为什么要提前查：真让它失败在 ``Process.start()`` 的 spawn 机制深处，报出来的是一句
不知从何下手的 ``PicklingError``（甚至只是子进程里的一行 traceback）。在这里拦下
就能直接告诉调用方"改成什么"。
"""
try:
pickle.dumps(
(
request.target,
request.args,
dict(request.kwargs),
request.artifact_dir,
dict(request.inputs),
)
)
except Exception as exc:
raise ConfigurationError(
f"进程执行器的 target / 参数 / 上游数据必须可 pickle："
f"{type(exc).__name__}: {exc}\n"
"lambda、闭包、局部函数都不可 pickle——这类任务请改用线程执行器"
"（ThreadExecutor）。进程执行器只在需要**强制停止、绕开 GIL、"
"崩溃隔离**时才值得付出这个代价。"
) from exc
./src\chongming_task\runtime\executors\thread.py
"""线程执行器：每个任务一个守护线程。
**这是本项目唯一的执行载体**（另有 :class:`~chongming_task.runtime.executors.inline.InlineExecutor`
用于测试同步化）。选它的理由：
* 线程创建开销在几十微秒量级，进程是毫秒到几十毫秒；
* 线程共享内存，可以直接复用连接池、缓存、已加载的模型；
* 跨进程的方案需要 pickle 参数与返回值，闭包、``lambda``、局部函数、装饰器产物
统统用不了——而 ``@task`` 装饰过的函数恰恰都是这类；参数还必须跨进程可传输。
代价只有一条：**无法强制停止**。Python 不能安全地终止线程，所以优雅窗口超时后
只能放弃等待，并如实报告 ``abandoned=True``（上层落 ``LOST``，不谎报 ``STOPPED``）。
于是约束落回任务本身，两条都要照做：
* 长耗时循环里**必须**调 ``ctx.checkpoint()``——否则停不了，也就会被判成卡死；
* 任务要**幂等或可断点续传**：超时与放弃之后它可能仍在后台跑，重跑是调用方的选择，
后果要能承受。
CPU 密集的活请在三处之一处理：把并行放进任务内部（例如任务里自建进程池），把它
拆成多个任务，或者干脆单独部署到一个节点上——**不要**指望线程能绕开 GIL。这里
不再提供进程执行器：一个"能强杀但要 pickle、还不能用闭包"的执行器，与这个项目的
任务形态（装饰器 + 闭包 + 共享连接池）长期相冲，维护两套语义的代价高于收益。
"""
from __future__ import annotations
import threading
from typing import Any
from ...clock import Clock
from ...models import Failure, FailureKind
from .base import (
ExecutionOutcome,
ExecutionRequest,
LocalExecutionHandle,
classify_exception,
run_target,
)
__all__ = ["ThreadExecutionHandle", "ThreadExecutor"]
class ThreadExecutor:
"""用线程执行任务。"""
name = "thread"
def __init__(self, *, clock: Clock | None = None) -> None:
self._clock = clock
def launch(self, request: ExecutionRequest) -> ThreadExecutionHandle:
"""启动守护线程并立即返回。
刻意**不用** ``concurrent.futures.ThreadPoolExecutor``，两个理由：
1. 它的工作线程会在解释器退出时被 ``join``。一个不响应取消的任务会让
整个进程**无法退出**——对一个专门用来监督"可能不配合"的任务的工具
而言，这是致命缺陷。
2. 它的线程无法放弃：池子会一直持有那个卡住的线程，无法回收。
并发上限由调度器统一控制（``ConcurrencyPolicy.max_parallel``），执行器
不再叠一层池。
"""
handle = ThreadExecutionHandle(request, clock=self._clock)
handle.start()
return handle
class ThreadExecutionHandle(LocalExecutionHandle):
"""线程执行的句柄。
线程安全：工作线程只写 ``_value`` / ``_error`` 并在最后 ``set`` 事件；
其余状态由调度线程读写，二者通过 ``_done`` 建立 happens-before 关系。
"""
__slots__ = ("_done", "_error", "_thread", "_value")
def __init__(self, request: ExecutionRequest, *, clock: Clock | None = None) -> None:
super().__init__(request, clock=clock)
self._done = threading.Event()
self._value: Any = None
self._error: BaseException | None = None
self._thread = threading.Thread(
target=self._body,
name=f"task-{request.task_id}",
daemon=True,
)
@property
def thread(self) -> threading.Thread:
"""底层线程，供诊断与测试使用。"""
return self._thread
def start(self) -> None:
"""启动线程。"""
self._thread.start()
def _body(self) -> None:
context = self.build_context()
try:
self._value = run_target(self._request, context)
except BaseException as exc:
# 必须捕获 BaseException：TaskCancelled 是控制流信号、不继承 Exception。
# 若任由它逃出工作线程，解释器只会打印一行日志然后丢弃它，调度器
# 将永远等不到结论。
self._error = exc
finally:
self._done.set()
@property
def running(self) -> bool:
"""是否仍在运行。被放弃的任务依然返回 ``True``。"""
return not self._done.is_set()
def _wait_impl(self, timeout: float | None) -> ExecutionOutcome | None:
if not self._done.wait(timeout):
return None
return self._collect()
def _on_stop_timeout(self, graceful_timeout: float) -> ExecutionOutcome:
"""线程无法强杀，因此这里只能放弃——架构级限制，不是实现偷懒。"""
return ExecutionOutcome(
abandoned=True,
failure=Failure(
kind=FailureKind.LOST,
message=(
f"任务在 {graceful_timeout} 秒内未响应停止请求；"
"线程无法强制终止，已放弃等待"
"（任务可能仍在后台运行并产生副作用）"
),
retryable=True,
),
)
def _collect(self) -> ExecutionOutcome:
if self._error is None:
return ExecutionOutcome(value=self._value)
return ExecutionOutcome(failure=classify_exception(self._error))
./src\chongming_task\runtime\executors\__init__.py
"""执行器实现。"""
./src\chongming_task\stores\codec.py
"""把领域对象与"扁平的原语映射"互相转换。**公开契约的一部分。**
任何要落地持久化的后端都躲不开这件事：:class:`~chongming_task.models.TaskSnapshot`
里有 ``datetime``、``frozenset``、枚举和嵌套的 ``Progress`` / ``Failure``，而数据库行、
JSON 文档、Redis Hash 只认原语。把这些细节收敛在这里，后端作者就不必各自重写一遍，
也不会各自重写错一遍。
映射的形状是**扁平的一层字典**（列名 → 原语），因为这是最容易落到任何后端的形式：
* 关系库 → 直接就是一行；
* 文档库 / JSON 文件 → 存成一个子对象，或按需展平；
* KV（Redis Hash）→ 键值对直接可写。
往返保真是**契约要求**，不是"尽量"：``row_to_snapshot(snapshot_to_row(s)) == s``
必须成立。``StoreContract`` 里有一条用例专门用"每个字段都填满"的快照验证它——
只测几个字段会漏掉时区、``frozenset``、嵌套结构这类静默丢失。
"""
from __future__ import annotations
import json
from collections.abc import Mapping
from datetime import UTC, datetime
from typing import Any
from ..core.state_machine import Transition
from ..models import Failure, FailureKind, Progress, RequestedAction, TaskSnapshot
from ..states import TaskStatus
from .protocol import Command, CommandKind
__all__ = [
"command_to_row",
"decode_dt",
"decode_failure",
"decode_progress",
"decode_tags",
"encode_dt",
"encode_failure",
"encode_progress",
"encode_tags",
"row_to_command",
"row_to_snapshot",
"row_to_transition",
"snapshot_to_row",
"transition_to_row",
]
# ----------------------------------------------------------------------
# 原语
# ----------------------------------------------------------------------
def encode_dt(value: datetime | None) -> str | None:
"""存成 UTC 的 ISO 字符串。
一律先转 UTC：跨时区部署时把本地时间原样写进去，排序与比较都会错，而且不报错。
"""
return value.astimezone(UTC).isoformat() if value is not None else None
def decode_dt(raw: str | None) -> datetime | None:
"""读回时间。``None`` 原样返回。"""
return datetime.fromisoformat(raw) if raw else None
def encode_tags(tags: frozenset[str]) -> str:
"""标签存成 JSON 数组。
排序后再写：不排序的话，内容相同的两个集合会产生不同的字节串，让"判断是否变化"
这类比较失效。
"""
return json.dumps(sorted(tags), ensure_ascii=False)
def decode_tags(raw: str) -> frozenset[str]:
"""读回标签集合。"""
return frozenset(json.loads(raw))
def _dump(value: Any) -> str | None:
return json.dumps(value, ensure_ascii=False, sort_keys=True) if value is not None else None
def _load(raw: str | None) -> Any:
return json.loads(raw) if raw else None
# ----------------------------------------------------------------------
# 嵌套结构
# ----------------------------------------------------------------------
def encode_progress(progress: Progress | None) -> str | None:
"""进度 → JSON 文本。"""
if progress is None:
return None
return _dump(
{
"done": progress.done,
"total": progress.total,
"unit": progress.unit,
"message": progress.message,
"updated_at": encode_dt(progress.updated_at),
}
)
def decode_progress(raw: str | None) -> Progress | None:
"""JSON 文本 → 进度。"""
data = _load(raw)
if data is None:
return None
return Progress(
done=data["done"],
total=data["total"],
unit=data["unit"],
message=data["message"],
updated_at=decode_dt(data["updated_at"]),
)
def encode_failure(failure: Failure | None) -> str | None:
"""失败信息 → JSON 文本。
``FailureKind`` 是 ``StrEnum``，必须显式取 ``.value``：直接把枚举丢给
``json.dumps`` 在部分版本上会写成一个不可读的形式。
"""
if failure is None:
return None
return _dump(
{
"kind": failure.kind.value,
"message": failure.message,
"traceback": failure.traceback,
"exit_code": failure.exit_code,
"retryable": failure.retryable,
}
)
def decode_failure(raw: str | None) -> Failure | None:
"""JSON 文本 → 失败信息。"""
data = _load(raw)
if data is None:
return None
return Failure(
kind=FailureKind(data["kind"]),
message=data["message"],
traceback=data["traceback"],
exit_code=data["exit_code"],
retryable=data["retryable"],
)
# ----------------------------------------------------------------------
# 快照
# ----------------------------------------------------------------------
def snapshot_to_row(snapshot: TaskSnapshot) -> dict[str, Any]:
"""快照 → 扁平原语映射。
键的顺序与 :data:`chongming_task.stores.sqlite._TASK_COLUMNS` 一致，关系库可以
直接拿它拼 SQL。
"""
return {
"task_id": snapshot.task_id,
"name": snapshot.name,
"pipeline_run_id": snapshot.pipeline_run_id,
"tags": encode_tags(snapshot.tags),
"bucket": snapshot.bucket,
"bucket_key": snapshot.bucket_key,
"status": snapshot.status.value,
"status_version": snapshot.status_version,
"requested_action": (
snapshot.requested_action.value if snapshot.requested_action else None
),
"attempt": snapshot.attempt,
"max_attempts": snapshot.max_attempts,
"created_at": encode_dt(snapshot.created_at),
"queued_at": encode_dt(snapshot.queued_at),
"started_at": encode_dt(snapshot.started_at),
"ended_at": encode_dt(snapshot.ended_at),
"last_heartbeat_at": encode_dt(snapshot.last_heartbeat_at),
"progress": encode_progress(snapshot.progress),
"failure": encode_failure(snapshot.failure),
"outcome": _dump(dict(snapshot.outcome)) if snapshot.outcome is not None else None,
"executor": snapshot.executor,
"pid": snapshot.pid,
"overridden": int(snapshot.overridden),
"lease_owner": snapshot.lease_owner,
"lease_until": encode_dt(snapshot.lease_until),
}
def row_to_snapshot(row: Mapping[str, Any]) -> TaskSnapshot:
"""扁平原语映射 → 快照。
``overridden`` 从整数读回布尔：SQLite 没有布尔类型，写进去的是 0/1。
"""
requested = row["requested_action"]
return TaskSnapshot(
task_id=row["task_id"],
name=row["name"],
pipeline_run_id=row["pipeline_run_id"],
tags=decode_tags(row["tags"]),
bucket=row["bucket"],
bucket_key=row["bucket_key"],
status=TaskStatus(row["status"]),
status_version=row["status_version"],
requested_action=RequestedAction(requested) if requested else None,
attempt=row["attempt"],
max_attempts=row["max_attempts"],
created_at=decode_dt(row["created_at"]),
queued_at=decode_dt(row["queued_at"]),
started_at=decode_dt(row["started_at"]),
ended_at=decode_dt(row["ended_at"]),
last_heartbeat_at=decode_dt(row["last_heartbeat_at"]),
progress=decode_progress(row["progress"]),
failure=decode_failure(row["failure"]),
outcome=_load(row["outcome"]),
executor=row["executor"],
pid=row["pid"],
overridden=bool(row["overridden"]),
lease_owner=row["lease_owner"],
lease_until=decode_dt(row["lease_until"]),
)
# ----------------------------------------------------------------------
# 迁移记录（审计轨迹，只追加）
# ----------------------------------------------------------------------
def transition_to_row(transition: Transition) -> dict[str, Any]:
"""迁移记录 → 扁平原语映射。"""
return {
"task_id": transition.task_id,
"version": transition.version,
"from_status": transition.from_status.value,
"to_status": transition.to_status.value,
"at": encode_dt(transition.at),
"trigger": transition.trigger,
"actor": transition.actor,
"reason": transition.reason,
}
def row_to_transition(row: Mapping[str, Any]) -> Transition:
"""扁平原语映射 → 迁移记录。
``at`` 缺失时回退到"现在"：宁可时间戳不准，也不要因为一条历史记录读不出来而让
整个任务查询失败。
"""
return Transition(
task_id=row["task_id"],
from_status=TaskStatus(row["from_status"]),
to_status=TaskStatus(row["to_status"]),
at=decode_dt(row["at"]) or datetime.now(UTC),
trigger=row["trigger"],
actor=row["actor"],
reason=row["reason"],
version=row["version"],
)
# ----------------------------------------------------------------------
# 控制指令
# ----------------------------------------------------------------------
def command_to_row(item: Command) -> dict[str, Any]:
"""控制指令 → 扁平原语映射。"""
return {
"command_id": item.command_id,
"task_id": item.task_id,
"kind": item.kind.value,
"created_at": encode_dt(item.created_at),
"requested_by": item.requested_by,
"reason": item.reason,
"lease_until": encode_dt(item.lease_until),
"finished_at": encode_dt(item.finished_at),
"outcome": item.outcome,
}
def row_to_command(row: Mapping[str, Any]) -> Command:
"""扁平原语映射 → 控制指令。"""
return Command(
command_id=row["command_id"],
task_id=row["task_id"],
kind=CommandKind(row["kind"]),
created_at=decode_dt(row["created_at"]) or datetime.min.replace(tzinfo=UTC),
requested_by=row["requested_by"],
reason=row["reason"],
lease_until=decode_dt(row["lease_until"]),
finished_at=decode_dt(row["finished_at"]),
outcome=row["outcome"],
)
./src\chongming_task\stores\memory.py
"""内存存储实现。
适合：单进程嵌入（不关心重启）、单元测试、以及作为其他实现的**参照物**——
新写一个后端时，先用它对照 :class:`~chongming_task.stores.protocol.TaskStore` 的
语义，比读文档可靠。
不适合：跨进程查询（内存不共享）。那正是数据库实现存在的理由。
"""
from __future__ import annotations
import threading
from dataclasses import replace
from datetime import datetime, timedelta
from ..core.state_machine import Transition
from ..errors import ConcurrencyConflict, StoreError
from ..models import Progress, TaskSnapshot
from ..states import TERMINAL, TaskStatus
from .protocol import Command, TaskQuery
__all__ = ["InMemoryTaskStore"]
class InMemoryTaskStore:
"""把任务、迁移历史与控制指令放在进程内存里。"""
def __init__(self) -> None:
self._lock = threading.RLock()
self._tasks: dict[str, TaskSnapshot] = {}
self._history: dict[str, dict[int, Transition]] = {}
self._commands: dict[str, Command] = {}
# ------------------------------------------------------------------
# 1. 任务快照
# ------------------------------------------------------------------
def create(self, snapshot: TaskSnapshot) -> None:
"""写入新任务；``task_id`` 已存在则抛错。"""
with self._lock:
if snapshot.task_id in self._tasks:
raise StoreError(f"任务已存在：{snapshot.task_id}")
self._tasks[snapshot.task_id] = snapshot
self._history.setdefault(snapshot.task_id, {})
def save(self, snapshot: TaskSnapshot, *, expect_version: int) -> None:
"""按乐观并发更新快照。
**保留租约列**：它们由 :meth:`touch` 独占。直接整体替换会让一次状态迁移把租约
清空，于是重启恢复把一个活得好好的任务误判为失去监督者——契约测试里有一条
用例专门盯这件事。
"""
with self._lock:
current = self._tasks.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)
self._tasks[snapshot.task_id] = replace(
snapshot,
lease_owner=current.lease_owner,
lease_until=current.lease_until,
)
def get(self, task_id: str) -> TaskSnapshot | None:
"""读取单个任务。"""
with self._lock:
return self._tasks.get(task_id)
def query(self, query: TaskQuery | None = None) -> list[TaskSnapshot]:
"""按条件查询。"""
condition = query if query is not None else TaskQuery()
with self._lock:
found = list(self._tasks.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._lock:
result: dict[TaskStatus, int] = {}
for snapshot in self._tasks.values():
result[snapshot.status] = result.get(snapshot.status, 0) + 1
return result
# ------------------------------------------------------------------
# 2. 审计轨迹
# ------------------------------------------------------------------
def append(self, transition: Transition) -> None:
"""追加一条迁移记录；同一版本重复写入是幂等的。"""
with self._lock:
entries = self._history.setdefault(transition.task_id, {})
entries.setdefault(transition.version, transition)
def history(self, task_id: str, *, since_version: int = 0) -> list[Transition]:
"""读取迁移历史。"""
with self._lock:
entries = self._history.get(task_id, {})
return [entry for version, entry in sorted(entries.items()) if version > since_version]
# ------------------------------------------------------------------
# 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._lock:
current = self._tasks.get(task_id)
if current is None:
return
changes: dict[str, object] = {}
if progress is not None:
changes["progress"] = progress
if heartbeat_at is not None:
changes["last_heartbeat_at"] = heartbeat_at
if lease_owner is not None:
changes["lease_owner"] = lease_owner
if lease_until is not None:
changes["lease_until"] = lease_until
if changes:
self._tasks[task_id] = replace(current, **changes) # type: ignore[arg-type]
# ------------------------------------------------------------------
# 4. 控制指令
# ------------------------------------------------------------------
def submit_command(self, command: Command) -> None:
"""写入一条控制指令。"""
with self._lock:
if command.command_id in self._commands:
raise StoreError(f"指令已存在：{command.command_id}")
self._commands[command.command_id] = command
def claim_commands(
self, *, now: datetime, lease_seconds: float, limit: int = 50
) -> list[Command]:
"""取走可执行指令并设置租约。"""
with self._lock:
claimable = [
command
for command in self._commands.values()
if command.is_open and (command.lease_until is None or command.lease_until <= now)
]
claimable.sort(key=lambda command: command.created_at)
claimed: list[Command] = []
for command in claimable[:limit]:
leased = replace(command, lease_until=now + timedelta(seconds=lease_seconds))
self._commands[command.command_id] = leased
claimed.append(leased)
return claimed
def complete_command(
self, command_id: str, *, now: datetime, outcome: str | None = None
) -> None:
"""标记指令已完成。"""
with self._lock:
command = self._commands.get(command_id)
if command is None:
return
self._commands[command_id] = replace(
command, finished_at=now, outcome=outcome, lease_until=None
)
def release_command(self, command_id: str, *, now: datetime) -> None:
"""交还一条已取走但无人认领的指令。"""
del now
with self._lock:
command = self._commands.get(command_id)
if command is None or not command.is_open:
return
self._commands[command_id] = replace(command, lease_until=None)
def commands(self) -> list[Command]:
"""全部指令，供诊断与测试使用（不属于 :class:`TaskStore` 契约）。"""
with self._lock:
return sorted(self._commands.values(), key=lambda item: item.created_at)
# ------------------------------------------------------------------
# 运维
# ------------------------------------------------------------------
def prune(self, *, keep_terminal: int) -> int:
"""丢弃最旧的终态任务及其历史。"""
with self._lock:
terminal = [
snapshot for snapshot in self._tasks.values() if snapshot.status in TERMINAL
]
terminal.sort(key=lambda item: item.created_at or datetime.min)
excess = len(terminal) - max(0, keep_terminal)
if excess <= 0:
return 0
for snapshot in terminal[:excess]:
del self._tasks[snapshot.task_id]
self._history.pop(snapshot.task_id, None)
return excess
def close(self) -> None:
"""空操作：内存实现没有需要释放的资源。"""
./src\chongming_task\stores\postgres.py
"""Postgres 存储实现（可选依赖：``psycopg`` 3 + ``psycopg-pool``）。
为什么是可选依赖
----------------
核心包保持**零第三方运行时依赖**：``import chongming_task`` 与内存 / SQLite 两个
内置实现都不需要驱动。本模块只在显式导入时才要求 psycopg::
pip install "chongming-task[postgres]"
from chongming_task import PostgresTaskStore
store = PostgresTaskStore("postgresql://user@host:5432/tasks")
controller = Controller(store=store, owner="worker-1")
它是**跨机部署的第 0 步**（SQLite 文件不能跨机共享，见分布式设计说明 §15）：
把持久层换成网络后端，上层（Controller、CLI、看板）一行都不用改。
与 SQLite 实现的三个刻意差异
----------------------------
1. **连接池**。SQLite 用"一把锁 + 单连接"串行化；Postgres 用
``psycopg_pool.ConnectionPool``，多个线程真的并行。``close()`` 的语义见方法文档。
2. **认领指令用 ``FOR UPDATE SKIP LOCKED``**：一条语句完成"挑选 + 打租约"，
并发认领互不阻塞、也不会重复取到同一条。SQLite 用 ``BEGIN IMMEDIATE`` 是同一
保证的不同实现。
3. **原生类型**：时间列 ``TIMESTAMPTZ``（无时区的 ``timestamp`` 跨机必然出错），
``tags`` / ``progress`` / ``failure`` / ``outcome`` 用 ``JSONB``，标签过滤下推为
``tags @> '["io"]'::jsonb`` 并吃到 GIN 索引。
三条硬约束（与 SQLite 实现相同，任何后端都不能破）
--------------------------------------------------
* ``save`` 的"校验 + 写入"是一条 ``UPDATE ... WHERE status_version = ?``，
"校验 + 写入"不是两步；
* ``save`` **不写租约列**（``lease_owner`` / ``lease_until`` 由 :meth:`touch`
独占）。否则一次状态迁移就会把租约覆盖回旧值，重启恢复会把活着的任务误判失联；
* ``append`` 幂等（``ON CONFLICT DO NOTHING``）：崩溃恢复会重放。
时钟：租约判定用**调用方传入的 now**
------------------------------------
契约把 ``now`` 作为参数传入（``claim_commands(*, now, ...)``），因此"租约过期了吗"
以调用方的时钟为准。分布式设计说明 §15.1 曾建议"租约判定一律用存储侧 ``now()``"——
那与契约冲突：可注入的 ``Clock`` 是这个项目可测试性的地基，服务端时钟会让
``StoreContract`` 里"过期 / 未过期"的用例无法构造。契约优先，此处以此为准。
表结构：建在 ``search_path`` 的第一个 schema
--------------------------------------------
DDL 使用不带 schema 限定的表名，因此建表位置由连接的 ``search_path`` 决定。要隔离
到独立 schema，把 ``options=-csearch_path=xxx`` 放进 DSN（见
:func:`chongming_task.stores.codec` 的用法示例），或用独立数据库。
**建表是条件性的**：三张表都已存在时不执行任何 DDL——只读账号（无 CREATE 权限）
因此可以正常使用本存储（查询、下发指令）。表缺失时才建，多节点同时首次启动由一把
咨询锁串行化。注意这**不是迁移**：表结构变更需要手工迁移，``IF NOT EXISTS`` 只负责
首次建表。
:: note
本实现跑通随包发布的 ``StoreContract``（``tests/integration/test_postgres_store.py``），
并在其上自建并发用例（N 连接并发 save / 并发认领 / 杀掉持有者后端）——契约测试
只覆盖单连接语义，跨进程原子性必须单独验证（见设计说明 §15.4）。
"""
from __future__ import annotations
import json
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from datetime import datetime, timedelta
from typing import Any
try:
import psycopg
from psycopg.conninfo import conninfo_to_dict
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool, PoolClosed, PoolTimeout
except ImportError as exc: # pragma: no cover - 依赖环境，无法在装好驱动的测试里走到
raise ImportError(
"PostgresTaskStore 需要可选依赖 psycopg 与 psycopg-pool。安装：\n"
' pip install "chongming-task[postgres]"\n'
'或：pip install "psycopg[binary]" psycopg-pool'
) from exc
from ..core.state_machine import Transition
from ..errors import ConcurrencyConflict, ConfigurationError, StoreError
from ..models import Progress, TaskSnapshot
from ..states import TERMINAL, TaskStatus
from .codec import (
command_to_row,
decode_dt,
encode_dt,
encode_progress,
row_to_command,
row_to_snapshot,
row_to_transition,
snapshot_to_row,
transition_to_row,
)
from .protocol import Command, TaskQuery
__all__ = ["PostgresTaskStore"]
#: 与 ``codec.snapshot_to_row`` 的键**逐字一致**（顺序也一致）。
#: 单元测试会把它与 codec 的输出比对，防止两边各自漂移。
_TASK_COLUMNS = (
"task_id",
"name",
"pipeline_run_id",
"tags",
"bucket",
"bucket_key",
"status",
"status_version",
"requested_action",
"attempt",
"max_attempts",
"created_at",
"queued_at",
"started_at",
"ended_at",
"last_heartbeat_at",
"progress",
"failure",
"outcome",
"executor",
"pid",
"overridden",
"lease_owner",
"lease_until",
)
#: ``save`` 可以写的列 = 全部列 **减去租约列**。理由见模块文档。
_SAVE_COLUMNS = tuple(column for column in _TASK_COLUMNS if not column.startswith("lease_"))
#: 存成 JSONB 的列。绑定参数时用 ``::jsonb`` 显式转型；读回来时包回 JSON 文本。
_JSON_COLUMNS = frozenset({"tags", "progress", "failure", "outcome"})
#: 存成 TIMESTAMPTZ 的列。写：ISO 文本 → datetime；读：datetime → ISO 文本。
#: 三张表的集合取并集即可：列名在各表间同义，没有歧义。
_DATETIME_COLUMNS = frozenset(
{
"created_at",
"queued_at",
"started_at",
"ended_at",
"last_heartbeat_at",
"lease_until",
"at",
"finished_at",
}
)
#: 建表前的建议锁 key。多个节点同时首次启动时，``CREATE TABLE IF NOT EXISTS``
#: 会竞争系统目录（``pg_type`` 的唯一键），拿它把 DDL 串行化。
_SCHEMA_LOCK_KEY = "chongming_task.schema"
#: 表是否齐全。齐全就完全跳过 DDL——只读账号因此不需要 CREATE 权限。
_SCHEMA_READY_SQL = """
SELECT to_regclass('tasks') IS NOT NULL
AND to_regclass('transitions') IS NOT NULL
AND to_regclass('commands') IS NOT NULL AS ready
"""
_DDL_STATEMENTS: tuple[str, ...] = (
"""
CREATE TABLE IF NOT EXISTS tasks (
task_id TEXT PRIMARY KEY,
name TEXT NOT NULL,
pipeline_run_id TEXT,
tags JSONB NOT NULL,
bucket INTEGER,
bucket_key TEXT,
status TEXT NOT NULL,
status_version INTEGER NOT NULL,
requested_action TEXT,
attempt INTEGER NOT NULL,
max_attempts INTEGER NOT NULL,
created_at TIMESTAMPTZ,
queued_at TIMESTAMPTZ,
started_at TIMESTAMPTZ,
ended_at TIMESTAMPTZ,
last_heartbeat_at TIMESTAMPTZ,
progress JSONB,
failure JSONB,
outcome JSONB,
executor TEXT,
pid INTEGER,
overridden BOOLEAN NOT NULL DEFAULT FALSE,
-- 监督者租约：与 last_heartbeat_at 是两件事，见 TaskSnapshot 的说明。
lease_owner TEXT,
lease_until TIMESTAMPTZ
)
""",
# 索引只建现有查询形状用得到的那些。给 tasks 建 `(lease_until) WHERE status IN (…)`
# 之类的部分索引是**没用的**：租约判定发生在 Python 里（reconcile 取回非终态任务后
# 逐条比较），而通用 query() 拼的是 `status = ANY($1)`，规划器证明不了它等价于单个
# 状态值——索引永远不会被选中，白付写入开销。理由与"什么时候该加"记在
# docs/distributed-design.md §15.3。
"CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks (status)",
"CREATE INDEX IF NOT EXISTS idx_tasks_created_at ON tasks (created_at)",
"CREATE INDEX IF NOT EXISTS idx_tasks_bucket ON tasks (bucket)",
"CREATE INDEX IF NOT EXISTS idx_tasks_pipeline_run ON tasks (pipeline_run_id)",
# 标签精确匹配（`tags @> '["io"]'`）走 GIN；这是 JSONB 相对 TEXT + LIKE 多出来的东西。
"CREATE INDEX IF NOT EXISTS idx_tasks_tags ON tasks USING GIN (tags)",
"""
CREATE TABLE IF NOT EXISTS transitions (
task_id TEXT NOT NULL,
version INTEGER NOT NULL,
from_status TEXT NOT NULL,
to_status TEXT NOT NULL,
at TIMESTAMPTZ NOT NULL,
trigger TEXT NOT NULL,
actor TEXT NOT NULL,
reason TEXT,
PRIMARY KEY (task_id, version)
)
""",
"""
CREATE TABLE IF NOT EXISTS commands (
command_id TEXT PRIMARY KEY,
task_id TEXT NOT NULL,
kind TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL,
requested_by TEXT NOT NULL,
reason TEXT,
lease_until TIMESTAMPTZ,
finished_at TIMESTAMPTZ,
outcome TEXT
)
""",
"CREATE INDEX IF NOT EXISTS idx_commands_task ON commands (task_id)",
# 认领查询只关心未完成的指令，而表在持续变大：部分索引更小也更快。
"CREATE INDEX IF NOT EXISTS idx_commands_pending ON commands (created_at) "
"WHERE finished_at IS NULL",
)
#: "挑选 + 打租约"必须是原子的。``FOR UPDATE SKIP LOCKED`` 让并发认领互不阻塞，
#: 也不会重复取到同一条指令——SQLite 版本用 ``BEGIN IMMEDIATE`` 换来的正是这个保证。
_CLAIM_SQL = """
WITH claimable AS (
SELECT command_id
FROM commands
WHERE finished_at IS NULL
AND (lease_until IS NULL OR lease_until <= %(now)s)
ORDER BY created_at ASC, command_id
LIMIT %(limit)s
FOR UPDATE SKIP LOCKED
)
UPDATE commands SET lease_until = %(until)s
FROM claimable
WHERE commands.command_id = claimable.command_id
RETURNING commands.*
"""
def _describe_dsn(dsn: str) -> str:
"""给报错用的 DSN 摘要：**不含密码**（连接串经常被整体写进日志）。"""
try:
info = conninfo_to_dict(dsn)
except psycopg.Error:
return "<无法解析的 DSN>"
user = info.get("user", "")
host = info.get("host", "localhost")
port = info.get("port", "5432")
dbname = info.get("dbname", user)
return f"{user}@{host}:{port}/{dbname}" if user else f"{host}:{port}/{dbname}"
def _to_db_params(row: Mapping[str, Any]) -> dict[str, Any]:
"""codec 的"扁平原语映射" → 可直接绑定的参数。
两处类型差异在这里收口，而不是散落到每条 SQL：
* 时间列：ISO 文本 → ``datetime``（让驱动按 TIMESTAMPTZ 绑定，而不是依赖
PG 对文本参数的隐式解析）；
* ``overridden``：codec 写的是 0/1（SQLite 没有布尔类型），PG 的 BOOLEAN
列**不接受** smallint（``DatatypeMismatch``），显式转 bool。
"""
params = dict(row)
for column in _DATETIME_COLUMNS:
if column in params:
params[column] = decode_dt(params[column])
if "overridden" in params:
params["overridden"] = bool(params["overridden"])
return params
def _from_db_row(row: Mapping[str, Any]) -> dict[str, Any]:
"""数据库行 → codec 认识的"扁平原语映射"。
``datetime`` 转回 UTC ISO 文本、JSONB 的 Python 对象转回 JSON 文本——codec 的
``decode_*`` 系列吃的是文本。JSONB 的规范化（键序、空白）不影响语义：解码后
比较的是对象而不是字节。
"""
data = dict(row)
for column in _DATETIME_COLUMNS:
value = data.get(column)
if isinstance(value, datetime):
data[column] = encode_dt(value)
for column in _JSON_COLUMNS:
value = data.get(column)
if value is not None:
data[column] = json.dumps(value, ensure_ascii=False)
return data
def _placeholder(column: str) -> str:
"""一个列的绑定占位符；JSON 列显式 ``::jsonb`` 转型（文本参数不会自动转）。"""
return f"%({column})s::jsonb" if column in _JSON_COLUMNS else f"%({column})s"
class PostgresTaskStore:
"""把任务、迁移历史与控制指令存进 PostgreSQL。
``dsn`` 是 libpq 连接串（URL 或 key=value 两种写法都行）。连接池默认
``min_size=1, max_size=8``：按需增长，空闲时只留一条。
"""
def __init__(
self,
dsn: str,
*,
min_size: int = 1,
max_size: int = 8,
timeout: float = 5.0,
) -> None:
if min_size < 1:
raise ConfigurationError(f"min_size 至少为 1，收到 {min_size}")
if max_size < min_size:
raise ConfigurationError(f"max_size({max_size}) 不能小于 min_size({min_size})")
if timeout <= 0:
raise ConfigurationError(f"timeout 必须为正数，收到 {timeout}")
self._dsn = dsn
self._closed = False
connect_timeout = max(1, int(timeout))
# 先直连一次：连接池的失败信息藏在日志里，直连能把真实原因（认证失败、
# 库不存在、网络不通）带进异常。
try:
with psycopg.connect(dsn, connect_timeout=connect_timeout):
pass
except psycopg.Error as exc:
raise StoreError(f"无法连接 Postgres（{_describe_dsn(dsn)}）：{exc}") from exc
self._pool: ConnectionPool[psycopg.Connection[dict[str, Any]]] = ConnectionPool(
dsn,
kwargs={
"row_factory": dict_row,
"application_name": "chongming-task",
"connect_timeout": connect_timeout,
},
min_size=min_size,
max_size=max_size,
timeout=timeout,
open=False,
# 空闲连接可能被对端悄悄杀掉（重启、超时踢线、网络抖动）。取出时做一次
# 轻量探活，坏连接会被丢弃并补新的——代价是每次借出多一次空查询。
check=ConnectionPool.check_connection,
)
try:
self._pool.open(wait=True, timeout=timeout)
except PoolTimeout as exc: # pragma: no cover - 直连已通过，这里基本只剩竞态
self._pool.close()
raise StoreError(f"连接池初始化超时（{_describe_dsn(dsn)}）") from exc
self._ensure_schema()
# ------------------------------------------------------------------
# 内部
# ------------------------------------------------------------------
@contextmanager
def _connection(self) -> Iterator[psycopg.Connection[dict[str, Any]]]:
"""借一条池连接；正常退出即提交。``close()`` 之后再用抛 :class:`StoreError`。"""
if self._closed:
raise StoreError("PostgresTaskStore 已关闭：close() 之后不能再使用。")
try:
with self._pool.connection() as conn:
yield conn
except PoolClosed as exc:
raise StoreError("PostgresTaskStore 已关闭：close() 之后不能再使用。") from exc
def _ensure_schema(self) -> None:
"""表缺失时才建；已齐全则一个 DDL 都不执行（只读账号也能用）。"""
try:
with self._connection() as conn:
if self._tables_ready(conn):
return
# 多节点同时首次启动会竞争系统目录。拿一把事务级建议锁串行化，
# 拿到锁后**再查一次**——别人可能已经建好了。
conn.execute(
"SELECT pg_advisory_xact_lock(hashtext(%(key)s))", {"key": _SCHEMA_LOCK_KEY}
)
if self._tables_ready(conn):
return
for statement in _DDL_STATEMENTS:
conn.execute(statement)
except psycopg.Error as exc:
raise StoreError(
f"初始化 Postgres 表结构失败（{_describe_dsn(self._dsn)}）：{exc}\n"
"需要建表权限；若用只读账号，请先由管理员在目标 schema 里建好三张表。"
) from exc
@staticmethod
def _tables_ready(conn: psycopg.Connection[dict[str, Any]]) -> bool:
row = conn.execute(_SCHEMA_READY_SQL).fetchone()
return bool(row is not None and row["ready"])
# ------------------------------------------------------------------
# 1. 任务快照
# ------------------------------------------------------------------
def create(self, snapshot: TaskSnapshot) -> None:
"""写入新任务；``task_id`` 已存在则抛错。"""
params = _to_db_params(snapshot_to_row(snapshot))
columns = ", ".join(_TASK_COLUMNS)
values = ", ".join(_placeholder(column) for column in _TASK_COLUMNS)
try:
with self._connection() as conn:
conn.execute(f"INSERT INTO tasks ({columns}) VALUES ({values})", params)
except psycopg.errors.UniqueViolation as exc:
raise StoreError(f"任务已存在：{snapshot.task_id}") from exc
def save(self, snapshot: TaskSnapshot, *, expect_version: int) -> None:
"""按乐观并发更新快照。
``UPDATE ... WHERE task_id = ? AND status_version = ?`` 让"版本校验 + 写入"
成为**同一条语句**——拆成先查后写，跨进程就必然丢更新。
**不写租约列**（见 ``_SAVE_COLUMNS``）。
"""
params = _to_db_params(snapshot_to_row(snapshot))
params["expect_version"] = expect_version
assignments = ", ".join(
f"{column} = {_placeholder(column)}" for column in _SAVE_COLUMNS if column != "task_id"
)
with self._connection() as conn:
cursor = conn.execute(
f"UPDATE tasks SET {assignments} "
"WHERE task_id = %(task_id)s AND status_version = %(expect_version)s",
params,
)
if cursor.rowcount == 1:
return
existing = conn.execute(
"SELECT status_version FROM tasks WHERE task_id = %(task_id)s",
{"task_id": snapshot.task_id},
).fetchone()
if existing is None:
raise StoreError(f"任务不存在：{snapshot.task_id}")
raise ConcurrencyConflict(snapshot.task_id, expect_version, existing["status_version"])
def get(self, task_id: str) -> TaskSnapshot | None:
"""读取单个任务。"""
with self._connection() as conn:
row = conn.execute(
"SELECT * FROM tasks WHERE task_id = %(task_id)s", {"task_id": task_id}
).fetchone()
return row_to_snapshot(_from_db_row(row)) if row is not None else None
def query(self, query: TaskQuery | None = None) -> list[TaskSnapshot]:
"""按条件查询。过滤**下推到 SQL**，不是在 Python 里筛。
排序与 SQLite 实现对齐：``created_at`` 升序时 NULL 在前、降序时在后
（SQLite 把 NULL 当作最小值，PG 默认相反，所以显式写 ``NULLS FIRST/LAST``）。
``created_at`` 相同时按 ``task_id`` 决出先后，保证结果稳定可复现。
"""
condition = query if query is not None else TaskQuery()
clauses: list[str] = []
params: dict[str, Any] = {}
if condition.statuses is not None:
clauses.append("status = ANY(%(statuses)s)")
params["statuses"] = [status.value for status in condition.statuses]
if condition.name is not None:
clauses.append("name = %(name)s")
params["name"] = condition.name
if condition.name_prefix is not None:
# 与 SQLite 实现同样不转义 LIKE 通配符：两边的语义保持一致。
clauses.append("name LIKE %(name_prefix)s")
params["name_prefix"] = f"{condition.name_prefix}%"
if condition.tag is not None:
# JSONB 包含：`["cpu", "io"] @> ["io"]` 为真，`["ios"] @> ["io"]` 为假。
# 这就是"标签精确匹配"的正确表达（不要退回 LIKE 子串匹配）。
clauses.append("tags @> %(tag)s::jsonb")
params["tag"] = json.dumps([condition.tag], ensure_ascii=False)
if condition.buckets is not None:
clauses.append("bucket = ANY(%(buckets)s)")
params["buckets"] = list(condition.buckets)
if condition.pipeline_run_id is not None:
clauses.append("pipeline_run_id = %(pipeline_run_id)s")
params["pipeline_run_id"] = condition.pipeline_run_id
where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
direction = "DESC NULLS LAST" if condition.newest_first else "ASC NULLS FIRST"
statement = f"SELECT * FROM tasks{where} ORDER BY created_at {direction}, task_id"
if condition.limit is not None:
statement += " LIMIT %(limit)s"
params["limit"] = condition.limit
with self._connection() as conn:
rows = conn.execute(statement, params).fetchall()
return [row_to_snapshot(_from_db_row(row)) for row in rows]
def counts(self) -> dict[TaskStatus, int]:
"""各状态计数（``GROUP BY``，不把全部记录读回来数）。"""
with self._connection() as conn:
rows = conn.execute(
"SELECT status, COUNT(*) AS total FROM tasks GROUP BY status"
).fetchall()
return {TaskStatus(row["status"]): int(row["total"]) for row in rows}
# ------------------------------------------------------------------
# 2. 审计轨迹
# ------------------------------------------------------------------
def append(self, transition: Transition) -> None:
"""追加迁移记录；同一 ``(task_id, version)`` 重复写入被忽略（幂等）。"""
params = _to_db_params(transition_to_row(transition))
with self._connection() as conn:
conn.execute(
"INSERT INTO transitions "
"(task_id, version, from_status, to_status, at, trigger, actor, reason) "
"VALUES (%(task_id)s, %(version)s, %(from_status)s, %(to_status)s, %(at)s, "
"%(trigger)s, %(actor)s, %(reason)s) "
"ON CONFLICT (task_id, version) DO NOTHING",
params,
)
def history(self, task_id: str, *, since_version: int = 0) -> list[Transition]:
"""读取迁移历史，``version > since_version``，按版本升序。"""
with self._connection() as conn:
rows = conn.execute(
"SELECT * FROM transitions "
"WHERE task_id = %(task_id)s AND version > %(since_version)s "
"ORDER BY version",
{"task_id": task_id, "since_version": since_version},
).fetchall()
return [row_to_transition(_from_db_row(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:
"""只更新进度、心跳与监督者租约，**不碰** ``status_version``。
任务不存在时静默忽略（心跳晚到一步不该报错）。四个参数各自独立可选。
"""
assignments: list[str] = []
params: dict[str, Any] = {"task_id": task_id}
if progress is not None:
assignments.append("progress = %(progress)s::jsonb")
params["progress"] = encode_progress(progress)
if heartbeat_at is not None:
assignments.append("last_heartbeat_at = %(heartbeat_at)s")
params["heartbeat_at"] = heartbeat_at
if lease_owner is not None:
assignments.append("lease_owner = %(lease_owner)s")
params["lease_owner"] = lease_owner
if lease_until is not None:
assignments.append("lease_until = %(lease_until)s")
params["lease_until"] = lease_until
if not assignments:
return
with self._connection() as conn:
conn.execute(
f"UPDATE tasks SET {', '.join(assignments)} WHERE task_id = %(task_id)s",
params,
)
# ------------------------------------------------------------------
# 4. 控制指令
# ------------------------------------------------------------------
def submit_command(self, command: Command) -> None:
"""写入一条控制指令。存储层只传递，不解释、不执行。"""
params = _to_db_params(command_to_row(command))
try:
with self._connection() as conn:
conn.execute(
"INSERT INTO commands "
"(command_id, task_id, kind, created_at, requested_by, reason, "
" lease_until, finished_at, outcome) "
"VALUES (%(command_id)s, %(task_id)s, %(kind)s, %(created_at)s, "
"%(requested_by)s, %(reason)s, %(lease_until)s, %(finished_at)s, "
"%(outcome)s)",
params,
)
except psycopg.errors.UniqueViolation as exc:
raise StoreError(f"指令已存在：{command.command_id}") from exc
def claim_commands(
self, *, now: datetime, lease_seconds: float, limit: int = 50
) -> list[Command]:
"""取走一批可执行指令并设置租约；单条 SQL，天然原子。
结果按 ``(created_at, command_id)`` 排序后再返回：``RETURNING`` 的行序没有
保证，而契约的用词是"按 created_at 升序取"，排序把它变成确定的。
"""
params = {
"now": now,
"until": now + timedelta(seconds=lease_seconds),
"limit": limit,
}
with self._connection() as conn:
rows = conn.execute(_CLAIM_SQL, params).fetchall()
rows.sort(key=lambda row: (row["created_at"], row["command_id"]))
return [row_to_command(_from_db_row(row)) for row in rows]
def complete_command(
self, command_id: str, *, now: datetime, outcome: str | None = None
) -> None:
"""标记指令已完成，使其不再被取走。"""
with self._connection() as conn:
conn.execute(
"UPDATE commands SET finished_at = %(now)s, outcome = %(outcome)s, "
"lease_until = NULL WHERE command_id = %(command_id)s",
{"command_id": command_id, "now": now, "outcome": outcome},
)
def release_command(self, command_id: str, *, now: datetime) -> None:
"""交还一条已取走但无人认领的指令：只清租约，**不**标记完成。
标记完成会让真正的持有者永远收不到这次请求；等租约过期又会被持续续租饿死。
所以需要第三种结局，见 :meth:`TaskStore.release_command`。
"""
del now # 交还与时间无关；接口保留参数以维持契约形状
with self._connection() as conn:
conn.execute(
"UPDATE commands SET lease_until = NULL "
"WHERE command_id = %(command_id)s AND finished_at IS NULL",
{"command_id": command_id},
)
# ------------------------------------------------------------------
# 运维
# ------------------------------------------------------------------
def prune(self, *, keep_terminal: int) -> int:
"""丢弃最旧的终态任务及其历史，返回丢弃数量。非终态任务绝不触碰。
一条 CTE 语句完成"选受害者 → 删历史 → 删任务"：两个 DELETE 与选择看到的是
同一份物化结果，不会出现"历史删了任务没删"的中间态（引用两次的 CTE 会被
Postgres 物化一次；数据修改型 CTE 即使不被引用也保证执行）。
顺序：按 ``created_at`` **降序**保留最新的 ``keep_terminal`` 条，``OFFSET``
跳过它们，剩下的就是"最旧的"要删的。升序 + ``OFFSET`` 会删反——那是把最新
的丢掉。NULL 在 SQLite 里是最小值（视为最旧），所以这里用 ``NULLS LAST``
对齐同一语义。
"""
params = {
"statuses": [status.value for status in TERMINAL],
"keep": max(0, keep_terminal),
}
statement = """
WITH doomed AS (
SELECT task_id FROM tasks
WHERE status = ANY(%(statuses)s)
ORDER BY created_at DESC NULLS LAST, task_id DESC
OFFSET %(keep)s
), purged AS (
DELETE FROM transitions WHERE task_id IN (SELECT task_id FROM doomed)
)
DELETE FROM tasks WHERE task_id IN (SELECT task_id FROM doomed)
RETURNING task_id
"""
with self._connection() as conn:
rows = conn.execute(statement, params).fetchall()
return len(rows)
def close(self) -> None:
"""关闭连接池。**幂等**；关闭后任何操作都会抛 :class:`StoreError`。
正在被借出的连接会在归还时关闭（``psycopg_pool`` 的语义），所以调用方应确保
没有线程还在使用——``Controller.close(close_store=True)`` 已经等任务收尾，
直接用它即可。
"""
if self._closed:
return
self._closed = True
self._pool.close()
./src\chongming_task\stores\protocol.py
"""存储抽象。
:class:`TaskStore` 是整个持久层的**唯一契约**。内存实现、SQLite 实现、以及将来任何
第三方实现（Postgres、Redis、对象存储……）都只需满足它，上层（:class:`Controller`、
CLI、看板）不必知道数据从哪来。
本模块只依赖标准库，**不依赖 pytest**：契约测试套件在
``chongming_task.testing.store_contract``，第三方实现应当拿它验证自己。
四类数据，写入频次与一致性要求完全不同，因此刻意分开
----------------------------------------------------
1. **任务快照**（``create`` / ``save`` / ``get`` / ``query``）——查询面，参与**乐观
并发**（``status_version``）。写入频率低（每次状态迁移一次）。
2. **审计轨迹**（``append`` / ``history``）——**只追加，永不修改**。这是状态机的
外部证据，删改它就等于篡改审计。
3. **高频旁路**（``touch``）——进度与心跳，可能每秒几十次。它们**不参与乐观并发**：
否则会与状态写入持续冲突，把版本号变成噪声。
4. **控制指令**（``submit_command`` / ``claim_commands`` / ``complete_command``）——
跨进程控制的媒介。CLI 无法给别的进程发信号，只能写一条指令让持有者取走。
关于"谁来执行指令"
------------------
存储层**只负责传递**，不解释、不执行。执行者是持有该任务的进程（见
:class:`~chongming_task.controller.Controller`）。这样存储实现可以保持极简，
也可以换成任何能 CAS 的后端。
"""
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from enum import StrEnum
from typing import Protocol, runtime_checkable
from ..core.state_machine import Transition
from ..models import Progress, TaskSnapshot
from ..states import TaskStatus
__all__ = [
"Command",
"CommandKind",
"TaskQuery",
"TaskStore",
]
class CommandKind(StrEnum):
"""跨进程控制动作。"""
STOP = "stop"
PAUSE = "pause"
RESUME = "resume"
@dataclass(frozen=True, slots=True)
class Command:
"""一条待执行的控制指令。
``lease_until`` 是**租约**：取走指令的一方获得一段独占时间，期间别的进程不会
重复取到它。若持有者在租约到期前没有 ``complete_command``（例如进程崩了），
指令会被重新放出——这样既不会重复执行，也不会因为一次崩溃就永久卡住。
"""
command_id: str
task_id: str
kind: CommandKind
created_at: datetime
requested_by: str = "unknown"
reason: str | None = None
lease_until: datetime | None = None
finished_at: datetime | None = None
outcome: str | None = None
@property
def is_open(self) -> bool:
"""是否尚未完成。"""
return self.finished_at is None
@dataclass(frozen=True, slots=True)
class TaskQuery:
"""查询条件。所有字段都是可选的，``None`` 表示不限制。"""
statuses: frozenset[TaskStatus] | None = None
name_prefix: str | None = None
tag: str | None = None
name: str | None = None
buckets: frozenset[int] | None = None
"""只取这些桶里的任务。
为将来的分布式预留：一台机器只关心自己拥有的那组桶，应当能在存储层直接过滤，
而不是取回全部再筛。单桶查询传单元素集合即可。
"""
limit: int | None = None
newest_first: bool = False
pipeline_run_id: str | None = None
"""只取属于某次管道运行的任务。
没有它，"把这条管道的所有节点列出来"就只能全表拉回来自己筛——而管道恰恰是最需要
按整体查看的东西。
"""
@runtime_checkable
class TaskStore(Protocol):
"""持久层契约。
实现者请用 ``chongming_task.testing.store_contract.StoreContract`` 验证自己：
同一套用例必须对每个实现都通过，否则"换后端"就是空话。
"""
# ------------------------------------------------------------------
# 1. 任务快照（参与乐观并发）
# ------------------------------------------------------------------
def create(self, snapshot: TaskSnapshot) -> None:
"""写入一个新任务。``task_id`` 已存在则抛 :class:`~chongming_task.errors.StoreError`。"""
...
def save(self, snapshot: TaskSnapshot, *, expect_version: int) -> None:
"""更新任务快照。
``expect_version`` 是写入前读到的 ``status_version``；不一致则抛
:class:`~chongming_task.errors.ConcurrencyConflict`，**不做自动重试**——
每个任务的写入者应当唯一，冲突意味着存在第二个写入者，那是设计缺陷。
"""
...
def get(self, task_id: str) -> TaskSnapshot | None:
"""读取单个任务；不存在返回 ``None``。"""
...
def query(self, query: TaskQuery | None = None) -> list[TaskSnapshot]:
"""按条件查询，按 ``created_at`` 升序（``newest_first=True`` 则降序）。"""
...
def counts(self) -> dict[TaskStatus, int]:
"""各状态的任务数。"""
...
# ------------------------------------------------------------------
# 2. 审计轨迹（只追加）
# ------------------------------------------------------------------
def append(self, transition: Transition) -> None:
"""追加一条状态迁移记录。
应当按 ``task_id`` + ``version`` **幂等**写入：重放同一批记录（例如崩溃后
恢复）不应产生重复行。
"""
...
def history(self, task_id: str, *, since_version: int = 0) -> list[Transition]:
"""读取迁移历史，``version > since_version``，按版本升序。"""
...
# ------------------------------------------------------------------
# 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:
"""仅更新进度、心跳与监督者租约，**不校验也不递增** ``status_version``。
两个可选信号是**独立的**，实现不得用一个去推导另一个：
* ``heartbeat_at`` 由任务自己的 ``checkpoint()`` 驱动，回答"任务在推进吗"。
* ``lease_owner`` / ``lease_until`` 由监督进程定期续租，回答"还有人在管它吗"。
任务不存在时静默忽略：任务可能刚被清理，而心跳晚到一步不该报错。
实现注意：``save`` **不得**写租约列（由本方法独占），否则一次状态迁移就会把
租约覆盖回旧值，让重启恢复把活着的任务误判为失联。
"""
...
# ------------------------------------------------------------------
# 4. 控制指令（跨进程）
# ------------------------------------------------------------------
def submit_command(self, command: Command) -> None:
"""写入一条控制指令。存储层不解释它。"""
...
def claim_commands(
self, *, now: datetime, lease_seconds: float, limit: int = 50
) -> list[Command]:
"""取走一批可执行指令，并为其设置租约。
应当**同时**满足两个条件才可被取走：未被完成，且（无租约或租约已过期）。
实现应保证并发调用下同一条指令不会被两个进程同时取到。
"""
...
def complete_command(
self, command_id: str, *, now: datetime, outcome: str | None = None
) -> None:
"""标记指令已完成，使其不再被取走。"""
...
def release_command(self, command_id: str, *, now: datetime) -> None:
"""放弃一条已取走的指令，交还给别人。
没有它，多进程下会出问题：任何**能取指令**的进程都可能取到不属于自己的
任务（它手上没有那个任务）。若它把指令标记完成，真正的持有者就永远收不到
这次停止请求——一个静默的丢指令。若它什么都不做、只等租约过期，那么只要
它持续轮询就会不断续租，真正的持有者会被无限期饿死。
因此需要第三种结局：``claimed → release``。取到不属于自己的指令时立即交还。
"""
...
# ------------------------------------------------------------------
# 运维
# ------------------------------------------------------------------
def prune(self, *, keep_terminal: int) -> int:
"""丢弃最旧的终态任务（含其历史），返回丢弃数量。"""
...
def close(self) -> None:
"""释放资源。对内存实现是空操作。"""
...
./src\chongming_task\stores\sqlite.py
"""SQLite 存储实现（标准库 ``sqlite3``，不引入任何第三方依赖）。
为什么默认选它
--------------
* **跨进程可见**——这正是它存在的首要理由：另一个进程里的 CLI 要能查询到任务。
* **无 IPC、无网络层**，天然跨平台，一个文件就是全部状态。
* **顺带就是持久化**，重启后可查。
并发模型
--------
* 进程内：所有操作由一把可重入锁串行化，使"读—判断—写"在进程内是原子的。
* 跨进程：可能被多个进程同时写（例如任务进程写状态、CLI 写指令），因此
取指令等复合操作用 ``BEGIN IMMEDIATE`` + ``COMMIT``，并依赖
``status_version`` 做乐观并发。``busy_timeout`` 让短暂争抢自动重试。
三个刻意的取舍
--------------
1. **不用 ORM。** 表结构与 :class:`~chongming_task.models.TaskSnapshot` 一一对应，
读起来就是 SQL，排障时不用先翻译一层。
2. **时间一律以 UTC ISO 字符串存储。** 本地时间原样写入会在跨时区排序时出错。
3. **指令只有租约，没有优先级。** 需要优先级时应当在上层排队，而不是让存储层
长出调度语义。
"""
from __future__ import annotations
import json
import sqlite3
import threading
from dataclasses import replace
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any
from ..core.state_machine import Transition
from ..errors import ConcurrencyConflict, StoreError
from ..models import Progress, TaskSnapshot
from ..states import TERMINAL, TaskStatus
from .codec import (
command_to_row,
encode_dt,
row_to_command,
row_to_snapshot,
row_to_transition,
snapshot_to_row,
transition_to_row,
)
from .protocol import Command, TaskQuery
__all__ = ["SqliteTaskStore"]
_TASK_COLUMNS = (
"task_id",
"name",
"pipeline_run_id",
"tags",
"bucket",
"bucket_key",
"status",
"status_version",
"requested_action",
"attempt",
"max_attempts",
"created_at",
"queued_at",
"started_at",
"ended_at",
"last_heartbeat_at",
"progress",
"failure",
"outcome",
"executor",
"pid",
"overridden",
"lease_owner",
"lease_until",
)
#: ``save`` 可以写的列 = 全部列 **减去租约列**。
#:
#: 租约由 :meth:`SqliteTaskStore.touch` 独占写入。若 ``save`` 也写它，一次状态迁移
#: 就会把租约覆盖成内存快照里的旧值（甚至 ``None``）——同一行两条写入路径互相打架，
#: 和心跳曾经踩过的是同一个坑：静默、且只有重启恢复时才会以"任务被误判失联"的形式
#: 暴露出来。
_SAVE_COLUMNS = tuple(column for column in _TASK_COLUMNS if not column.startswith("lease_"))
_SCHEMA = """
CREATE TABLE IF NOT EXISTS tasks (
task_id TEXT PRIMARY KEY,
name TEXT NOT NULL,
pipeline_run_id TEXT,
tags TEXT NOT NULL,
bucket INTEGER,
bucket_key TEXT,
status TEXT NOT NULL,
status_version INTEGER NOT NULL,
requested_action TEXT,
attempt INTEGER NOT NULL,
max_attempts INTEGER NOT NULL,
created_at TEXT,
queued_at TEXT,
started_at TEXT,
ended_at TEXT,
last_heartbeat_at TEXT,
progress TEXT,
failure TEXT,
outcome TEXT,
executor TEXT,
pid INTEGER,
overridden INTEGER NOT NULL DEFAULT 0,
-- 监督者租约：与 last_heartbeat_at 是两件事，见 TaskSnapshot 的说明。
lease_owner TEXT,
lease_until TEXT
);
CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status);
CREATE INDEX IF NOT EXISTS idx_tasks_created_at ON tasks(created_at);
-- 分桶索引是为将来的分布式预留的：一台机器"取我拥有的那些桶"会是一条高频查询，
-- 没有索引就只能全表扫。
CREATE INDEX IF NOT EXISTS idx_tasks_bucket ON tasks(bucket);
-- 管道索引同理：看板/CLI 最常问的就是"这条管道的节点现在怎样了"。
CREATE INDEX IF NOT EXISTS idx_tasks_pipeline_run ON tasks(pipeline_run_id);
CREATE TABLE IF NOT EXISTS transitions (
task_id TEXT NOT NULL,
version INTEGER NOT NULL,
from_status TEXT NOT NULL,
to_status TEXT NOT NULL,
at TEXT NOT NULL,
trigger TEXT NOT NULL,
actor TEXT NOT NULL,
reason TEXT,
PRIMARY KEY (task_id, version)
);
CREATE TABLE IF NOT EXISTS commands (
command_id TEXT PRIMARY KEY,
task_id TEXT NOT NULL,
kind TEXT NOT NULL,
created_at TEXT NOT NULL,
requested_by TEXT NOT NULL,
reason TEXT,
lease_until TEXT,
finished_at TEXT,
outcome TEXT
);
CREATE INDEX IF NOT EXISTS idx_commands_pending ON commands(finished_at, lease_until);
CREATE INDEX IF NOT EXISTS idx_commands_task ON commands(task_id);
"""
class SqliteTaskStore:
"""把任务、迁移历史与控制指令存进一个 SQLite 文件。
传 ``":memory:"`` 可得到一个进程内容量库，语义与文件库完全一致——因此同一套
契约测试可以同时跑在两者上。
"""
def __init__(self, path: str | Path = "chongming-tasks.db", *, timeout: float = 5.0) -> None:
self._path = str(path)
if self._path != ":memory:":
parent = Path(self._path).expanduser().resolve().parent
if not parent.is_dir():
# SQLite 只会报 "unable to open database file"，完全看不出是目录不存在。
# 这是最容易浪费时间的错误之一，所以在这里就说清楚。
raise StoreError(f"数据库所在目录不存在：{parent}。请先创建它，或改用别的路径。")
self._lock = threading.RLock()
self._conn = sqlite3.connect(
self._path,
timeout=timeout,
check_same_thread=False, # 由自己的锁串行化，见模块文档
isolation_level=None, # 显式管理事务
)
self._conn.row_factory = sqlite3.Row
# WAL 让读不阻塞写；busy_timeout 让跨进程争抢自动重试而不是立刻报错。
self._conn.execute("PRAGMA journal_mode=WAL")
self._conn.execute("PRAGMA synchronous=NORMAL")
self._conn.execute(f"PRAGMA busy_timeout={int(timeout * 1000)}")
self._conn.executescript(_SCHEMA)
@property
def path(self) -> str:
"""数据库文件路径，供诊断使用。"""
return self._path
# ------------------------------------------------------------------
# 1. 任务快照
# ------------------------------------------------------------------
def create(self, snapshot: TaskSnapshot) -> None:
"""写入新任务；``task_id`` 已存在则抛错。"""
row = snapshot_to_row(snapshot)
placeholders = ", ".join("?" for _ in _TASK_COLUMNS)
with self._lock:
try:
self._conn.execute(
f"INSERT INTO tasks ({', '.join(_TASK_COLUMNS)}) VALUES ({placeholders})",
tuple(row[column] for column in _TASK_COLUMNS),
)
except sqlite3.IntegrityError as exc:
raise StoreError(f"任务已存在：{snapshot.task_id}") from exc
def save(self, snapshot: TaskSnapshot, *, expect_version: int) -> None:
"""按乐观并发更新快照。
``UPDATE ... WHERE status_version = ?`` 让"版本校验"和"写入"成为同一条
SQL——拆成先查后写就会在跨进程场景下产生竞态。
**不写租约列**（见 ``_SAVE_COLUMNS``）：那些列由 :meth:`touch` 独占，否则一次
状态迁移就会把租约覆盖回旧值，导致重启恢复误判任务失联。
"""
row = snapshot_to_row(snapshot)
assignments = ", ".join(f"{column} = ?" for column in _SAVE_COLUMNS if column != "task_id")
values = [row[column] for column in _SAVE_COLUMNS if column != "task_id"]
with self._lock:
cursor = self._conn.execute(
f"UPDATE tasks SET {assignments} WHERE task_id = ? AND status_version = ?",
(*values, snapshot.task_id, expect_version),
)
if cursor.rowcount == 1:
return
existing = self._conn.execute(
"SELECT status_version FROM tasks WHERE task_id = ?", (snapshot.task_id,)
).fetchone()
if existing is None:
raise StoreError(f"任务不存在：{snapshot.task_id}")
raise ConcurrencyConflict(snapshot.task_id, expect_version, existing["status_version"])
def get(self, task_id: str) -> TaskSnapshot | None:
"""读取单个任务。"""
with self._lock:
row = self._conn.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,)).fetchone()
return row_to_snapshot(dict(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()
clauses: list[str] = []
params: list[Any] = []
if condition.statuses is not None:
placeholders = ", ".join("?" for _ in condition.statuses)
clauses.append(f"status IN ({placeholders})")
params.extend(status.value for status in condition.statuses)
if condition.name is not None:
clauses.append("name = ?")
params.append(condition.name)
if condition.name_prefix is not None:
clauses.append("name LIKE ?")
params.append(f"{condition.name_prefix}%")
if condition.tag is not None:
# 标签存成 JSON 数组；用 LIKE 匹配 "tag" 带引号的形式，避免 "io" 命中 "ios"。
clauses.append("tags LIKE ?")
params.append(f'%"{condition.tag}"%')
if condition.buckets is not None:
placeholders = ", ".join("?" for _ in condition.buckets)
clauses.append(f"bucket IN ({placeholders})")
params.extend(condition.buckets)
if condition.pipeline_run_id is not None:
clauses.append("pipeline_run_id = ?")
params.append(condition.pipeline_run_id)
where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
direction = "DESC" if condition.newest_first else "ASC"
statement = f"SELECT * FROM tasks{where} ORDER BY created_at {direction}"
if condition.limit is not None:
statement += " LIMIT ?"
params.append(condition.limit)
with self._lock:
rows = self._conn.execute(statement, params).fetchall()
return [row_to_snapshot(dict(row)) for row in rows]
def counts(self) -> dict[TaskStatus, int]:
"""各状态计数。"""
with self._lock:
rows = self._conn.execute(
"SELECT status, COUNT(*) AS total FROM tasks GROUP BY status"
).fetchall()
return {TaskStatus(row["status"]): row["total"] for row in rows}
# ------------------------------------------------------------------
# 2. 审计轨迹
# ------------------------------------------------------------------
def append(self, transition: Transition) -> None:
"""追加迁移记录；同一版本重复写入被忽略（幂等）。"""
row = transition_to_row(transition)
with self._lock:
self._conn.execute(
"INSERT OR IGNORE INTO transitions "
"(task_id, version, from_status, to_status, at, trigger, actor, reason) "
"VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
(
row["task_id"],
row["version"],
row["from_status"],
row["to_status"],
row["at"],
row["trigger"],
row["actor"],
row["reason"],
),
)
def history(self, task_id: str, *, since_version: int = 0) -> list[Transition]:
"""读取迁移历史。"""
with self._lock:
rows = self._conn.execute(
"SELECT * FROM transitions WHERE task_id = ? AND version > ? ORDER BY version",
(task_id, since_version),
).fetchall()
return [row_to_transition(dict(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:
"""只更新进度、心跳与监督者租约。
刻意不碰 ``status_version``：它每次跳动都会让状态写入的乐观并发失败，
把版本号变成噪声。任务不存在时静默忽略——心跳晚到一步不该报错。
租约与心跳是**两个独立的信号**，见 :class:`~chongming_task.models.TaskSnapshot`
的说明：心跳回答"任务在推进吗"，租约回答"还有人在管它吗"。用一个去推另一个
都会出错。
"""
assignments: list[str] = []
params: list[Any] = []
if progress is not None:
assignments.append("progress = ?")
params.append(
json.dumps(
{
"done": progress.done,
"total": progress.total,
"unit": progress.unit,
"message": progress.message,
"updated_at": encode_dt(progress.updated_at),
},
ensure_ascii=False,
)
)
if heartbeat_at is not None:
assignments.append("last_heartbeat_at = ?")
params.append(encode_dt(heartbeat_at))
if lease_owner is not None:
assignments.append("lease_owner = ?")
params.append(lease_owner)
if lease_until is not None:
assignments.append("lease_until = ?")
params.append(encode_dt(lease_until))
if not assignments:
return
with self._lock:
self._conn.execute(
f"UPDATE tasks SET {', '.join(assignments)} WHERE task_id = ?",
(*params, task_id),
)
# ------------------------------------------------------------------
# 4. 控制指令
# ------------------------------------------------------------------
def submit_command(self, command: Command) -> None:
"""写入一条控制指令。"""
row = command_to_row(command)
with self._lock:
try:
self._conn.execute(
"INSERT INTO commands "
"(command_id, task_id, kind, created_at, requested_by, reason, "
" lease_until, finished_at, outcome) "
"VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
(
row["command_id"],
row["task_id"],
row["kind"],
row["created_at"],
row["requested_by"],
row["reason"],
row["lease_until"],
row["finished_at"],
row["outcome"],
),
)
except sqlite3.IntegrityError as exc:
raise StoreError(f"指令已存在：{command.command_id}") from exc
def claim_commands(
self, *, now: datetime, lease_seconds: float, limit: int = 50
) -> list[Command]:
"""取走可执行指令并设置租约。
整个"挑选 + 打租约"必须在一个 ``BEGIN IMMEDIATE`` 事务里完成，否则两个
进程可能同时取到同一条指令，于是同一次停止被下发两次。
"""
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
try:
rows = self._conn.execute(
"SELECT * FROM commands "
"WHERE finished_at IS NULL AND (lease_until IS NULL OR lease_until <= ?) "
"ORDER BY created_at LIMIT ?",
(encode_dt(now), limit),
).fetchall()
claimed: list[Command] = []
for row in rows:
command = row_to_command(dict(row))
leased = replace(command, lease_until=now + timedelta(seconds=lease_seconds))
self._conn.execute(
"UPDATE commands SET lease_until = ? WHERE command_id = ?",
(encode_dt(leased.lease_until), command.command_id),
)
claimed.append(leased)
self._conn.execute("COMMIT")
except BaseException:
self._conn.execute("ROLLBACK")
raise
return claimed
def complete_command(
self, command_id: str, *, now: datetime, outcome: str | None = None
) -> None:
"""标记指令已完成。"""
with self._lock:
self._conn.execute(
"UPDATE commands SET finished_at = ?, outcome = ?, lease_until = NULL "
"WHERE command_id = ?",
(encode_dt(now), outcome, command_id),
)
def release_command(self, command_id: str, *, now: datetime) -> None:
"""交还一条已取走但无人认领的指令。"""
del now
with self._lock:
self._conn.execute(
"UPDATE commands SET lease_until = NULL "
"WHERE command_id = ? AND finished_at IS NULL",
(command_id,),
)
def commands(self) -> list[Command]:
"""全部指令，供诊断与测试使用（不属于 :class:`TaskStore` 契约）。"""
with self._lock:
rows = self._conn.execute("SELECT * FROM commands ORDER BY created_at").fetchall()
return [row_to_command(dict(row)) for row in rows]
# ------------------------------------------------------------------
# 运维
# ------------------------------------------------------------------
def prune(self, *, keep_terminal: int) -> int:
"""丢弃最旧的终态任务及其历史。"""
terminal_values = tuple(status.value for status in TERMINAL)
placeholders = ", ".join("?" for _ in terminal_values)
with self._lock:
self._conn.execute("BEGIN IMMEDIATE")
try:
rows = self._conn.execute(
f"SELECT task_id FROM tasks WHERE status IN ({placeholders}) "
"ORDER BY created_at ASC",
terminal_values,
).fetchall()
excess = len(rows) - max(0, keep_terminal)
if excess <= 0:
self._conn.execute("COMMIT")
return 0
victims = [row["task_id"] for row in rows[:excess]]
marks = ", ".join("?" for _ in victims)
self._conn.execute(f"DELETE FROM transitions WHERE task_id IN ({marks})", victims)
self._conn.execute(f"DELETE FROM tasks WHERE task_id IN ({marks})", victims)
self._conn.execute("COMMIT")
except BaseException:
self._conn.execute("ROLLBACK")
raise
return excess
def close(self) -> None:
"""关闭连接。"""
with self._lock:
self._conn.close()
./src\chongming_task\stores\__init__.py
"""持久层实现。
:mod:`~chongming_task.stores.protocol` 定义契约，本包提供：
* :class:`~chongming_task.stores.memory.InMemoryTaskStore` —— 单进程 / 测试 / 参照物；
* :class:`~chongming_task.stores.sqlite.SqliteTaskStore` —— 跨进程与持久化（标准库）；
* :class:`~chongming_task.stores.postgres.PostgresTaskStore` —— 跨机部署（可选依赖）。
**换持久层只需要做两件事**：
1. 实现 :class:`TaskStore` 契约里的方法（语义要求见 ``docs/store-guide.md``）；
2. 用 ``chongming_task.testing.store_contract.StoreContract`` 自证。
代码里用不到继承或注册——``Controller(store=...)`` 只做结构化检查，任何满足契约的
对象都能直接塞进去。序列化细节可以复用 :mod:`~chongming_task.stores.codec`，不必
自己重写一遍 ``datetime`` / ``frozenset`` / 枚举 / 嵌套结构的往返。
``PostgresTaskStore`` 需要可选的 psycopg，因此它**惰性导入**（既不在 ``__all__`` 里，
也不在包导入时求值）——装没装驱动都不影响 ``import chongming_task``。
"""
from __future__ import annotations
from typing import Any
from .codec import (
command_to_row,
decode_dt,
encode_dt,
row_to_command,
row_to_snapshot,
row_to_transition,
snapshot_to_row,
transition_to_row,
)
from .memory import InMemoryTaskStore
from .protocol import Command, CommandKind, TaskQuery, TaskStore
from .sqlite import SqliteTaskStore
__all__ = [
"Command",
"CommandKind",
"InMemoryTaskStore",
"SqliteTaskStore",
"TaskQuery",
"TaskStore",
"command_to_row",
"decode_dt",
"encode_dt",
"row_to_command",
"row_to_snapshot",
"row_to_transition",
"snapshot_to_row",
"transition_to_row",
]
def __getattr__(name: str) -> Any:
"""惰性导出可选依赖的实现。
见模块文档：``PostgresTaskStore`` 需要 psycopg，把它放进 ``__all__`` 会让
``from chongming_task.stores import *`` 在没装驱动时报错——那是核心包不该有的
依赖。懒加载保证"用到了才需要装"。
"""
if name == "PostgresTaskStore":
from .postgres import PostgresTaskStore
return PostgresTaskStore
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
./src\chongming_task\testing\clock.py
"""确定性时钟。"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
__all__ = ["FakeClock"]
class FakeClock:
"""可显式推进的时钟。
刻意**不**自动前进：测试必须显式调用 :meth:`advance`。这样"时间没有流逝"
与"忘记推进时间"不会混为一谈——后者会让超时测试静默地永远不触发。
"""
__slots__ = ("_now",)
def __init__(self, start: datetime | None = None) -> None:
self._now = start if start is not None else datetime(2026, 1, 1, tzinfo=UTC)
def __call__(self) -> datetime:
return self._now
@property
def now(self) -> datetime:
"""当前时刻。"""
return self._now
def advance(self, seconds: float) -> datetime:
"""向前推进 ``seconds`` 秒并返回新时刻。"""
if seconds < 0:
raise ValueError(f"时间不能倒流：{seconds}")
self._now += timedelta(seconds=seconds)
return self._now
def set(self, moment: datetime) -> None:
"""直接设为 ``moment``（只能向前）。"""
if moment < self._now:
raise ValueError(f"时间不能倒流：{moment} < {self._now}")
self._now = moment
./src\chongming_task\testing\store_contract.py
"""``TaskStore`` 契约测试套件。
**随包发布**，因为"可扩展的持久层"只有配上可执行的契约才有意义：第三方实现
（Postgres、Redis、对象存储……）只要能通过这一套用例，就可以放心替换掉内置实现。
刻意**不依赖 pytest**：只用普通断言与一个 :func:`expect` 助手，因此这个模块不会
给运行期引入任何第三方依赖。pytest 仍会自动收集子类里的 ``test_*`` 方法。
用法::
from chongming_task.testing.store_contract import StoreContract
class TestMyStore(StoreContract):
def make_store(self):
return MyStore(...)
"""
from __future__ import annotations
from dataclasses import replace
from datetime import UTC, datetime, timedelta
from typing import Any
from ..core.state_machine import Transition
from ..errors import ConcurrencyConflict, StoreError
from ..models import Failure, FailureKind, Progress, TaskSnapshot
from ..states import TaskStatus
from ..stores.protocol import Command, CommandKind, TaskQuery, TaskStore
__all__ = ["StoreContract"]
_T0 = datetime(2026, 1, 1, tzinfo=UTC)
def expect(exception: type[BaseException], action: Any, *, because: str = "") -> BaseException:
"""断言 ``action`` 抛出 ``exception``，并把它返回以便进一步检查。
不用 ``pytest.raises``，是为了让本模块保持零依赖。
"""
try:
action()
except exception as caught:
return caught
except BaseException as other:
raise AssertionError(
f"{because}：期望 {exception.__name__}，实际 {type(other).__name__}: {other}"
) from other
raise AssertionError(f"{because}：期望抛出 {exception.__name__}，但没有")
def full_snapshot(task_id: str = "t-full") -> TaskSnapshot:
"""一个**每个字段都填满**的快照。
契约测试用它做往返比对：只测几个字段会漏掉序列化里的静默丢失（例如时区、
frozenset、嵌套的 progress / failure）。
"""
return TaskSnapshot(
task_id=task_id,
name="抽取",
pipeline_run_id="run-1",
tags=frozenset({"io", "nightly"}),
bucket=7,
bucket_key="tenant-42",
status=TaskStatus.FAILED,
status_version=7,
requested_action=None,
attempt=2,
max_attempts=3,
created_at=_T0,
queued_at=_T0 + timedelta(seconds=1),
started_at=_T0 + timedelta(seconds=1),
ended_at=_T0 + timedelta(seconds=5),
last_heartbeat_at=_T0 + timedelta(seconds=4),
progress=Progress(
done=3,
total=10,
unit="rows",
message="半途",
updated_at=_T0 + timedelta(seconds=3),
),
failure=Failure(
kind=FailureKind.TIMEOUT_ATTEMPT,
message="超时了",
traceback="Traceback (most recent call last): ...",
exit_code=3,
retryable=True,
),
outcome={"rows": 3},
executor="process",
pid=4321,
overridden=False,
lease_owner="worker-1",
lease_until=_T0 + timedelta(seconds=30),
)
def simple_snapshot(
task_id: str,
*,
status: TaskStatus = TaskStatus.RUNNING,
version: int = 1,
created_offset: float = 0.0,
name: str = "任务",
tags: frozenset[str] = frozenset(),
bucket: int | None = None,
) -> TaskSnapshot:
"""一个便于构造查询场景的最小快照。"""
return TaskSnapshot(
task_id=task_id,
name=name,
tags=tags,
bucket=bucket,
bucket_key=None if bucket is None else f"key-{task_id}",
status=status,
status_version=version,
attempt=1,
max_attempts=1,
created_at=_T0 + timedelta(seconds=created_offset),
started_at=_T0 + timedelta(seconds=created_offset),
)
def transition(task_id: str, version: int, *, reason: str | None = None) -> Transition:
"""构造一条迁移记录。"""
return Transition(
task_id=task_id,
from_status=TaskStatus.QUEUED,
to_status=TaskStatus.RUNNING,
at=_T0 + timedelta(seconds=version),
trigger="executor:launched",
actor="system",
reason=reason,
version=version,
)
def command(
command_id: str,
task_id: str = "t1",
*,
kind: CommandKind = CommandKind.STOP,
created_offset: float = 0.0,
) -> Command:
"""构造一条控制指令。"""
return Command(
command_id=command_id,
task_id=task_id,
kind=kind,
created_at=_T0 + timedelta(seconds=created_offset),
requested_by="ops",
reason="测试",
)
class StoreContract:
"""每个 :class:`TaskStore` 实现都必须通过的行为契约。
子类只需实现 :meth:`make_store`，并保证每次调用返回**互相隔离**的干净存储。
"""
store: TaskStore
def make_store(self) -> TaskStore:
"""返回一个全新的、隔离的存储实例。"""
raise NotImplementedError
def setup_method(self) -> None:
self.store = self.make_store()
def teardown_method(self) -> None:
self.store.close()
# ------------------------------------------------------------------
# 契约的形式要件
# ------------------------------------------------------------------
def test_implements_the_protocol(self) -> None:
assert isinstance(self.store, TaskStore)
# ------------------------------------------------------------------
# 1. 任务快照
# ------------------------------------------------------------------
def test_full_snapshot_round_trips_without_loss(self) -> None:
"""每个字段都要原样回来：时区、frozenset、嵌套的 progress / failure。"""
original = full_snapshot()
self.store.create(original)
restored = self.store.get(original.task_id)
assert restored is not None
assert restored == original
assert restored.tags == frozenset({"io", "nightly"})
assert restored.progress is not None
assert restored.progress.message == "半途"
assert restored.failure is not None
assert restored.failure.kind is FailureKind.TIMEOUT_ATTEMPT
assert restored.failure.retryable is True
assert restored.outcome == {"rows": 3}
def test_get_returns_none_for_unknown_task(self) -> None:
assert self.store.get("nope") is None
def test_create_rejects_duplicate_task_id(self) -> None:
self.store.create(simple_snapshot("dup"))
expect(StoreError, lambda: self.store.create(simple_snapshot("dup")), because="重复创建")
def test_save_updates_when_version_matches(self) -> None:
self.store.create(simple_snapshot("t1", version=1))
updated = simple_snapshot("t1", status=TaskStatus.SUCCEEDED, version=2)
self.store.save(updated, expect_version=1)
restored = self.store.get("t1")
assert restored is not None
assert restored.status is TaskStatus.SUCCEEDED
assert restored.status_version == 2
def test_save_rejects_stale_version(self) -> None:
"""乐观并发是"换后端"能成立的前提：冲突必须被拒绝，而不是静默覆盖。"""
self.store.create(simple_snapshot("t1", version=1))
self.store.save(simple_snapshot("t1", version=2), expect_version=1)
conflict = expect(
ConcurrencyConflict,
lambda: self.store.save(simple_snapshot("t1", version=3), expect_version=1),
because="用过期的版本号写入",
)
assert isinstance(conflict, ConcurrencyConflict)
assert self.store.get("t1").status_version == 2 # type: ignore[union-attr]
def test_save_on_missing_task_raises(self) -> None:
expect(
StoreError,
lambda: self.store.save(simple_snapshot("ghost"), expect_version=1),
because="更新不存在的任务",
)
# ------------------------------------------------------------------
# 2. 查询
# ------------------------------------------------------------------
def test_query_filters_by_status(self) -> None:
self.store.create(simple_snapshot("a", status=TaskStatus.RUNNING))
self.store.create(simple_snapshot("b", status=TaskStatus.SUCCEEDED))
found = self.store.query(TaskQuery(statuses=frozenset({TaskStatus.RUNNING})))
assert [item.task_id for item in found] == ["a"]
def test_query_filters_by_tag_exactly(self) -> None:
"""标签必须精确匹配：``io`` 不能命中 ``ios``。
存储层常见的退化实现是子串匹配，它在标签里会悄悄错配。
"""
self.store.create(simple_snapshot("a", tags=frozenset({"io"})))
self.store.create(simple_snapshot("b", tags=frozenset({"ios"})))
self.store.create(simple_snapshot("c", tags=frozenset({"cpu", "io"})))
found = self.store.query(TaskQuery(tag="io"))
assert sorted(item.task_id for item in found) == ["a", "c"]
def test_query_filters_by_name_and_prefix(self) -> None:
self.store.create(simple_snapshot("a", name="抽取-甲"))
self.store.create(simple_snapshot("b", name="抽取-乙"))
self.store.create(simple_snapshot("c", name="加载"))
assert [i.task_id for i in self.store.query(TaskQuery(name="加载"))] == ["c"]
assert sorted(i.task_id for i in self.store.query(TaskQuery(name_prefix="抽取"))) == [
"a",
"b",
]
def test_query_orders_and_limits(self) -> None:
for index, task_id in enumerate(("a", "b", "c")):
self.store.create(simple_snapshot(task_id, created_offset=float(index)))
assert [i.task_id for i in self.store.query()] == ["a", "b", "c"]
assert [i.task_id for i in self.store.query(TaskQuery(newest_first=True))] == [
"c",
"b",
"a",
]
assert [i.task_id for i in self.store.query(TaskQuery(limit=2))] == ["a", "b"]
def test_query_filters_by_bucket_set(self) -> None:
"""分布式下"取我拥有的那组桶"会是一条高频查询，因此它必须能下推到存储层。"""
for index, (task_id, bucket) in enumerate((("a", 1), ("b", 2), ("c", 1))):
self.store.create(simple_snapshot(task_id, created_offset=float(index), bucket=bucket))
assert [i.task_id for i in self.store.query(TaskQuery(buckets=frozenset({1})))] == [
"a",
"c",
]
assert [i.task_id for i in self.store.query(TaskQuery(buckets=frozenset({1, 2})))] == [
"a",
"b",
"c",
]
assert self.store.query(TaskQuery(buckets=frozenset({9}))) == []
def test_task_without_bucket_key_has_no_bucket(self) -> None:
self.store.create(simple_snapshot("a"))
found = self.store.query()[0]
assert found.bucket is None
assert found.bucket_key is None
def test_query_on_empty_store(self) -> None:
assert self.store.query() == []
assert self.store.counts() == {}
def test_counts_by_status(self) -> None:
self.store.create(simple_snapshot("a", status=TaskStatus.RUNNING))
self.store.create(simple_snapshot("b", status=TaskStatus.RUNNING))
self.store.create(simple_snapshot("c", status=TaskStatus.FAILED))
counts = self.store.counts()
assert counts[TaskStatus.RUNNING] == 2
assert counts[TaskStatus.FAILED] == 1
# ------------------------------------------------------------------
# 3. 审计轨迹
# ------------------------------------------------------------------
def test_history_is_ordered_by_version(self) -> None:
self.store.create(simple_snapshot("t1"))
for version in (3, 1, 2):
self.store.append(transition("t1", version))
assert [item.version for item in self.store.history("t1")] == [1, 2, 3]
def test_history_filters_by_since_version(self) -> None:
self.store.create(simple_snapshot("t1"))
for version in (1, 2, 3):
self.store.append(transition("t1", version))
assert [i.version for i in self.store.history("t1", since_version=1)] == [2, 3]
def test_append_is_idempotent(self) -> None:
"""重放同一批记录（崩溃恢复）不应产生重复行。"""
self.store.create(simple_snapshot("t1"))
self.store.append(transition("t1", 1))
self.store.append(transition("t1", 1))
assert len(self.store.history("t1")) == 1
def test_history_round_trips_all_fields(self) -> None:
self.store.create(simple_snapshot("t1"))
original = transition("t1", 5, reason="发布窗口")
self.store.append(original)
restored = self.store.history("t1")[0]
assert restored == original
assert restored.reason == "发布窗口"
def test_history_of_unknown_task_is_empty(self) -> None:
assert self.store.history("nope") == []
# ------------------------------------------------------------------
# 4. 高频旁路
# ------------------------------------------------------------------
def test_touch_updates_progress_and_heartbeat_without_bumping_version(self) -> None:
"""进度与心跳若参与乐观并发，会与状态写入持续冲突。"""
self.store.create(simple_snapshot("t1", version=4))
moment = _T0 + timedelta(seconds=30)
self.store.touch(
"t1",
progress=Progress(done=5, total=10, updated_at=moment),
heartbeat_at=moment,
)
restored = self.store.get("t1")
assert restored is not None
assert restored.status_version == 4, "touch 不应改动版本号"
assert restored.last_heartbeat_at == moment
assert restored.progress is not None
assert restored.progress.done == 5
def test_touch_on_missing_task_is_ignored(self) -> None:
"""心跳晚到一步（任务已被清理）不该报错。"""
self.store.touch("ghost", heartbeat_at=_T0)
def test_touch_writes_the_supervisor_lease(self) -> None:
"""租约与心跳是两个独立信号，都能经 touch 落库。"""
self.store.create(simple_snapshot("t1"))
until = _T0 + timedelta(seconds=30)
self.store.touch("t1", lease_owner="worker-1", lease_until=until)
restored = self.store.get("t1")
assert restored is not None
assert restored.lease_owner == "worker-1"
assert restored.lease_until == until
def test_touch_can_set_owner_and_lease_independently(self) -> None:
"""只给 owner 不给 until（或反之）也应当被接受，不能要求调用方凑齐两个。"""
self.store.create(simple_snapshot("t1"))
self.store.touch("t1", lease_owner="worker-1")
restored = self.store.get("t1")
assert restored is not None
assert restored.lease_owner == "worker-1"
assert restored.lease_until is None
def test_save_does_not_clobber_the_lease(self) -> None:
"""**这条是重启恢复的关键。** ``save`` 不写租约列。
否则一次状态迁移就会把租约覆盖回内存快照里的旧值（甚至 ``None``），于是
重启恢复会把一个活得好好的任务误判为"失去监督者"并判 ``LOST``——把一个正常
任务杀掉，而且只在重启时才暴露。
"""
self.store.create(simple_snapshot("t1", version=1))
until = _T0 + timedelta(seconds=30)
self.store.touch("t1", lease_owner="worker-1", lease_until=until)
# 内存里的快照并不知道租约（它来自 touch），于是保存时会带一个空租约。
plain = simple_snapshot("t1", status=TaskStatus.SUCCEEDED, version=2)
assert plain.lease_owner is None
self.store.save(plain, expect_version=1)
restored = self.store.get("t1")
assert restored is not None
assert restored.status is TaskStatus.SUCCEEDED
assert restored.lease_owner == "worker-1", "save 覆盖了租约，重启恢复会误判失联"
assert restored.lease_until == until
# ------------------------------------------------------------------
# 5. 控制指令
# ------------------------------------------------------------------
def test_claim_then_complete(self) -> None:
self.store.create(simple_snapshot("t1"))
self.store.submit_command(command("c1"))
claimed = self.store.claim_commands(now=_T0, lease_seconds=30)
assert [item.command_id for item in claimed] == ["c1"]
assert claimed[0].lease_until is not None
self.store.complete_command("c1", now=_T0, outcome="ok")
assert self.store.claim_commands(now=_T0 + timedelta(seconds=60), lease_seconds=30) == []
def test_leased_command_is_not_handed_out_twice(self) -> None:
"""同一次停止被下发两次会让任务收到重复指令。"""
self.store.create(simple_snapshot("t1"))
self.store.submit_command(command("c1"))
assert len(self.store.claim_commands(now=_T0, lease_seconds=30)) == 1
assert self.store.claim_commands(now=_T0 + timedelta(seconds=1), lease_seconds=30) == []
def test_expired_lease_can_be_reclaimed(self) -> None:
"""持有者崩溃后，指令必须能被重新取走，否则会永久卡住。"""
self.store.create(simple_snapshot("t1"))
self.store.submit_command(command("c1"))
assert len(self.store.claim_commands(now=_T0, lease_seconds=10)) == 1
reclaimed = self.store.claim_commands(now=_T0 + timedelta(seconds=11), lease_seconds=10)
assert [item.command_id for item in reclaimed] == ["c1"]
def test_claims_respect_limit_and_creation_order(self) -> None:
self.store.create(simple_snapshot("t1"))
for index in range(3):
self.store.submit_command(command(f"c{index}", created_offset=float(index)))
claimed = self.store.claim_commands(now=_T0, lease_seconds=30, limit=2)
assert [item.command_id for item in claimed] == ["c0", "c1"]
def test_command_kind_and_metadata_survive(self) -> None:
self.store.create(simple_snapshot("t1"))
original = command("c1", kind=CommandKind.PAUSE)
self.store.submit_command(original)
claimed = self.store.claim_commands(now=_T0, lease_seconds=30)[0]
assert claimed.kind is CommandKind.PAUSE
assert claimed.task_id == "t1"
assert claimed.requested_by == "ops"
assert claimed.reason == "测试"
def test_released_command_becomes_claimable_immediately(self) -> None:
"""取到不属于自己的指令时应当立即交还。
否则要么被错误地标记完成（真正的持有者永远收不到这次停止请求），要么被
持续续租（真正的持有者被无限期饿死）。
"""
self.store.create(simple_snapshot("t1"))
self.store.submit_command(command("c1"))
assert len(self.store.claim_commands(now=_T0, lease_seconds=300)) == 1
self.store.release_command("c1", now=_T0)
reclaimed = self.store.claim_commands(now=_T0, lease_seconds=30)
assert [item.command_id for item in reclaimed] == ["c1"]
def test_released_command_is_still_open(self) -> None:
self.store.create(simple_snapshot("t1"))
self.store.submit_command(command("c1"))
self.store.claim_commands(now=_T0, lease_seconds=300)
self.store.release_command("c1", now=_T0)
claimed = self.store.claim_commands(now=_T0, lease_seconds=300)
assert claimed[0].is_open
def test_release_does_not_revive_a_finished_command(self) -> None:
self.store.create(simple_snapshot("t1"))
self.store.submit_command(command("c1"))
self.store.claim_commands(now=_T0, lease_seconds=30)
self.store.complete_command("c1", now=_T0, outcome="ok")
self.store.release_command("c1", now=_T0)
assert self.store.claim_commands(now=_T0, lease_seconds=30) == []
def test_query_filters_by_pipeline_run(self) -> None:
self.store.create(replace(simple_snapshot("t1"), pipeline_run_id="run-a"))
self.store.create(replace(simple_snapshot("t2"), pipeline_run_id="run-a"))
self.store.create(replace(simple_snapshot("t3"), pipeline_run_id="run-b"))
self.store.create(simple_snapshot("t4"))
found = self.store.query(TaskQuery(pipeline_run_id="run-a"))
assert [item.task_id for item in found] == ["t1", "t2"]
assert self.store.query(TaskQuery(pipeline_run_id="nope")) == []
def test_duplicate_command_id_is_rejected(self) -> None:
self.store.create(simple_snapshot("t1"))
self.store.submit_command(command("c1"))
expect(StoreError, lambda: self.store.submit_command(command("c1")), because="重复指令")
# ------------------------------------------------------------------
# 6. 运维
# ------------------------------------------------------------------
def test_prune_keeps_only_the_newest_terminal_tasks(self) -> None:
for index, task_id in enumerate(("a", "b", "c")):
self.store.create(
simple_snapshot(
task_id,
status=TaskStatus.SUCCEEDED,
created_offset=float(index),
)
)
self.store.append(transition(task_id, 1))
assert self.store.prune(keep_terminal=2) == 1
assert [item.task_id for item in self.store.query()] == ["b", "c"]
assert self.store.history("a") == [], "被清理任务的历史也应一并删除"
assert self.store.prune(keep_terminal=2) == 0
def test_prune_never_touches_unfinished_tasks(self) -> None:
"""运行中的任务被清掉就等于把状态机扔了。"""
self.store.create(simple_snapshot("running", status=TaskStatus.RUNNING))
self.store.create(simple_snapshot("done", status=TaskStatus.SUCCEEDED))
assert self.store.prune(keep_terminal=0) == 1
assert self.store.get("running") is not None
assert self.store.get("done") is None
./src\chongming_task\testing\__init__.py
"""随包发布的测试工具。
放在包内而非 ``tests/`` 之下，是为了让**使用者的测试**也能用上确定性时钟与
断言助手——否则每个使用方都要自己重写一遍 :class:`FakeClock`，而重写出来的
版本往往无法与运行时的取时点完全对齐。
"""
from __future__ import annotations
from .clock import FakeClock
__all__ = ["FakeClock"]
./tests\contract\test_store_contract.py
"""内置存储实现的契约一致性测试。
两个实现（内存 / SQLite）跑**完全相同**的用例。这是"持久层可替换"这句话能成立的
唯一证据——否则换后端就是在赌。
"""
from __future__ import annotations
import tempfile
from pathlib import Path
from chongming_task.stores.memory import InMemoryTaskStore
from chongming_task.stores.protocol import TaskStore
from chongming_task.stores.sqlite import SqliteTaskStore
from chongming_task.testing.store_contract import StoreContract
class TestInMemoryStore(StoreContract):
"""内存实现。"""
def make_store(self) -> TaskStore:
return InMemoryTaskStore()
class TestSqliteFileStore(StoreContract):
"""SQLite 文件库——真实使用场景（跨进程、可重启）。"""
_tempdir: tempfile.TemporaryDirectory[str]
def make_store(self) -> TaskStore:
self._tempdir = tempfile.TemporaryDirectory(prefix="chongming-store-")
return SqliteTaskStore(Path(self._tempdir.name) / "tasks.db")
class TestSqliteMemoryStore(StoreContract):
"""SQLite 内存库——语义必须与文件库一致。"""
def make_store(self) -> TaskStore:
return SqliteTaskStore(":memory:")
./tests\integration\test_cli.py
"""命令行测试。
CLI 是**另一个进程**，所以这里全部通过 ``--db`` 指向真实的 SQLite 文件，而不是把
``Controller`` 塞进去——那样测的就不是 CLI 了。
重点覆盖两件容易静默出错的事：
* **输出契约**（``--json`` 的形状、退出码），因为脚本要依赖它；
* **错误路径**，尤其是"给一个不存在的任务下指令"——它曾经会成功并留下一条永远不会
被执行的指令。
"""
from __future__ import annotations
import json
import os
from datetime import UTC, datetime
from pathlib import Path
import pytest
from chongming_task.cli.main import (
EXIT_ERROR,
EXIT_OK,
EXIT_TASK_NOT_SUCCESSFUL,
build_parser,
main,
)
from chongming_task.errors import ConfigurationError
from chongming_task.models import TaskSnapshot
from chongming_task.states import TaskStatus
from chongming_task.stores.sqlite import SqliteTaskStore
_TASKS_MODULE = "tasks_demo"
@pytest.fixture
def db(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> str:
"""把示例任务模块加进导入路径，并返回一个全新的库路径。"""
import sys
examples = Path(__file__).resolve().parents[2] / "examples"
monkeypatch.syspath_prepend(str(examples))
sys.modules.pop(_TASKS_MODULE, None)
return str(tmp_path / "tasks.db")
def _run_cli(db: str, *argv: str) -> int:
return main(["--db", db, *argv])
# ---------------------------------------------------------------------------
# 参数解析
# ---------------------------------------------------------------------------
def test_parser_requires_a_subcommand() -> None:
with pytest.raises(SystemExit):
build_parser().parse_args([])
def test_parser_help_mentions_every_command() -> None:
text = build_parser().format_help()
for command in (
"tasks",
"run",
"ps",
"status",
"history",
"stats",
"buckets",
"stop",
"pause",
"resume",
):
assert command in text
# ---------------------------------------------------------------------------
# 定义
# ---------------------------------------------------------------------------
def test_tasks_lists_the_registered_definitions(
db: str, capsys: pytest.CaptureFixture[str]
) -> None:
assert main(["tasks", "--tasks", _TASKS_MODULE]) == EXIT_OK
out = capsys.readouterr().out
assert "sync" in out
assert "add" in out
assert "两数相加。" in out, "注册项的 docstring 应当作为描述显示"
def test_tasks_on_an_unimportable_module_fails_clearly(
capsys: pytest.CaptureFixture[str],
) -> None:
assert main(["tasks", "--tasks", "no.such.module"]) == EXIT_ERROR
assert "无法导入任务模块" in capsys.readouterr().err
def test_a_module_not_found_error_explains_sys_path(
monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
"""模块自身没找到时，必须说清"为什么找不到"。
最常踩的坑：``chongming-task`` 是**入口脚本**，不会把当前目录放进 ``sys.path``——
把模块放在当前目录、然后直接 ``--tasks mymodule`` 会失败，而 "No module named"
本身完全不解释这件事。
"""
monkeypatch.delenv("PYTHONPATH", raising=False)
assert main(["tasks", "--tasks", "definitely_not_installed"]) == EXIT_ERROR
err = capsys.readouterr().err
assert "PYTHONPATH" in err
assert "不会**把当前目录放进 sys.path" in err
assert "python -m chongming_task.cli.main" in err
def test_a_module_found_but_failing_to_import_says_so(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
"""模块存在、但内部导入失败：不能把矛头指向 sys.path，那是它自己的依赖问题。"""
module = tmp_path / "broken_tasks.py"
module.write_text("import no_such_dependency_here\n", encoding="utf-8")
monkeypatch.syspath_prepend(str(tmp_path))
assert main(["tasks", "--tasks", "broken_tasks"]) == EXIT_ERROR
err = capsys.readouterr().err
assert "找到了，但它导入失败" in err
assert "PYTHONPATH" not in err
def test_the_cli_module_can_be_run_with_python_m() -> None:
"""备选入口：``python -m`` 会把当前目录放进 ``sys.path``，因此根目录的模块也能用。
这条用例钉住"备选入口存在"这个事实——它已经写进了错误提示里，不能悄悄消失。
"""
source = Path(__file__).resolve().parents[2] / "src" / "chongming_task" / "cli" / "main.py"
assert source.is_file()
assert 'if __name__ == "__main__"' in source.read_text(encoding="utf-8")
def test_tasks_on_a_module_without_register_fails_clearly(
capsys: pytest.CaptureFixture[str],
) -> None:
"""约定必须说清楚——否则用户只会看到一句 AttributeError。"""
assert main(["tasks", "--tasks", "json"]) == EXIT_ERROR
assert "register(controller)" in capsys.readouterr().err
# ---------------------------------------------------------------------------
# 提交与执行
# ---------------------------------------------------------------------------
def test_run_can_execute_in_a_child_process(db: str, capsys: pytest.CaptureFixture[str]) -> None:
"""``--executor process``：任务在子进程里跑，快照因此带上 pid。
这是"能被真的杀掉"那条能力在命令行上的入口——前台 ``run`` 遇到不配合的任务时，
另一个终端发 ``stop`` 就能把它连根拔掉（线程只能放弃等待）。
"""
code = _run_cli(
db,
"run",
"--tasks",
_TASKS_MODULE,
"add",
"--arg",
"left=2",
"--arg",
"right=3",
"--executor",
"process",
)
assert code == EXIT_OK
assert "succeeded" in capsys.readouterr().out
store = SqliteTaskStore(db)
try:
snapshot = store.query()[0]
assert snapshot.executor == "process"
assert snapshot.pid is not None
assert snapshot.pid != os.getpid(), "任务确实跑在另一个进程里"
finally:
store.close()
def test_run_executes_and_exits_zero(db: str, capsys: pytest.CaptureFixture[str]) -> None:
code = _run_cli(
db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=2", "--arg", "right=3"
)
assert code == EXIT_OK
out = capsys.readouterr().out
assert "succeeded" in out
store = SqliteTaskStore(db)
try:
snapshots = store.query()
assert len(snapshots) == 1
assert snapshots[0].outcome == {"value": 5}
finally:
store.close()
def test_run_parses_json_typed_arguments(db: str) -> None:
"""``--arg n=3`` 应当得到整数 3，而不是字符串 "3"。"""
assert (
_run_cli(
db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=20", "--arg", "right=22"
)
== EXIT_OK
)
store = SqliteTaskStore(db)
try:
assert store.query()[0].outcome == {"value": 42}
finally:
store.close()
def test_run_on_a_failing_task_exits_three(db: str) -> None:
"""命令成功、任务失败——两者用不同退出码区分，脚本才好判分支。"""
code = _run_cli(db, "run", "--tasks", _TASKS_MODULE, "flaky", "--arg", "fail_times=9")
assert code == EXIT_TASK_NOT_SUCCESSFUL
def test_run_rejects_a_malformed_arg(db: str, capsys: pytest.CaptureFixture[str]) -> None:
code = _run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "noequals")
assert code == EXIT_ERROR
assert "K=V" in capsys.readouterr().err
def test_run_records_the_bucket(db: str) -> None:
code = _run_cli(
db,
"run",
"--tasks",
_TASKS_MODULE,
"add",
"--arg",
"left=1",
"--arg",
"right=1",
"--bucket-key",
"tenant-9",
)
assert code == EXIT_OK
store = SqliteTaskStore(db)
try:
snapshot = store.query()[0]
assert snapshot.bucket_key == "tenant-9"
assert snapshot.bucket is not None
finally:
store.close()
# ---------------------------------------------------------------------------
# 查询
# ---------------------------------------------------------------------------
def test_ps_defaults_to_unfinished_and_all_shows_everything(
db: str, capsys: pytest.CaptureFixture[str]
) -> None:
assert (
_run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=1", "--arg", "right=1")
== EXIT_OK
)
capsys.readouterr()
assert _run_cli(db, "ps") == EXIT_OK
assert "没有匹配的任务" in capsys.readouterr().out, "默认只显示未落终态的"
assert _run_cli(db, "ps", "--all") == EXIT_OK
assert "succeeded" in capsys.readouterr().out
def test_status_shows_detail_and_history(db: str, capsys: pytest.CaptureFixture[str]) -> None:
_run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=1", "--arg", "right=1")
capsys.readouterr()
assert _run_cli(db, "status", "x") == EXIT_ERROR # 不存在
store = SqliteTaskStore(db)
try:
task_id = store.query()[0].task_id
finally:
store.close()
assert _run_cli(db, "status", task_id) == EXIT_OK
out = capsys.readouterr().out
assert "status" in out
assert "迁移历史" in out
assert "executor:completed" in out
def test_stats_and_buckets(db: str, capsys: pytest.CaptureFixture[str]) -> None:
_run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=1", "--arg", "right=1")
capsys.readouterr()
assert _run_cli(db, "stats") == EXIT_OK
out = capsys.readouterr().out
assert "succeeded" in out
assert "total" in out
assert _run_cli(db, "buckets") == EXIT_OK
assert "BUCKET" in capsys.readouterr().out
def test_ps_rejects_an_unknown_status(db: str, capsys: pytest.CaptureFixture[str]) -> None:
assert _run_cli(db, "ps", "--status", "no-such-status") == EXIT_ERROR
assert "未知状态" in capsys.readouterr().err
def test_ps_filters_by_pipeline_run(db: str, capsys: pytest.CaptureFixture[str]) -> None:
"""按管道查看是看板最常问的问题，过滤下推到存储层。"""
from datetime import UTC, datetime
from chongming_task.models import TaskSnapshot
store = SqliteTaskStore(db)
try:
for task_id, run_id in (("a1", "run-a"), ("a2", "run-a"), ("b1", "run-b")):
store.create(
TaskSnapshot(
task_id=task_id,
name="节点",
status=TaskStatus.SUCCEEDED,
status_version=1,
attempt=1,
max_attempts=1,
created_at=datetime(2026, 1, 1, tzinfo=UTC),
ended_at=datetime(2026, 1, 1, 1, tzinfo=UTC),
pipeline_run_id=run_id,
)
)
finally:
store.close()
assert _run_cli(db, "ps", "--all", "--pipeline", "run-a") == EXIT_OK
out = capsys.readouterr().out
assert "a1" in out
assert "a2" in out
assert "b1" not in out
assert _run_cli(db, "ps", "--all", "--pipeline", "no-such-run") == EXIT_OK
assert "没有匹配的任务" in capsys.readouterr().out
# ---------------------------------------------------------------------------
# 输出契约
# ---------------------------------------------------------------------------
def test_json_output_is_primitive_only(db: str, capsys: pytest.CaptureFixture[str]) -> None:
"""``--json`` 的输出必须全是原语，脚本才不必了解 datetime / 枚举的序列化规则。"""
_run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=2", "--arg", "right=3")
capsys.readouterr()
assert main(["--db", db, "--json", "ps", "--all"]) == EXIT_OK
payload = json.loads(capsys.readouterr().out)
assert isinstance(payload, list)
entry = payload[0]
assert entry["status"] == "succeeded"
assert entry["outcome"] == '{"value": 5}'
assert all(
value is None or isinstance(value, (str, int, float, bool)) for value in entry.values()
), "出现了非原语字段"
def test_json_error_path_uses_stderr_not_stdout(capsys: pytest.CaptureFixture[str]) -> None:
assert main(["--json", "ps"]) == EXIT_ERROR
captured = capsys.readouterr()
assert captured.out == ""
assert "未指定任务存储" in captured.err
# ---------------------------------------------------------------------------
# 错误路径
# ---------------------------------------------------------------------------
def test_missing_db_is_reported_clearly(
monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
monkeypatch.delenv("CHONGMING_TASK_DB", raising=False)
assert main(["ps"]) == EXIT_ERROR
assert "CHONGMING_TASK_DB" in capsys.readouterr().err
def test_db_can_come_from_the_environment(db: str, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("CHONGMING_TASK_DB", db)
_run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=1", "--arg", "right=1")
assert main(["ps", "--all"]) == EXIT_OK
def test_stop_on_an_unknown_task_fails_instead_of_writing_a_dead_command(
db: str, capsys: pytest.CaptureFixture[str]
) -> None:
"""曾经这里的 bug：给不存在的任务下指令会成功，留下一条永远不会被执行的指令，
而调用方还以为"停止已下发"。
"""
_run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=1", "--arg", "right=1")
capsys.readouterr()
assert _run_cli(db, "stop", "typo-task-id") == EXIT_ERROR
assert "没有找到任务" in capsys.readouterr().err
store = SqliteTaskStore(db)
try:
assert store.commands() == [], "不该留下任何指令"
finally:
store.close()
def test_control_on_an_existing_task_is_accepted(
db: str, capsys: pytest.CaptureFixture[str]
) -> None:
_run_cli(db, "run", "--tasks", _TASKS_MODULE, "add", "--arg", "left=1", "--arg", "right=1")
capsys.readouterr()
store = SqliteTaskStore(db)
try:
task_id = store.query()[0].task_id
finally:
store.close()
# 任务已落终态，但"找到并受理"仍然成立。
assert _run_cli(db, "stop", task_id, "--reason", "收尾") == EXIT_OK
assert "已请求停止" in capsys.readouterr().out
def test_store_directory_missing_is_reported_clearly(
tmp_path: Path, capsys: pytest.CaptureFixture[str]
) -> None:
"""SQLite 只会报 "unable to open database file"，完全看不出是目录不存在。"""
missing = tmp_path / "not-created" / "tasks.db"
assert main(["--db", str(missing), "ps"]) == EXIT_ERROR
assert "目录不存在" in capsys.readouterr().err
# ---------------------------------------------------------------------------
# 重启恢复
# ---------------------------------------------------------------------------
def _plant_zombie(db: str, task_id: str = "zombie") -> None:
"""造一个"上次崩溃留下的"任务行：running 且租约早已过期。"""
from datetime import UTC, datetime, timedelta
from chongming_task.models import TaskSnapshot
store = SqliteTaskStore(db)
try:
old = datetime.now(UTC) - timedelta(minutes=10)
store.create(
TaskSnapshot(
task_id=task_id,
name="遗留任务",
status=TaskStatus.RUNNING,
status_version=3,
attempt=1,
max_attempts=1,
created_at=old,
started_at=old,
lease_owner="worker-old",
lease_until=datetime.now(UTC) - timedelta(minutes=9),
)
)
finally:
store.close()
def test_reconcile_dry_run_writes_nothing(db: str, capsys: pytest.CaptureFixture[str]) -> None:
_plant_zombie(db)
assert _run_cli(db, "reconcile", "--dry-run") == EXIT_OK
out = capsys.readouterr().out
assert "marked-lost" in out
assert "预演" in out
store = SqliteTaskStore(db)
try:
snapshot = store.get("zombie")
assert snapshot is not None
assert snapshot.status is TaskStatus.RUNNING, "预演不该写状态"
assert store.history("zombie") == []
finally:
store.close()
def test_reconcile_marks_zombies(db: str, capsys: pytest.CaptureFixture[str]) -> None:
_plant_zombie(db)
assert _run_cli(db, "reconcile", "--reason", "worker-old 崩溃") == EXIT_OK
assert "已把 1 个僵尸任务标记为 lost" in capsys.readouterr().out
store = SqliteTaskStore(db)
try:
snapshot = store.get("zombie")
assert snapshot is not None
assert snapshot.status is TaskStatus.LOST
assert store.history("zombie")[0].reason == "worker-old 崩溃"
finally:
store.close()
def test_reconcile_with_nothing_to_do(db: str, capsys: pytest.CaptureFixture[str]) -> None:
assert _run_cli(db, "reconcile") == EXIT_OK
assert "没有需要处理的任务" in capsys.readouterr().out
def test_reconcile_json_report(db: str, capsys: pytest.CaptureFixture[str]) -> None:
_plant_zombie(db)
assert main(["--db", db, "--json", "reconcile", "--dry-run"]) == EXIT_OK
payload = json.loads(capsys.readouterr().out)
assert payload["dry_run"] is True
assert payload["entries"][0]["task_id"] == "zombie"
assert payload["entries"][0]["previous_status"] == "running"
assert payload["entries"][0]["outcome"] == "marked-lost"
# ---------------------------------------------------------------------------
# 节点与桶归属（自治节点 + 统一平面）
# ---------------------------------------------------------------------------
def test_bucket_spec_supports_ranges_and_lists() -> None:
from chongming_task.cli.main import _parse_bucket_spec
assert _parse_bucket_spec("0-3") == frozenset({0, 1, 2, 3})
assert _parse_bucket_spec("0-1,5,7-8") == frozenset({0, 1, 5, 7, 8})
assert _parse_bucket_spec(" 2 , 4 ") == frozenset({2, 4})
assert _parse_bucket_spec("3-3") == frozenset({3})
def test_bucket_spec_rejects_bad_input() -> None:
from chongming_task.cli.main import _parse_bucket_spec
for bad in ("3-1", "a", "0-", "1,x"):
with pytest.raises(ConfigurationError):
_parse_bucket_spec(bad)
with pytest.raises(ConfigurationError, match="没有解析出任何桶"):
_parse_bucket_spec(" , ")
def test_ps_mine_filters_to_owned_buckets(db: str, capsys: pytest.CaptureFixture[str]) -> None:
"""自治架构里最常看的视图："我这一份现在怎样"。"""
_plant_zombie(db, "zombie") # 无桶任务
store = SqliteTaskStore(db)
try:
for task_id, bucket in (("in-a", 1), ("in-b", 40)):
store.create(
TaskSnapshot(
task_id=task_id,
name="节点任务",
bucket=bucket,
bucket_key=f"key-{bucket}",
status=TaskStatus.SUCCEEDED,
status_version=1,
attempt=1,
max_attempts=1,
created_at=datetime(2026, 1, 1, tzinfo=UTC),
ended_at=datetime(2026, 1, 1, 1, tzinfo=UTC),
)
)
finally:
store.close()
assert main(["--db", db, "--owned-buckets", "0-31", "ps", "--all", "--mine"]) == EXIT_OK
out = capsys.readouterr().out
assert "in-a" in out
assert "in-b" not in out, "别把别的桶的任务算进来"
assert "zombie" not in out, "无桶任务不属于任何限定了归属的节点"
def test_ps_mine_requires_owned_buckets(db: str, capsys: pytest.CaptureFixture[str]) -> None:
assert main(["--db", db, "ps", "--mine"]) == EXIT_ERROR
assert "--owned-buckets" in capsys.readouterr().err
def test_reconcile_is_scoped_to_owned_buckets(db: str, capsys: pytest.CaptureFixture[str]) -> None:
"""节点只清理自己桶里的僵尸——这是自治架构的核心不变量。"""
store = SqliteTaskStore(db)
try:
for task_id, bucket in (("mine", 1), ("theirs", 40)):
old = datetime(2026, 1, 1, tzinfo=UTC)
store.create(
TaskSnapshot(
task_id=task_id,
name="遗留任务",
bucket=bucket,
bucket_key=f"key-{bucket}",
status=TaskStatus.RUNNING,
status_version=3,
attempt=1,
max_attempts=1,
created_at=old,
started_at=old,
lease_owner="dead",
lease_until=old,
)
)
finally:
store.close()
assert main(["--db", db, "--owned-buckets", "0-31", "reconcile"]) == EXIT_OK
out = capsys.readouterr().out
assert "mine" in out
assert "theirs" not in out
store = SqliteTaskStore(db)
try:
assert store.get("mine").status is TaskStatus.LOST # type: ignore[union-attr]
assert store.get("theirs").status is TaskStatus.RUNNING, ( # type: ignore[union-attr]
"动了别的节点的任务"
)
finally:
store.close()
def test_buckets_view_marks_owned_buckets(db: str, capsys: pytest.CaptureFixture[str]) -> None:
assert main(["--db", db, "--num-buckets", "4", "--owned-buckets", "0-1", "buckets"]) == EXIT_OK
out = capsys.readouterr().out
marked = [line for line in out.splitlines() if line.startswith("*")]
assert len(marked) == 2, f"应当恰好标出两个桶，实际：{marked}"
def test_an_out_of_range_owned_bucket_is_reported(
db: str, capsys: pytest.CaptureFixture[str]
) -> None:
"""桶号越界必须在启动时报错——否则症状是"任务永远排队且没有任何报错"。"""
assert main(["--db", db, "--owned-buckets", "0-999", "ps"]) == EXIT_ERROR
assert "桶号越界" in capsys.readouterr().err
def test_run_respects_owned_buckets(db: str, capsys: pytest.CaptureFixture[str]) -> None:
"""``run`` 必须**真的**受归属约束。
曾经 ``run`` 自己构造 ``Controller(...)``，把 ``--owned-buckets`` 静默忽略了——
参数被接受却不起作用，比报错更难查。
"""
from chongming_task import BucketPolicy
policy = BucketPolicy(num_buckets=64)
mine = next(f"tenant-{i}" for i in range(999) if policy.bucket_for(f"tenant-{i}") < 16)
theirs = next(f"tenant-{i}" for i in range(999) if policy.bucket_for(f"tenant-{i}") >= 48)
# 限定了归属却没给 bucket_key → 必须报错（否则会产出没人认领的任务）
assert (
main(
[
"--db",
db,
"--num-buckets",
"64",
"--owned-buckets",
"0-15",
"run",
"--tasks",
_TASKS_MODULE,
"add",
"--arg",
"left=1",
"--arg",
"right=1",
"--task-id",
"r1",
]
)
== EXIT_ERROR
)
assert "必须给出 bucket_key" in capsys.readouterr().err
# 落在别人的桶 → 必须报错
assert (
main(
[
"--db",
db,
"--num-buckets",
"64",
"--owned-buckets",
"0-15",
"run",
"--tasks",
_TASKS_MODULE,
"add",
"--arg",
"left=1",
"--arg",
"right=1",
"--task-id",
"r2",
"--bucket-key",
theirs,
]
)
== EXIT_ERROR
)
assert "不属于本节点" in capsys.readouterr().err
# 落在自己的桶 → 正常执行
assert (
main(
[
"--db",
db,
"--num-buckets",
"64",
"--owned-buckets",
"0-15",
"run",
"--tasks",
_TASKS_MODULE,
"add",
"--arg",
"left=20",
"--arg",
"right=22",
"--task-id",
"r3",
"--bucket-key",
mine,
]
)
== EXIT_OK
)
assert "succeeded" in capsys.readouterr().out
def test_a_module_with_only_task_decorators_works(
db: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
"""只用 ``@task`` 装饰、**没有** ``register()`` 的模块也必须能用。
``@task`` 是后来加的写法；CLI 若只认 ``register()``，那个 API 在命令行下就废了。
"""
module = tmp_path / "decorated_only.py"
module.write_text(
"from chongming_task import task\n"
"\n"
"@task(tags={'io'})\n"
"def doubled(ctx, value: int) -> int:\n"
" '''把数字翻倍。'''\n"
" del ctx\n"
" return value * 2\n",
encoding="utf-8",
)
monkeypatch.syspath_prepend(str(tmp_path))
assert main(["tasks", "--tasks", "decorated_only"]) == EXIT_OK
out = capsys.readouterr().out
assert "doubled" in out
assert "把数字翻倍。" in out
assert (
main(["--db", db, "run", "--tasks", "decorated_only", "doubled", "--arg", "value=21"])
== EXIT_OK
)
assert "succeeded" in capsys.readouterr().out
def test_a_module_with_neither_register_nor_task_is_reported(
capsys: pytest.CaptureFixture[str],
) -> None:
assert main(["tasks", "--tasks", "json"]) == EXIT_ERROR
assert "既没有 register" in capsys.readouterr().err
./tests\integration\test_custom_store.py
"""换掉默认持久层的**可执行范例**：一个 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 上"
./tests\integration\test_node_buckets.py
"""节点与桶归属测试：**自治节点 + 统一平面**架构的划分单位。
这个架构里，每个节点有自己的任务定义与触发方式，共用同一套存储与控制层，
靠 ``owned_buckets`` 区分"哪一份是我的"。所以这一组用例的重点是两条：
1. **划界**——节点只接受、只清理自己负责的桶；
2. **不许静默**——任务落进别人的桶时必须在提交那一刻报错，而不是产出一个人人都不认领
的任务（它会永远停在 ``queued``，且没有任何报错）。
"""
from __future__ import annotations
from datetime import UTC, datetime
from typing import Any
import pytest
from chongming_task.controller import Controller, ReconcileOutcome
from chongming_task.errors import ConfigurationError
from chongming_task.models import TaskSnapshot
from chongming_task.policies import BucketPolicy
from chongming_task.states import TaskStatus
from chongming_task.stores.memory import InMemoryTaskStore
_POLICY = BucketPolicy()
_OLD = datetime(2026, 1, 1, tzinfo=UTC)
#: 两个节点各自的桶（按业务划分）。
_NODE_A = frozenset(range(0, 32))
_NODE_B = frozenset(range(32, 64))
def _key_for(policy: BucketPolicy, wanted: frozenset[int]) -> str:
"""找一个落在指定桶集合里的业务键。"""
for index in range(5000):
key = f"tenant-{index}"
if policy.bucket_for(key) in wanted:
return key
raise AssertionError("找不到落在目标桶集合里的键")
# ---------------------------------------------------------------------------
# 归属的表达与校验
# ---------------------------------------------------------------------------
def test_without_owned_buckets_there_is_no_constraint() -> None:
"""单机（或负责全部桶的节点）不该被迫给每个任务写 bucket_key。"""
controller = Controller()
controller.register("noop", _noop)
assert controller.owned_buckets is None
assert controller.owns_all_buckets is True
assert controller.owns_bucket(None) is True
assert controller.owns_bucket(63) is True
run = controller.submit("noop")
assert run.snapshot.bucket is None
controller.wait_all(timeout=10)
controller.close(timeout=5)
def test_owned_buckets_are_recorded_as_a_set() -> None:
controller = Controller(owned_buckets=[1, 2, 2, 3])
assert controller.owned_buckets == frozenset({1, 2, 3})
assert controller.owns_all_buckets is False
assert controller.owns_bucket(2) is True
assert controller.owns_bucket(4) is False
assert controller.owns_bucket(None) is False
controller.close(timeout=5)
def test_an_empty_bucket_set_is_rejected() -> None:
"""空集通常是配置写漏了——那意味着这个节点什么都不负责。"""
with pytest.raises(ConfigurationError, match="不能是空集"):
Controller(owned_buckets=[])
def test_an_out_of_range_bucket_is_rejected_at_startup() -> None:
"""越界桶号必须在启动时报错。
放过去的症状是"某些任务永远停在 queued，且没有任何报错"——最难查的一类问题。
"""
limit = _POLICY.num_buckets
with pytest.raises(ConfigurationError, match="桶号越界"):
Controller(bucket_policy=_POLICY, owned_buckets=[limit])
with pytest.raises(ConfigurationError, match="桶号越界"):
Controller(bucket_policy=_POLICY, owned_buckets=[-1])
def test_a_non_integer_bucket_is_rejected() -> None:
with pytest.raises(ConfigurationError, match="必须是整数"):
Controller(owned_buckets=[1, "2"]) # type: ignore[list-item]
# ---------------------------------------------------------------------------
# 提交时划界：不许产出"没人认领"的任务
# ---------------------------------------------------------------------------
def test_a_task_without_a_bucket_is_rejected_when_buckets_are_limited() -> None:
controller = Controller(owned_buckets=_NODE_A)
controller.register("noop", _noop)
with pytest.raises(ConfigurationError, match="必须给出 bucket_key"):
controller.submit("noop")
controller.close(timeout=5)
def test_a_task_in_a_foreign_bucket_is_rejected() -> None:
controller = Controller(bucket_policy=_POLICY, owned_buckets=_NODE_A)
controller.register("noop", _noop)
foreign = _key_for(_POLICY, _NODE_B)
with pytest.raises(ConfigurationError, match="不属于本节点"):
controller.submit("noop", bucket_key=foreign)
controller.close(timeout=5)
def test_a_task_in_an_owned_bucket_is_accepted() -> None:
controller = Controller(bucket_policy=_POLICY, owned_buckets=_NODE_A)
controller.register("noop", _noop)
mine = _key_for(_POLICY, _NODE_A)
run = controller.submit("noop", bucket_key=mine)
assert run.snapshot.bucket in _NODE_A
controller.wait_all(timeout=10)
controller.close(timeout=5)
def test_a_node_owning_every_bucket_accepts_unbucketed_tasks() -> None:
"""常见形态：一个节点显式列举了全部桶（生成出来的配置）。
此时没有任何"别的节点"，无桶任务同样会且只会由它执行——再强制每个任务写
``bucket_key`` 就只是纯摩擦。
"""
policy = BucketPolicy(num_buckets=8)
controller = Controller(bucket_policy=policy, owned_buckets=range(8))
controller.register("noop", _noop)
assert controller.owns_all_buckets is True, "覆盖全部桶 = 不构成划分"
assert controller.owned_buckets == frozenset(range(8)), "但归属集合仍然记着"
run = controller.submit("noop")
assert run.snapshot.bucket is None
controller.wait_all(timeout=10)
controller.close(timeout=5)
def test_owning_one_bucket_less_is_a_real_partition() -> None:
"""差一个桶就构成划分，此时必须约束——边界不能含糊。"""
policy = BucketPolicy(num_buckets=8)
controller = Controller(bucket_policy=policy, owned_buckets=range(7))
controller.register("noop", _noop)
assert controller.owns_all_buckets is False
with pytest.raises(ConfigurationError, match="必须给出 bucket_key"):
controller.submit("noop")
controller.close(timeout=5)
# ---------------------------------------------------------------------------
# reconcile 的作用范围（本组最重要）
# ---------------------------------------------------------------------------
def _plant_zombie(store: InMemoryTaskStore, task_id: str, bucket: int) -> None:
"""种一个"上次崩溃留下的"任务：running 且租约早已过期。"""
store.create(
TaskSnapshot(
task_id=task_id,
name="遗留任务",
bucket=bucket,
bucket_key=f"key-{bucket}",
status=TaskStatus.RUNNING,
status_version=3,
attempt=1,
max_attempts=1,
created_at=_OLD,
started_at=_OLD,
lease_owner="dead-worker",
lease_until=_OLD,
)
)
def test_reconcile_only_touches_its_own_buckets() -> None:
"""**这是自治架构的核心不变量。**
共享存储里有别的节点的任务，去动它既越界（那是它的责任）也没有意义
（它的租约由它自己续）。
"""
store = InMemoryTaskStore()
mine = min(_NODE_A)
theirs = min(_NODE_B)
_plant_zombie(store, "mine", mine)
_plant_zombie(store, "theirs", theirs)
node_a = Controller(store=store, bucket_policy=_POLICY, owned_buckets=_NODE_A)
report = node_a.reconcile()
assert [item.task_id for item in report.lost] == ["mine"]
assert store.get("theirs").status is TaskStatus.RUNNING, ( # type: ignore[union-attr]
"node A 动了 node B 的任务"
)
node_b = Controller(store=store, bucket_policy=_POLICY, owned_buckets=_NODE_B)
assert [item.task_id for item in node_b.reconcile().lost] == ["theirs"]
assert store.get("theirs").status is TaskStatus.LOST # type: ignore[union-attr]
node_a.close(timeout=5)
node_b.close(timeout=5)
def test_reconcile_without_a_limit_still_covers_everything() -> None:
"""单机语义不能被这次改动破坏。"""
store = InMemoryTaskStore()
_plant_zombie(store, "one", min(_NODE_A))
_plant_zombie(store, "two", min(_NODE_B))
controller = Controller(store=store, bucket_policy=_POLICY)
report = controller.reconcile()
assert sorted(item.task_id for item in report.lost) == ["one", "two"]
assert {item.outcome for item in report.entries} == {ReconcileOutcome.MARKED_LOST}
controller.close(timeout=5)
def test_reconcile_ignores_a_task_with_no_bucket_when_limited() -> None:
"""没有归属桶的任务不属于任何节点——限定归属的节点不该去动它。"""
store = InMemoryTaskStore()
store.create(
TaskSnapshot(
task_id="loose",
name="无归属",
status=TaskStatus.RUNNING,
status_version=3,
attempt=1,
max_attempts=1,
created_at=_OLD,
started_at=_OLD,
lease_owner="dead",
lease_until=_OLD,
)
)
node = Controller(store=store, bucket_policy=_POLICY, owned_buckets=_NODE_A)
assert node.reconcile().entries == ()
assert store.get("loose").status is TaskStatus.RUNNING # type: ignore[union-attr]
node.close(timeout=5)
# ---------------------------------------------------------------------------
# 与统一平面的协同
# ---------------------------------------------------------------------------
def test_two_nodes_share_one_store_without_stepping_on_each_other() -> None:
"""每个节点跑自己的任务，共享存储提供统一的观测面。"""
store = InMemoryTaskStore()
key_a = _key_for(_POLICY, _NODE_A)
key_b = _key_for(_POLICY, _NODE_B)
node_a = Controller(store=store, bucket_policy=_POLICY, owned_buckets=_NODE_A, owner="a")
node_b = Controller(store=store, bucket_policy=_POLICY, owned_buckets=_NODE_B, owner="b")
node_a.register("noop", _noop)
node_b.register("noop", _noop)
node_a.submit("noop", task_id="a1", bucket_key=key_a)
node_b.submit("noop", task_id="b1", bucket_key=key_b)
node_a.wait_all(timeout=10)
node_b.wait_all(timeout=10)
# 统一观测面：任一个节点都能看到全部任务。
assert {s.task_id for s in node_a.list_tasks()} == {"a1", "b1"}
# 各自的"我这一份"：靠桶筛选，已有下推查询支持。
mine = node_a.list_tasks(buckets=node_a.owned_buckets)
assert [s.task_id for s in mine] == ["a1"]
assert node_a.owns_bucket(store.get("a1").bucket) # type: ignore[union-attr]
assert not node_a.owns_bucket(store.get("b1").bucket) # type: ignore[union-attr]
node_a.close(timeout=5)
node_b.close(timeout=5)
def test_stopping_a_task_owned_by_another_node_goes_through_the_command_channel() -> None:
"""控制层统一：一个节点可以停止另一个节点正在跑的任务——靠指令通道，不靠抢占。"""
store = InMemoryTaskStore()
node_b = Controller(
store=store,
bucket_policy=_POLICY,
owned_buckets=_NODE_B,
command_poll_interval=0.05,
owner="b",
)
node_b.register("block", _block)
node_b.submit("block", task_id="b1", bucket_key=_key_for(_POLICY, _NODE_B))
assert _wait_until(lambda: node_b.get("b1").status is TaskStatus.RUNNING) # type: ignore[union-attr]
node_a = Controller(store=store, bucket_policy=_POLICY, owned_buckets=_NODE_A, owner="a")
assert node_a.stop("b1", reason="另一个节点发起的停止") is True
assert _wait_until(lambda: node_a.get("b1").status is TaskStatus.STOPPED) # type: ignore[union-attr]
node_a.close(timeout=5)
node_b.close(timeout=5)
def test_a_node_that_declares_buckets_polls_commands_from_startup() -> None:
"""声明了归属 = 真 worker，从启动就轮询指令。
否则会有一个延迟生效的坑：别人 `stop` 一个属于本节点、当前还排在队列里的任务时，
指令已经写进存储，但本节点要等到**第一次提交**才启动轮询线程——对定时任务节点
来说那可能是几小时之后。
"""
store = InMemoryTaskStore()
node = Controller(store=store, bucket_policy=_POLICY, owned_buckets=_NODE_A)
assert node._worker is not None, "声明了归属的节点必须从启动就轮询"
node.close(timeout=5)
def test_a_pure_query_client_still_does_not_poll() -> None:
"""没有声明归属的实例不轮询——否则它会抢走别人的指令。"""
store = InMemoryTaskStore()
client = Controller(store=store, bucket_policy=_POLICY)
assert client._worker is None
client.close(timeout=5)
def _wait_until(predicate: Any, timeout: float = 15.0) -> bool:
import time
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if predicate():
return True
time.sleep(0.01)
return False
def _noop(ctx: object) -> str:
del ctx
return "ok"
def _block(ctx: object) -> str:
import time
for _ in range(1_000_000):
ctx.checkpoint() # type: ignore[attr-defined]
time.sleep(0.005)
return "never"
./tests\integration\test_persist_recovery.py
"""写穿（write-through）的恢复能力测试。
这一组用例存在的理由：**写穿失败曾经会让任务的存储状态永久冻结**。
故障链是：``_persist`` 用 ``expect_version = transition.version - 1`` 写入，这个式子
假设此前的每一次迁移都成功落库。一旦某次失败，存储停在 k-1 而内存已到 k，此后**每一次**
写入都因版本对不上而冲突——一次瞬时故障让任务在存储里永远停住，最后被 ``reconcile()``
判成 ``LOST``：把一个真正跑完的任务报成"结果未知"。
所以这里的重点不是"失败时不会崩"，而是"**失败之后能追上**"。
"""
from __future__ import annotations
import threading
import time
from typing import Any
import pytest
from chongming_task.controller import Controller
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 = 15.0
class FlakyStore(InMemoryTaskStore):
"""可以按需让 ``save`` 抖动若干次，之后恢复正常。
``fail_saves`` 计数：每次 ``save`` 尝试消耗一个；为 0 时完全正常。这样能精确模拟
"网络抖了一下"以及"存储长时间不可用"两种情况。
"""
def __init__(self, *, fail_saves: int = 0, fail_forever: bool = False) -> None:
super().__init__()
self.fail_saves = fail_saves
self.fail_forever = fail_forever
self.save_calls = 0
self._mutex = threading.RLock()
def save(self, snapshot: TaskSnapshot, *, expect_version: int) -> None:
with self._mutex:
self.save_calls += 1
if self.fail_forever or self.fail_saves > 0:
if not self.fail_forever:
self.fail_saves -= 1
raise OSError("模拟存储抖动")
super().save(snapshot, expect_version=expect_version)
def _wait_until_running(controller: Controller, task_id: str, *, timeout: float = 5.0) -> None:
"""等任务真的进入 ``RUNNING``。
**为什么必须显式等**：``TaskRun.pause()`` 对还没进入 ``RUNNING`` 的任务会
**立刻返回 False 且不产生任何迁移**。提交后马上开循环，若执行线程还没被调度起来，
整个循环会在微秒级跑完、一次都没暂停成功——用例看起来"通过到了断言那一步"，
实则前提从未成立（它曾因此间歇性失败）。机器越忙，这个窗口越大。
**为什么不能用 ``controller.get()``**：它**优先读存储**，而本文件的存储多半是故意
写不进去的 ``FlakyStore``——那样永远等不到 ``RUNNING``。这里只看本地执行句柄。
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
for run in controller.runs():
if run.snapshot.task_id == task_id and run.status is TaskStatus.RUNNING:
return
time.sleep(0.005)
raise AssertionError(f"{task_id} 没能在 {timeout}s 内进入 RUNNING")
# ---------------------------------------------------------------------------
# 核心：一次抖动不能让存储永久冻结
# ---------------------------------------------------------------------------
def test_a_single_transient_failure_does_not_freeze_the_store() -> None:
"""**这条是本次修复的核心回归用例。**
以前：第 3 次 save 抖动一次 → 存储永久停在 ``starting``，而任务其实跑完了。
现在：下一次迁移会把存储一次追上。
"""
store = FlakyStore(fail_saves=1)
controller = Controller(store=store)
controller.register("ok", _ok)
controller.submit("ok", task_id="t1")
outcome = controller.wait("t1", timeout=_WINDOW)
assert outcome.ok, "存储抖动不该让任务失败"
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.status is TaskStatus.SUCCEEDED, "存储必须追上内存"
assert snapshot.outcome == {"value": "done"}
history = [record.to_status for record in store.history("t1")]
assert history == [
TaskStatus.QUEUED,
TaskStatus.STARTING,
TaskStatus.RUNNING,
TaskStatus.SUCCEEDED,
], "迁移轨迹必须完整——缓冲的意义就在这里"
assert controller.desynced_tasks == {}, "瞬时故障不该造成永久分叉"
assert controller.persist_failures, "但失败必须留痕"
controller.close(timeout=_WINDOW)
def test_the_store_version_catches_up_after_a_gap() -> None:
"""版本会**跳号**（这是允许的），但必须单调且最终一致。"""
store = FlakyStore(fail_saves=2)
controller = Controller(store=store)
controller.register("ok", _ok)
controller.submit("ok", task_id="t1")
controller.wait("t1", timeout=_WINDOW)
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.status is TaskStatus.SUCCEEDED
assert snapshot.status_version == 4, "内存算出的版本"
versions = [record.version for record in store.history("t1")]
assert versions == [1, 2, 3, 4], "补写后轨迹不缺项"
assert versions == sorted(versions), "版本必须单调"
controller.close(timeout=_WINDOW)
def test_failures_spread_over_several_transitions_still_converge() -> None:
"""抖动散布在多次迁移上（而不是集中一次）也要能追上。"""
store = FlakyStore(fail_saves=3)
controller = Controller(store=store)
controller.register("ok", _ok)
controller.submit("ok", task_id="t1")
assert controller.wait("t1", timeout=_WINDOW).ok
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.status is TaskStatus.SUCCEEDED
controller.close(timeout=_WINDOW)
def test_a_long_outage_recovered_by_close_catches_up() -> None:
"""存储整个任务期间都不可用：任务照常成功，``close()`` 时尽力补一次。
补不上的话，任务在存储里就是落后的——这仍然比"任务失败"好，但必须在
``persist_failures`` 里留痕。
"""
store = FlakyStore(fail_forever=True)
controller = Controller(store=store)
controller.register("ok", _ok)
controller.submit("ok", task_id="t1")
outcome = controller.wait("t1", timeout=_WINDOW)
assert outcome.ok, "存储不可用不该让任务失败"
# 一次都没写进去：行停在 create 时的初始状态（pending）。
assert store.get("t1").status is TaskStatus.PENDING # type: ignore[union-attr]
assert len(controller.persist_failures) >= 1
# 存储恢复后 close() 会补上
store.fail_forever = False
controller.close(timeout=_WINDOW)
assert store.get("t1").status is TaskStatus.SUCCEEDED # type: ignore[union-attr]
def test_the_pending_buffer_is_bounded() -> None:
"""存储长期不可用时内存不能跟着无限增长。"""
from chongming_task.controller import _PENDING_TRANSITIONS
store = FlakyStore(fail_forever=True)
controller = Controller(store=store)
controller.register("spinner", _spinner)
controller.submit("spinner", task_id="t1")
_wait_until_running(controller, "t1")
# 让任务反复暂停/恢复，制造远超上限的迁移次数。
for _ in range(_PENDING_TRANSITIONS + 20):
assert controller.pause("t1", timeout=2.0), "任务没到达检查点：迁移没发生，用例失去意义"
controller.resume("t1")
state = controller._persist_states["t1"]
assert len(state.pending) <= _PENDING_TRANSITIONS
assert state.dropped > 0, "丢弃必须被计数，否则轨迹缺口就不可见了"
controller.stop("t1")
controller.wait_all(timeout=_WINDOW)
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 外部写入者：该放手时放手
# ---------------------------------------------------------------------------
def test_an_external_writer_marks_the_task_desynced_instead_of_being_overwritten() -> None:
"""用"上次确实写进去的版本"仍冲突 → 是**别人**写的，覆盖它会抹掉对方的状态。
这与"我们自己漏写了一笔"必须区分开：前者要放手，后者要追上去。
"""
store = InMemoryTaskStore()
controller = Controller(store=store, command_poll_interval=0.05)
controller.register("block", _block)
run = controller.submit("block", task_id="t1")
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
# 模拟另一个写入者把版本推进了（例如另一个进程认领了同一个任务）。
current = store.get("t1")
assert current is not None
intruder = _bump(current)
store.save(intruder, expect_version=current.status_version)
# 触发一次我们自己的迁移：应当发现外部写入并放手。
controller.stop("t1")
assert (
_wait_until(lambda: bool(controller.desynced_tasks))
or controller.wait_all(timeout=_WINDOW) is not None
)
desynced = controller.desynced_tasks
assert "t1" in desynced, "必须检测到外部写入者并停止写穿"
assert "外部写入者" in desynced["t1"]
# 对方的状态没有被我们覆盖。
assert store.get("t1").status is TaskStatus.RUNNING # type: ignore[union-attr]
controller.close(timeout=_WINDOW)
def test_desynced_tasks_are_reported_not_hidden() -> None:
"""永久分叉必须可见——只靠 persist_failures（最近 20 条）是不够的。"""
store = InMemoryTaskStore()
controller = Controller(store=store)
assert controller.desynced_tasks == {}
state = controller._persist_state("ghost")
assert controller.desynced_tasks == {}
state.desynced = "测试"
assert controller.desynced_tasks == {"ghost": "测试"}
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 非法返回值：响亮失败，而不是静默损坏
# ---------------------------------------------------------------------------
def test_a_non_serializable_return_value_fails_the_task_loudly() -> None:
"""以前：进程内 succeeded、存储停在 running、重启后被判 LOST。
现在：任务**失败**，原因写明是返回值不满足持久化契约。
"""
store = InMemoryTaskStore()
controller = Controller(store=store)
controller.register("bad_result", _bad_result)
controller.submit("bad_result", task_id="t1")
outcome = controller.wait("t1", timeout=_WINDOW)
assert not outcome.ok, "违反契约必须是失败，而不是静默的'成功'"
assert outcome.status is TaskStatus.FAILED
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.status is TaskStatus.FAILED, "存储与内存一致，不再是 running"
assert snapshot.failure is not None
assert snapshot.failure.kind is FailureKind.RESULT_INVALID
assert "不可 JSON 序列化" in snapshot.failure.message
assert snapshot.failure.retryable is False, "重试解决不了：同一个值还是不可序列化"
controller.close(timeout=_WINDOW)
def test_a_non_serializable_result_is_not_retried() -> None:
"""重试没有意义，所以只尝试一次。"""
store = InMemoryTaskStore()
controller = Controller(store=store)
controller.register("bad_result", _bad_result)
controller.submit("bad_result", task_id="t1", retry=_retry_three())
controller.wait("t1", timeout=_WINDOW)
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.attempt == 1, "不可序列化不是暂时性问题，不该重试"
controller.close(timeout=_WINDOW)
def test_serializable_results_are_unaffected() -> None:
"""正常路径不能被这次收紧影响。"""
store = InMemoryTaskStore()
controller = Controller(store=store)
controller.register("rich", _rich_result)
controller.submit("rich", task_id="t1")
outcome = controller.wait("t1", timeout=_WINDOW)
assert outcome.ok
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.outcome is not None
assert snapshot.outcome["value"]["nested"] == [1, 2, 3]
controller.close(timeout=_WINDOW)
def test_a_non_string_tag_is_rejected_at_submit_time() -> None:
"""标签会被写进存储（JSON 数组），所以必须在**入口**按类型契约拒绝。
真实类型是 ``frozenset[str]``。非字符串元素既不能正确往返，又只会在写穿路径上才
炸——而那条路径的失败不会让任务失败，问题会以"看板上的标签不对"的形式出现。
注意这条**必须用会序列化的后端**才测得出存储层的反应：``InMemoryTaskStore`` 直接
持有对象、从不编码。"内存后端能跑"不等于"真后端能跑"。
"""
store = SqliteTaskStore(":memory:")
try:
controller = Controller(store=store)
controller.register("ok", _ok)
with pytest.raises(ConfigurationError, match="标签必须是字符串"):
controller.submit("ok", task_id="t1", tags=[object()])
controller.close(timeout=_WINDOW)
finally:
store.close()
def test_a_non_serializable_progress_message_raises_at_the_call_site() -> None:
"""进度字段会被写进存储；类型错误必须在调用点炸掉，而不是流到写穿路径上。"""
store = InMemoryTaskStore()
controller = Controller(store=store)
controller.register("bad_progress", _bad_progress)
outcome = controller.wait(
controller.submit("bad_progress", task_id="t1").snapshot.task_id, timeout=_WINDOW
)
assert not outcome.ok
assert "message 必须是 str" in (outcome.failure.message if outcome.failure else "")
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.status is TaskStatus.FAILED
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 工具
# ---------------------------------------------------------------------------
def _bump(snapshot: TaskSnapshot) -> TaskSnapshot:
"""构造一个"别人写的"快照：版本前进一步。"""
from dataclasses import replace
return replace(snapshot, status_version=snapshot.status_version + 1)
def _retry_three() -> Any:
from chongming_task.policies import FixedBackoff, RetryPolicy
return RetryPolicy(max_attempts=3, backoff=FixedBackoff(0.0))
def _wait_until(predicate: Any, timeout: float = _WINDOW) -> bool:
import time
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if predicate():
return True
time.sleep(0.01)
return False
def _ok(ctx: Any) -> str:
del ctx
return "done"
def _bad_result(ctx: Any) -> Any:
del ctx
return {1, 2, 3}
def _rich_result(ctx: Any) -> dict[str, Any]:
del ctx
return {"nested": [1, 2, 3], "flag": True, "text": "中文"}
def _bad_progress(ctx: Any) -> str:
ctx.progress(1, total=10, message=object()) # type: ignore[arg-type]
return "unreachable"
def _spinner(ctx: Any) -> str:
import time
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.002)
return "never"
def _block(ctx: Any) -> str:
import time
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.005)
return "never"
./tests\integration\test_postgres_store.py
"""Postgres 后端的契约与跨连接测试。
需要一台可连的 Postgres，DSN 从环境变量给出::
CHONGMING_TEST_POSTGRES='postgresql://user@127.0.0.1:5432/tasks' \\
uv run pytest tests/integration/test_postgres_store.py
未设置时整组跳过——它依赖外部服务，与内置实现的契约测试（纯本地）刻意分开。
这里覆盖三类只有真数据库才能证明的事：
1. **完整契约**——与内存 / SQLite 跑的是同一套 ``StoreContract``；
2. **跨连接原子性**——设计说明 §15.4 明确要求后端自建并发用例：
N 个连接并发 ``save`` 恰好一个成功、并发 ``claim_commands`` 不重复下发、
持有者被强杀后租约仍按时间过期且连接池自愈；
3. **接线**——``Controller`` 真的能把它当存储用，跨 store 的控制走指令通道。
"""
from __future__ import annotations
import os
import threading
import time
from datetime import UTC, datetime, timedelta
from typing import Any
from uuid import uuid4
import pytest
psycopg = pytest.importorskip("psycopg")
pytest.importorskip("psycopg_pool")
from psycopg import sql # noqa: E402
from psycopg.conninfo import make_conninfo # noqa: E402
from chongming_task.controller import Controller # noqa: E402
from chongming_task.errors import ConcurrencyConflict, StoreError # noqa: E402
from chongming_task.states import TaskStatus # noqa: E402
from chongming_task.stores.postgres import PostgresTaskStore # noqa: E402
from chongming_task.stores.protocol import TaskStore # noqa: E402
from chongming_task.testing.store_contract import ( # noqa: E402
StoreContract,
command,
simple_snapshot,
)
_DSN = os.environ.get("CHONGMING_TEST_POSTGRES")
_T0 = datetime(2026, 1, 1, tzinfo=UTC)
pytestmark = pytest.mark.skipif(
not _DSN,
reason="未设置 CHONGMING_TEST_POSTGRES（DSN）：Postgres 集成测试需要一台可连的数据库",
)
# ---------------------------------------------------------------------------
# 隔离：每个存储实例独占一个 schema（与 SQLite 的临时文件等价）
# ---------------------------------------------------------------------------
def _isolated_dsn() -> tuple[str, str]:
"""建一个全新 schema，返回 ``(schema 名, 指向它的 DSN)``。"""
assert _DSN is not None
schema = f"cstore_{uuid4().hex[:12]}"
with psycopg.connect(_DSN, autocommit=True) as conn:
conn.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema)))
return schema, make_conninfo(_DSN, options=f"-csearch_path={schema}")
def _drop_schema(schema: str) -> None:
assert _DSN is not None
with psycopg.connect(_DSN, autocommit=True) as conn:
conn.execute(sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(sql.Identifier(schema)))
# ---------------------------------------------------------------------------
# 1. 完整契约
# ---------------------------------------------------------------------------
class TestPostgresStoreContract(StoreContract):
"""与内存 / SQLite 实现跑**完全相同**的用例。"""
_sandboxes: list[str]
def setup_method(self) -> None:
self._sandboxes = []
super().setup_method()
def make_store(self) -> TaskStore:
schema, dsn = _isolated_dsn()
self._sandboxes.append(schema)
return PostgresTaskStore(dsn)
def teardown_method(self) -> None:
super().teardown_method()
for schema in self._sandboxes:
_drop_schema(schema)
# ---------------------------------------------------------------------------
# 2. 跨连接原子性（StoreContract 覆盖不到，设计说明 §15.4 要求）
# ---------------------------------------------------------------------------
def test_concurrent_saves_have_exactly_one_winner() -> None:
"""N 个连接并发 save 同一任务：恰好一个成功，其余收到冲突。
乐观并发在**跨连接**下真正成立，而不是只在单连接语义里成立。
"""
schema, dsn = _isolated_dsn()
try:
seed = PostgresTaskStore(dsn)
seed.create(simple_snapshot("t1", version=1))
seed.close()
workers = 8
barrier = threading.Barrier(workers)
outcomes: list[str] = []
lock = threading.Lock()
def attempt() -> None:
store = PostgresTaskStore(dsn)
try:
barrier.wait(timeout=10)
try:
store.save(
simple_snapshot("t1", status=TaskStatus.SUCCEEDED, version=2),
expect_version=1,
)
except ConcurrencyConflict:
outcome = "conflict"
else:
outcome = "ok"
finally:
store.close()
with lock:
outcomes.append(outcome)
threads = [threading.Thread(target=attempt) for _ in range(workers)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout=30)
assert sorted(outcomes) == ["conflict"] * (workers - 1) + ["ok"]
finally:
_drop_schema(schema)
def test_concurrent_claims_hand_out_each_command_exactly_once() -> None:
"""N 个连接并发认领：同一条指令绝不会被两个连接取到。
"挑选 + 打租约"是一条 ``UPDATE ... FOR UPDATE SKIP LOCKED``，所以并发认领互不
阻塞、也不会重复。这是 SQLite 用 ``BEGIN IMMEDIATE`` 换来的同一个保证。
"""
schema, dsn = _isolated_dsn()
try:
total = 20
seed = PostgresTaskStore(dsn)
for index in range(total):
seed.submit_command(command(f"c{index}", created_offset=float(index)))
seed.close()
claimed: list[str] = []
lock = threading.Lock()
def drain() -> None:
store = PostgresTaskStore(dsn)
try:
while True:
batch = store.claim_commands(now=_T0, lease_seconds=300, limit=3)
if not batch:
return
with lock:
claimed.extend(item.command_id for item in batch)
finally:
store.close()
threads = [threading.Thread(target=drain) for _ in range(8)]
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout=30)
assert sorted(claimed) == sorted(f"c{index}" for index in range(total))
assert len(claimed) == len(set(claimed)), "同一条指令被下发了两次"
finally:
_drop_schema(schema)
def test_a_killed_holder_does_not_wedge_the_lease() -> None:
"""持有者被强杀：租约仍按时间过期，别的连接可以接管；连接池自愈。
租约是**存储里的数据**，不是连接上的会话状态——所以杀掉持有者的后端连接既不会
提前释放它，也不会永久锁住它。
"""
schema, dsn = _isolated_dsn()
try:
store_a = PostgresTaskStore(dsn, min_size=1, max_size=1)
store_a.create(simple_snapshot("t1"))
store_a.submit_command(command("c1"))
claimed = store_a.claim_commands(now=_T0, lease_seconds=30)
assert [item.command_id for item in claimed] == ["c1"]
# 模拟持有者进程消失：杀掉它的后端连接（本测试此刻只有 store_a 在用）。
killed: list[bool] = []
try:
with psycopg.connect(_DSN, autocommit=True) as admin:
rows = admin.execute(
"SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
"WHERE application_name = 'chongming-task' AND pid <> pg_backend_pid()"
).fetchall()
killed = [bool(row[0]) for row in rows]
except psycopg.errors.InsufficientPrivilege: # pragma: no cover - 权限受限的环境
pytest.skip("当前角色无权 pg_terminate_backend，跳过")
assert any(killed), "没有杀掉任何连接，测试前提不成立"
store_b = PostgresTaskStore(dsn)
# 租约未到期：接管者取不到。
assert store_b.claim_commands(now=_T0 + timedelta(seconds=10), lease_seconds=30) == []
# 租约到期：可以接管。
reclaimed = store_b.claim_commands(now=_T0 + timedelta(seconds=31), lease_seconds=30)
assert [item.command_id for item in reclaimed] == ["c1"]
# 被杀掉的持有者：连接池把死连接换掉，操作照常。
assert store_a.get("t1") is not None
store_a.close()
store_b.close()
finally:
_drop_schema(schema)
# ---------------------------------------------------------------------------
# 3. close() 的语义（模块文档承诺过，就要有用例守着）
# ---------------------------------------------------------------------------
def test_close_is_idempotent_and_operations_fail_afterwards() -> None:
schema, dsn = _isolated_dsn()
try:
store = PostgresTaskStore(dsn)
store.close()
store.close() # 幂等：重复关闭是空操作
with pytest.raises(StoreError, match="已关闭"):
store.get("t1")
finally:
_drop_schema(schema)
def test_a_role_without_ddl_privileges_can_use_an_existing_schema() -> None:
"""建表是**条件性**的：表已存在时一个 DDL 都不执行。
因此运维控制台可以用只读账号接进来（只给 SELECT 就够看板了）。若实现无条件跑
``CREATE TABLE IF NOT EXISTS``，这个账号会因为缺 CREATE 权限直接连不上。
"""
schema, dsn = _isolated_dsn()
role = f"cstore_ro_{uuid4().hex[:8]}"
password = uuid4().hex
try:
seed = PostgresTaskStore(dsn)
seed.create(simple_snapshot("t1", name="只读也能看到"))
seed.close()
try:
with psycopg.connect(_DSN, autocommit=True) as admin:
admin.execute(
sql.SQL("CREATE ROLE {} LOGIN PASSWORD {}").format(
sql.Identifier(role), sql.Literal(password)
)
)
except psycopg.errors.InsufficientPrivilege: # pragma: no cover - 权限受限的环境
pytest.skip("当前角色无权创建角色，跳过")
with psycopg.connect(_DSN, autocommit=True) as admin:
admin.execute(
sql.SQL("GRANT USAGE ON SCHEMA {} TO {}").format(
sql.Identifier(schema), sql.Identifier(role)
)
)
admin.execute(
sql.SQL(
"GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {} TO {}"
).format(sql.Identifier(schema), sql.Identifier(role))
)
info = psycopg.conninfo.conninfo_to_dict(dsn) if dsn else {}
info.update({"user": role, "password": password})
readonly_dsn = psycopg.conninfo.make_conninfo(**info)
store = PostgresTaskStore(readonly_dsn) # 没有 CREATE 权限，也必须能起来
try:
snapshot = store.get("t1")
assert snapshot is not None and snapshot.name == "只读也能看到"
assert len(store.query()) == 1
finally:
store.close()
finally:
with psycopg.connect(_DSN, autocommit=True) as admin:
admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role)))
admin.execute(sql.SQL("DROP ROLE IF EXISTS {}").format(sql.Identifier(role)))
_drop_schema(schema)
# ---------------------------------------------------------------------------
# 4. 接线：Controller 真的能把它当存储用
# ---------------------------------------------------------------------------
def _add(ctx: Any, left: int, right: int) -> int:
del ctx
return left + right
def _block(ctx: Any) -> str:
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.005)
return "done"
def _wait_until(predicate: Any, timeout: float = 10.0) -> bool:
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if predicate():
return True
time.sleep(0.01)
return False
def test_state_transitions_persist_and_are_auditable() -> None:
schema, dsn = _isolated_dsn()
try:
store = PostgresTaskStore(dsn)
controller = Controller(store=store)
controller.register("add", _add)
controller.submit("add", 2, 3, task_id="t1")
controller.wait_all(timeout=10)
snapshot = store.get("t1")
assert snapshot is not None
assert snapshot.status is TaskStatus.SUCCEEDED
assert snapshot.outcome == {"value": 5}
assert [item.to_status for item in store.history("t1")] == [
TaskStatus.QUEUED,
TaskStatus.STARTING,
TaskStatus.RUNNING,
TaskStatus.SUCCEEDED,
]
store.close()
finally:
_drop_schema(schema)
def test_stop_crosses_stores_through_postgres() -> None:
"""两个 store = 两个"进程"：worker 在跑，运维从另一套连接停掉它。"""
schema, dsn = _isolated_dsn()
try:
owner_store = PostgresTaskStore(dsn)
owner = Controller(store=owner_store, command_poll_interval=0.05, owner="worker-pg")
owner.register("block", _block)
run = owner.submit("block", task_id="t1")
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
client_store = PostgresTaskStore(dsn)
client = Controller(store=client_store, owner="ops-pg")
assert client.runs() == (), "客户端并不持有这个任务"
assert client.stop("t1", reason="跨 store 停止") is True
assert _wait_until(lambda: run.status is TaskStatus.STOPPED), "持有者没有取走指令"
final = client.get("t1")
assert final is not None and final.status is TaskStatus.STOPPED
owner.close(timeout=10)
client.close(timeout=5)
finally:
_drop_schema(schema)
./tests\integration\test_process_execution.py
"""进程执行器的端到端接线：``TaskRun`` / ``Controller`` / 存储。
单元测试盯的是句柄本身（暂停、强杀、心跳、崩溃）。这里盯的是**接进主线之后**的
三件承诺：
1. 快照里真的记下了子进程的 pid（排障时 `ps` 对得上号的就是它）；
2. 不响应停止的任务会被**超时机制真的杀掉**，而不是留一个卡死的线程（线程执行器
只能把它记成 `LOST`，这是进程执行器存在的理由）；
3. 这套东西写进存储之后仍然对得上（pid 是 store 契约里的字段，不是内存里的临时值）。
"""
from __future__ import annotations
import multiprocessing
import os
import time
from pathlib import Path
from typing import Any
from chongming_task import (
Controller,
FailureKind,
FixedBackoff,
ProcessExecutor,
RetryPolicy,
TaskStatus,
Timeout,
run_task,
)
from chongming_task.stores.sqlite import SqliteTaskStore
_WINDOW = 20.0
def _add(ctx: Any, left: int, right: int) -> int:
"""两数相加。"""
del ctx
return left + right
def _boom(ctx: Any) -> None:
"""总是失败：用来让任务反复重试。"""
del ctx
raise RuntimeError("总是失败")
def _deaf_loop(ctx: Any) -> str:
"""**不**调用 checkpoint：停不掉，只能强杀。"""
del ctx
end = time.monotonic() + 120
while time.monotonic() < end:
pass
return "never"
def test_a_task_run_records_the_child_pid_in_its_snapshot() -> None:
outcome = run_task(
_add,
2,
40,
task_id="proc-1",
executor=ProcessExecutor(),
timeout=Timeout(per_attempt=_WINDOW),
)
assert outcome.ok, outcome.failure
assert outcome.value == 42
snapshot = outcome.snapshot
assert snapshot.executor == "process"
assert snapshot.pid is not None
assert snapshot.pid != os.getpid(), "任务跑在另一个进程里"
def test_an_idle_task_is_actually_killed_by_the_timeout() -> None:
"""线程执行器下这条只能"放弃等待"并落 `LOST`（线程杀不掉）；进程执行器**真的**把它停掉。
超时仍然如实记成超时（不是"已停止"）——但载体确实没了，资源也确实释放了。
"""
started = time.monotonic()
outcome = run_task(
_deaf_loop,
task_id="proc-deaf",
executor=ProcessExecutor(),
timeout=Timeout(idle=1),
retry=RetryPolicy(max_attempts=1),
graceful_stop=0.5,
)
elapsed = time.monotonic() - started
assert outcome.status is TaskStatus.FAILED
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.TIMEOUT_IDLE
assert outcome.snapshot.pid is not None
# 它自己那个循环要跑两分钟，这里几秒就回来了——只有"真的被停掉"解释得通。
assert elapsed < _WINDOW
# 载体不在了：那个 pid 已经不再是本进程活着的子进程。
alive = {child.pid for child in multiprocessing.active_children()}
assert outcome.snapshot.pid not in alive, "子进程必须真的被终止并回收"
def test_a_controller_can_run_tasks_in_child_processes(tmp_path: Path) -> None:
store = SqliteTaskStore(tmp_path / "proc.db")
controller = Controller(store=store, executor=ProcessExecutor(), owner="node-a")
try:
controller.register("add", _add)
controller.submit("add", 20, 22, task_id="t1")
outcome = controller.wait("t1", timeout=_WINDOW)
assert outcome.ok and outcome.value == 42
assert outcome.snapshot.executor == "process"
# 落库之后仍然对得上：pid 是 store 契约里的字段，不是内存里的临时值。
stored = store.get("t1")
assert stored is not None
assert stored.pid == outcome.snapshot.pid
assert stored.executor == "process"
finally:
controller.close(timeout=_WINDOW, close_store=True)
def test_a_target_that_cannot_cross_the_boundary_fails_fast() -> None:
"""不可 pickle 的目标：任务落 ``failed`` 并带上诊断，而不是让调用方干等。
（直接调 ``ProcessExecutor.launch()`` 会当场抛 ``ConfigurationError``；这条测的是
走 ``run_task`` 那条路——异常在那里必须变成任务的结论。）
"""
def local_closure(ctx: Any) -> str: # 局部函数：子进程里找不到它
del ctx
return "nope"
started = time.monotonic()
outcome = run_task(local_closure, task_id="bad-1", executor=ProcessExecutor())
elapsed = time.monotonic() - started
assert outcome.status is TaskStatus.FAILED
assert outcome.failure is not None
assert "可 pickle" in outcome.failure.message
assert outcome.failure.retryable is False
assert elapsed < _WINDOW, "它必须立刻有结论，而不是等到天荒地老"
def test_a_retry_does_not_carry_the_previous_attempts_pid() -> None:
"""重试之间采一次快照：上一次的进程号必须已经清掉。
留着它的危害很具体：运维看到一个 pid 就去 `ps`，而对上的可能是一个早在别处跑着的
无关进程——比"没有 pid"更容易误导。
"""
controller = Controller(executor=ProcessExecutor())
controller.register("boom", _boom)
try:
controller.submit(
"boom",
task_id="t-retry",
retry=RetryPolicy(max_attempts=3, backoff=FixedBackoff(2.0)),
)
deadline = time.monotonic() + _WINDOW
while time.monotonic() < deadline:
snapshot = controller.get("t-retry")
if snapshot is not None and snapshot.status is TaskStatus.RETRY_WAIT:
break
time.sleep(0.02)
else: # pragma: no cover - 没等到就说明退避策略没生效
raise AssertionError("任务没有进入 retry_wait")
assert snapshot is not None
assert snapshot.pid is None, "上一次尝试的进程号必须清掉"
assert snapshot.executor == "process", "执行器是这一轮的配置，不随尝试变化"
finally:
controller.close(timeout=_WINDOW)
./tests\integration\test_reconcile.py
"""重启恢复（``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"
./tests\integration\test_store_wiring.py
"""存储接线测试：写穿、读穿、跨进程指令。
这里用**两个 ``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]
./tests\property\test_state_machine_properties.py
"""用 hypothesis 随机生成事件序列，验证状态机的核心不变量。
比手写用例更强的地方：它能生成人手想不到的组合，例如
"暂停中收到停止"、"重试等待中被重试"、"重试预算耗尽后又被派发"这类长尾序列。
被验证的不变量
--------------
1. 被拒绝的动作**完全不产生副作用**（快照对象不变、历史不增长）。
2. 每次被接受的动作使 ``status_version`` 恰好 +1，并使历史恰好 +1 条。
3. 除显式审计的复活路径外，实际发生的转移必在转移表内。
4. 终态不可逃逸。
5. 每次被接受的动作之后，快照都通过 :meth:`TaskSnapshot.validate` 的跨字段一致性检查。
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from hypothesis import HealthCheck, given, settings
from hypothesis import strategies as st
from chongming_task.core.state_machine import StateMachine, Trigger
from chongming_task.errors import AttemptsExhausted, InvalidTransition
from chongming_task.models import Failure, FailureKind
from chongming_task.states import TERMINAL, TRANSITIONS
from chongming_task.testing import FakeClock
if TYPE_CHECKING:
from collections.abc import Callable
_FAILURE = Failure(kind=FailureKind.EXCEPTION, message="boom")
#: 刻意绕过转移表的触发源，只有这三条审计路径合法。
_NON_TABLE_TRIGGERS = frozenset(
{
Trigger.RETRY_REQUESTED,
Trigger.RESET_REQUESTED,
Trigger.ADMIN_OVERRIDE,
}
)
#: 随机动作集合。
#:
#: 刻意不含 ``force_transition``：它按定义可跳到任意状态，与"随机事件序列"这一
#: 模型正交，其行为由 tests/unit/test_state_machine.py 的定向用例覆盖。
ACTIONS: tuple[tuple[str, Callable[[StateMachine], object]], ...] = (
("mark_blocked", lambda m: m.mark_blocked()),
("mark_queued", lambda m: m.mark_queued()),
("mark_starting", lambda m: m.mark_starting(executor="inline")),
("mark_running", lambda m: m.mark_running()),
("mark_succeeded", lambda m: m.mark_succeeded()),
("mark_failed", lambda m: m.mark_failed(_FAILURE)),
("mark_retry_wait", lambda m: m.mark_retry_wait(_FAILURE)),
("mark_retry_exhausted", lambda m: m.mark_retry_exhausted(_FAILURE)),
("mark_skipped", lambda m: m.mark_skipped()),
("mark_cancelled", lambda m: m.mark_cancelled()),
("mark_lost", lambda m: m.mark_lost()),
("request_stop", lambda m: m.request_stop()),
("mark_stopped", lambda m: m.mark_stopped()),
("request_pause", lambda m: m.request_pause()),
("mark_paused", lambda m: m.mark_paused()),
("abort_pause", lambda m: m.abort_pause()),
("resume", lambda m: m.resume()),
("retry", lambda m: m.retry()),
("force_retry", lambda m: m.retry(force=True)),
("reset", lambda m: m.reset()),
)
@settings(
max_examples=500,
deadline=None,
suppress_health_check=[HealthCheck.too_slow],
)
@given(
max_attempts=st.integers(min_value=1, max_value=4),
picks=st.lists(st.integers(min_value=0, max_value=len(ACTIONS) - 1), max_size=30),
)
def test_invariants_hold_for_any_action_sequence(max_attempts: int, picks: list[int]) -> None:
clock = FakeClock()
machine = StateMachine.create("t", "x", max_attempts=max_attempts, clock=clock)
for index in picks:
name, action = ACTIONS[index]
clock.advance(1)
before = machine.snapshot
history_len = len(machine.history)
try:
action(machine)
except (InvalidTransition, AttemptsExhausted):
# 不变量 1：被拒绝的动作必须完全无副作用。
assert machine.snapshot is before, f"{name} 被拒绝却改动了快照"
assert len(machine.history) == history_len, f"{name} 被拒绝却写了历史"
continue
after = machine.snapshot
record = machine.history[-1]
# 不变量 2：版本号与历史同步严格 +1。
assert len(machine.history) == history_len + 1, name
assert after.status_version == before.status_version + 1, name
# 不变量 3：只有审计路径可以绕过转移表。
if record.trigger not in _NON_TABLE_TRIGGERS:
assert after.status in TRANSITIONS[before.status], (
f"{name} 绕过了转移表：{before.status} -> {after.status}"
)
# 不变量 4：终态不可逃逸——**复活除外**，后者必须走审计路径。
# 终态出边为空集，因此这一条其实是上一条的推论；显式写出是为了失败时
# 能一眼分清"逃逸"与"绕表"。
if record.trigger not in _NON_TABLE_TRIGGERS:
assert before.status not in TERMINAL, f"{name} 从终态 {before.status} 逃逸"
# 不变量 5：快照跨字段一致。
problems = after.validate()
assert problems == (), f"{name} 产生不一致快照：{problems}"
./tests\unit\test_buckets.py
"""分桶策略测试。
最重要的一条是**跨进程稳定性**：内置 ``hash()`` 带进程级随机盐，若图省事用它，
同一个键在不同进程里会落到不同的桶，分布式下"同桶任务在同一台机器"这条约定会
**静默**失效。这里用真实子进程与不同的 ``PYTHONHASHSEED`` 把它钉死。
"""
from __future__ import annotations
import os
import subprocess
import sys
import pytest
from chongming_task.policies import BucketPolicy
def _run_with_hash_seed(script: str, seed: str) -> str:
"""在指定 ``PYTHONHASHSEED`` 的子进程里跑一段脚本，返回其标准输出。"""
environment = {**os.environ, "PYTHONHASHSEED": seed}
completed = subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
env=environment,
check=True,
timeout=60,
)
return completed.stdout.strip()
# ---------------------------------------------------------------------------
# 稳定性（分桶能用于分布式的全部前提）
# ---------------------------------------------------------------------------
def test_same_key_always_maps_to_the_same_bucket() -> None:
policy = BucketPolicy(64)
for _ in range(50):
assert policy.bucket_for("tenant-42") == policy.bucket_for("tenant-42")
def test_assignment_is_stable_across_processes_with_different_hash_seeds() -> None:
"""内置 ``hash()`` 带随机盐，本实现必须不受它影响。"""
script = (
"from chongming_task.policies import BucketPolicy;"
"print(' '.join(str(BucketPolicy(64).bucket_for(f'tenant-{i}')) for i in range(20)))"
)
first = _run_with_hash_seed(script, "0")
second = _run_with_hash_seed(script, "1")
assert first == second, "同一个键在不同进程里落到了不同的桶"
local = " ".join(str(BucketPolicy(64).bucket_for(f"tenant-{i}")) for i in range(20))
assert first == local, "子进程与本进程的结果不一致"
def test_builtin_hash_would_have_been_unusable() -> None:
"""留一条反例，说明为什么不能图省事用内置 ``hash()``。
不是"可能不一样"，而是**确实不一样**——这就是那个坑本身。
"""
script = "print(' '.join(str(hash(f'k{i}') % 64) for i in range(20)))"
assert _run_with_hash_seed(script, "0") != _run_with_hash_seed(script, "1")
# ---------------------------------------------------------------------------
# 取值范围与分布
# ---------------------------------------------------------------------------
def test_bucket_is_always_within_range() -> None:
for num_buckets in (1, 2, 7, 64, 1000):
policy = BucketPolicy(num_buckets)
for index in range(200):
bucket = policy.bucket_for(f"key-{index}")
assert 0 <= bucket < num_buckets
def test_single_bucket_policy_puts_everything_in_bucket_zero() -> None:
policy = BucketPolicy(1)
assert all(policy.bucket_for(key) == 0 for key in ("a", "b", "c"))
def test_distribution_is_reasonably_uniform() -> None:
"""分桶明显倾斜会让某些机器空转、另一些排队。
这里只做一个粗检：64 个桶、4096 个键，每个桶的落点不应当偏离均值太多。
"""
policy = BucketPolicy(64)
counts = [0] * 64
for index in range(4096):
counts[policy.bucket_for(f"tenant-{index}")] += 1
expected = 4096 / 64
assert all(count > 0 for count in counts), "有桶完全没分到键，分布明显异常"
assert max(counts) < expected * 3, f"分布过于集中：max={max(counts)}"
def test_similar_keys_do_not_cluster() -> None:
"""``tenant-1`` 与 ``tenant-2`` 这类只差一个字符的键必须被打散。"""
policy = BucketPolicy(64)
buckets = {policy.bucket_for(f"tenant-{index}") for index in range(50)}
assert len(buckets) > 20, "相邻键挤在少数几个桶里，哈希的雪崩性不足"
def test_policy_validates_bucket_range() -> None:
policy = BucketPolicy(8)
assert policy.is_valid(0)
assert policy.is_valid(7)
assert not policy.is_valid(8)
assert not policy.is_valid(-1)
def test_policy_exposes_all_buckets() -> None:
assert list(BucketPolicy(4).buckets) == [0, 1, 2, 3]
# ---------------------------------------------------------------------------
# 参数校验
# ---------------------------------------------------------------------------
def test_num_buckets_must_be_positive() -> None:
with pytest.raises(ValueError, match="num_buckets"):
BucketPolicy(0)
def test_empty_key_is_rejected() -> None:
"""空键会让所有"没填租户"的任务挤进同一个桶。"""
with pytest.raises(ValueError, match="不能为空"):
BucketPolicy().bucket_for("")
def test_policy_is_hashable_and_comparable() -> None:
"""策略会作为配置传递，值语义让它能被比较与当作字典键。"""
assert BucketPolicy(64) == BucketPolicy(64)
assert BucketPolicy(64) != BucketPolicy(128)
assert len({BucketPolicy(64), BucketPolicy(64)}) == 1
./tests\unit\test_cancellation.py
"""协作式取消令牌的行为测试。
这里必须用**真实线程**：令牌的价值恰恰在于跨线程正确握手（任务线程阻塞在
暂停点、调度线程请求停止并唤醒它）。用假的线程模型测不出竞态。
因此这些用例带小超时并断言线程确实退出——**同时也在验证没有死锁**。
"""
from __future__ import annotations
import asyncio
import contextlib
import threading
import pytest
from chongming_task.errors import TaskCancelled
from chongming_task.runtime.cancellation import CancellationToken
_JOIN_TIMEOUT = 5.0
def _run_in_thread(target: object) -> threading.Thread:
thread = threading.Thread(target=target) # type: ignore[arg-type]
thread.daemon = True
thread.start()
return thread
# ---------------------------------------------------------------------------
# 停止
# ---------------------------------------------------------------------------
def test_checkpoint_is_noop_without_any_request() -> None:
"""没有请求时 checkpoint 必须立即返回，否则每轮循环都在付代价。"""
CancellationToken().checkpoint()
def test_checkpoint_raises_when_stop_requested() -> None:
token = CancellationToken()
token.request_stop("发布窗口")
with pytest.raises(TaskCancelled) as info:
token.checkpoint()
assert info.value.reason == "发布窗口"
def test_stop_is_idempotent_and_keeps_the_first_reason() -> None:
"""首个原因通常最接近真实起因，后续请求不应覆盖它。"""
token = CancellationToken()
token.request_stop("first")
token.request_stop("second")
with pytest.raises(TaskCancelled) as info:
token.checkpoint()
assert info.value.reason == "first"
def test_raise_if_stopped_ignores_pause() -> None:
"""用于清理路径：只响应停止，绝不阻塞。"""
token = CancellationToken()
token.request_pause()
token.raise_if_stopped() # 不应阻塞
token.request_stop()
with pytest.raises(TaskCancelled):
token.raise_if_stopped()
def test_is_stop_requested_reflects_state() -> None:
token = CancellationToken()
assert not token.is_stop_requested
token.request_stop()
assert token.is_stop_requested
# ---------------------------------------------------------------------------
# 暂停
# ---------------------------------------------------------------------------
def test_pause_blocks_until_resume() -> None:
token = CancellationToken()
token.request_pause() # 先请求，消除"任务先跑到检查点"的竞态
observed: list[str] = []
def worker() -> None:
token.checkpoint()
observed.append("resumed")
thread = _run_in_thread(worker)
assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True
assert observed == [], "暂停请求生效前任务不应继续执行"
token.resume()
thread.join(_JOIN_TIMEOUT)
assert not thread.is_alive(), "resume() 必须唤醒阻塞中的任务"
assert observed == ["resumed"]
def test_stop_wakes_a_paused_task() -> None:
"""暂停中的任务必须能被停止，否则"暂停"就成了不可撤销的挂起。"""
token = CancellationToken()
token.request_pause()
reasons: list[str | None] = []
def worker() -> None:
try:
token.checkpoint()
except TaskCancelled as exc:
reasons.append(exc.reason)
thread = _run_in_thread(worker)
assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True
token.request_stop("运维介入")
thread.join(_JOIN_TIMEOUT)
assert not thread.is_alive(), "停止请求必须能唤醒暂停点"
assert reasons == ["运维介入"]
def test_repeated_pause_does_not_clear_acknowledgement() -> None:
"""重复请求暂停不应把"已确认"重置掉，否则 wait_paused 会误判为未生效。"""
token = CancellationToken()
token.request_pause()
def worker() -> None:
token.checkpoint()
thread = _run_in_thread(worker)
assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True
token.request_pause()
assert token.is_pause_acknowledged is True
token.resume()
thread.join(_JOIN_TIMEOUT)
def test_wait_paused_times_out_when_task_never_checkpoints() -> None:
"""调度器必须能区分"已暂停"与"暂停没生效"——后者要回滚为 RUNNING 并告警。"""
token = CancellationToken()
token.request_pause()
assert token.wait_paused(timeout=0.05) is False
assert token.is_pause_acknowledged is False
def test_wait_paused_reports_false_when_stopped_instead() -> None:
"""请求暂停后又被停止：这不是"暂停成功"，必须返回 False。"""
token = CancellationToken()
token.request_pause()
token.request_stop()
assert token.wait_paused(timeout=_JOIN_TIMEOUT) is False
def test_resume_without_pause_is_noop() -> None:
token = CancellationToken()
token.resume()
assert not token.is_pause_requested
token.checkpoint() # 仍应立即返回
def _start_pausing_worker(token: CancellationToken) -> threading.Event:
"""起一个持续调用 ``checkpoint()`` 的工作线程，返回"完成过一轮"的事件。
线程必须**持续**调用 checkpoint——只调一次就退出的话，之后的暂停请求自然没人
确认，那是测试写错而不是缺陷。
"""
def worker() -> None:
# 收尾时的 stop 会让 checkpoint 抛 TaskCancelled：那正是预期行为，
# 但线程里未捕获的异常会被 pytest 转成错误。
with contextlib.suppress(TaskCancelled):
while True:
token.checkpoint()
ran.set()
ran = threading.Event()
threading.Thread(target=worker, daemon=True).start()
return ran
def test_pause_is_acknowledged_again_after_a_resume() -> None:
"""**回归用例**：暂停 → 恢复 → 再暂停里，第二次暂停曾经永久失效。
``checkpoint()`` 原先只在**进入**暂停分支时设置 ``_ack``。``resume()`` 把
``_pause`` 置回 False 并通知后，任务线程还没来得及重新取到锁，新的
``request_pause()`` 就把 ``_pause`` 又置成了 True——线程重新检查循环条件时发现
"还要暂停"，于是继续等待，而它已经越过了设置 ``_ack`` 的那一行。
结果 ``wait_paused()`` 永远等不到确认，暂停超时回滚。表现出来就是"暂停一次成功，
之后再也暂停不了"。
"""
token = CancellationToken()
ran = _start_pausing_worker(token)
assert ran.wait(_JOIN_TIMEOUT), "工作线程没有跑起来"
for round_index in range(5):
token.request_pause()
assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True, (
f"第 {round_index + 1} 轮暂停没有得到确认"
)
ran.clear()
token.resume()
assert ran.wait(_JOIN_TIMEOUT), f"第 {round_index + 1} 轮恢复后任务没有继续推进"
token.request_stop()
def test_a_pause_right_after_a_resume_is_still_acknowledged() -> None:
"""恢复与再次暂停之间不留任何间隔——最容易踩中的时序。"""
token = CancellationToken()
ran = _start_pausing_worker(token)
assert ran.wait(_JOIN_TIMEOUT)
token.request_pause()
assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True
# 立刻恢复再立刻暂停：不给任务线程任何重新取锁的机会。
token.resume()
token.request_pause()
assert token.wait_paused(timeout=_JOIN_TIMEOUT) is True, "紧接着的第二次暂停丢了"
token.resume()
token.request_stop()
# ---------------------------------------------------------------------------
# 异步路径
# ---------------------------------------------------------------------------
def test_acheckpoint_fast_path_does_not_offload_to_thread() -> None:
"""紧密循环里的每次检查都丢线程会拖垮事件循环，快速路径必须留在本线程。"""
async def main() -> str:
token = CancellationToken()
for _ in range(2000):
await token.acheckpoint()
return "done"
assert asyncio.run(main()) == "done"
def test_acheckpoint_raises_when_stopped() -> None:
token = CancellationToken()
token.request_stop("已取消")
async def main() -> None:
await token.acheckpoint()
with pytest.raises(TaskCancelled) as info:
asyncio.run(main())
assert info.value.reason == "已取消"
def test_acheckpoint_suspends_and_resumes_without_blocking_loop() -> None:
"""暂停期间事件循环必须仍能推进——这正是异步版本存在的理由。"""
token = CancellationToken()
token.request_pause()
async def main() -> list[str]:
events: list[str] = []
async def heartbeat() -> None:
while True:
events.append("tick")
await asyncio.sleep(0.01)
ticker = asyncio.create_task(heartbeat())
waiter = asyncio.create_task(token.acheckpoint())
while not token.is_pause_acknowledged:
await asyncio.sleep(0.01)
ticks_while_paused = len(events)
assert ticks_while_paused > 0, "暂停期间事件循环不应被阻塞"
token.resume()
await asyncio.wait_for(waiter, timeout=_JOIN_TIMEOUT)
ticker.cancel()
return events
assert asyncio.run(main())
./tests\unit\test_controller.py
"""任务管理门面（Controller / Registry）的测试。
重点覆盖三件事：**白名单**（只能启动已注册的任务）、**并发额度**（超限排队而不是
无限起线程）、以及**收尾的诚实性**（排队中的任务被取消时不应真的执行任务体）。
"""
from __future__ import annotations
import time
from typing import Any
import pytest
from chongming_task.controller import Controller, StatusCounts, TaskRegistry, render_table
from chongming_task.errors import ConfigurationError
from chongming_task.policies import BucketPolicy, ConcurrencyPolicy
from chongming_task.run import TaskRun
from chongming_task.states import TaskStatus
_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 _wait_for_status(run: TaskRun, status: TaskStatus, timeout: float = _WINDOW) -> bool:
return _wait_until(lambda: run.status is status, timeout)
# --- 注册用任务（模块级，便于复用） ---
def add(ctx: Any, left: int, right: int) -> int:
"""两数相加。"""
del ctx
return left + right
def slow_loop(ctx: Any) -> str:
"""持续小步推进直到被停止。"""
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.005)
return "done"
def never_runs(ctx: Any) -> str:
"""排队期间被取消时，这个函数体不应该被执行。"""
marker.append("ran")
return "ran"
marker: list[str] = []
# ---------------------------------------------------------------------------
# 注册表：白名单
# ---------------------------------------------------------------------------
def test_registry_rejects_duplicate_names() -> None:
"""静默覆盖会让"我改了实现，跑的还是旧的"这类问题极难发现。"""
registry = TaskRegistry()
registry.register("add", add)
with pytest.raises(ConfigurationError, match="已注册"):
registry.register("add", add)
def test_registry_requires_a_callable() -> None:
with pytest.raises(ConfigurationError, match="可调用"):
TaskRegistry().register("bad", "not-callable") # type: ignore[arg-type]
def test_registry_rejects_blank_name() -> None:
with pytest.raises(ConfigurationError, match="不能为空"):
TaskRegistry().register(" ", add)
def test_registry_picks_up_the_docstring_as_description() -> None:
registration = TaskRegistry().register("add", add)
assert registration.description == "两数相加。"
def test_submitting_an_unregistered_task_is_rejected() -> None:
"""安全底线：入口只接受**名字**，绝不接受代码。
一个能启动任意可调用对象的入口等于任意代码执行——这是将来接 CLI 还是
HTTP 都不能破的规则。
"""
controller = Controller()
with pytest.raises(ConfigurationError, match="未注册"):
controller.submit("nope")
assert controller.registry.names() == ()
def test_registry_lists_registered_names() -> None:
controller = Controller()
controller.register("add", add)
controller.register("slow", slow_loop)
assert controller.registry.names() == ("add", "slow")
assert len(controller.registry) == 2
assert "add" in controller.registry
# ---------------------------------------------------------------------------
# 提交、查询、统计
# ---------------------------------------------------------------------------
def test_submit_runs_and_reports_snapshot() -> None:
controller = Controller()
controller.register("add", add)
run = controller.submit("add", 2, 3, task_id="t1", label="加法")
assert isinstance(run, TaskRun)
outcome = controller.wait("t1", timeout=_WINDOW)
assert outcome.ok
assert outcome.value == 5
snapshot = controller.get("t1")
assert snapshot is not None
assert snapshot.task_id == "t1"
assert snapshot.name == "加法"
assert snapshot.status is TaskStatus.SUCCEEDED
def test_get_returns_none_for_unknown_task() -> None:
assert Controller().get("nope") is None
def test_stats_counts_by_status() -> None:
controller = Controller()
controller.register("add", add)
for index in range(3):
controller.submit("add", index, 1, task_id=f"t{index}")
controller.wait_all(timeout=_WINDOW)
stats = controller.stats()
assert isinstance(stats, StatusCounts)
assert stats.total == 3
assert stats.of(TaskStatus.SUCCEEDED) == 3
assert stats.finished == 3
assert stats.unfinished == 0
def test_list_filters_by_status_and_name_prefix() -> None:
controller = Controller()
controller.register("add", add)
controller.submit("add", 1, 1, task_id="a1", label="抽取-甲")
controller.submit("add", 1, 1, task_id="a2", label="抽取-乙")
controller.wait_all(timeout=_WINDOW)
assert len(controller.list_tasks()) == 2
assert len(controller.list_tasks(name_prefix="抽取")) == 2
assert controller.list_tasks(name_prefix="加载") == []
assert len(controller.list_tasks(status=TaskStatus.SUCCEEDED)) == 2
assert controller.list_tasks(status=TaskStatus.FAILED) == []
assert len(controller.list_tasks(limit=1)) == 1
def test_list_filters_by_tag() -> None:
controller = Controller()
controller.register("add", add, tags={"io"})
controller.submit("add", 1, 1, task_id="a1")
controller.submit("add", 1, 1, task_id="a2", tags={"cpu"})
controller.wait_all(timeout=_WINDOW)
assert len(controller.list_tasks(tag="io")) == 2, "注册时的标签应当被继承"
assert [s.task_id for s in controller.list_tasks(tag="cpu")] == ["a2"]
def test_duplicate_task_id_is_rejected() -> None:
controller = Controller()
controller.register("add", add)
controller.submit("add", 1, 1, task_id="dup")
with pytest.raises(ConfigurationError, match="已存在"):
controller.submit("add", 1, 1, task_id="dup")
# ---------------------------------------------------------------------------
# 并发额度
# ---------------------------------------------------------------------------
def test_excess_submissions_queue_instead_of_starting() -> None:
"""超限的任务状态是 ``queued``（已被受理、等待执行槽），不是 ``pending``。"""
controller = Controller(max_parallel=1)
controller.register("slow", slow_loop)
first = controller.submit("slow", task_id="t1")
assert _wait_for_status(first, TaskStatus.RUNNING)
second = controller.submit("slow", task_id="t2")
third = controller.submit("slow", task_id="t3")
assert second.status is TaskStatus.QUEUED
assert third.status is TaskStatus.QUEUED
assert not second.started, "排队中的任务不应被启动"
assert controller.active_count == 1
assert controller.queued_count == 2
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
assert controller.queued_count == 0
def test_queued_tasks_start_as_slots_free_up() -> None:
controller = Controller(max_parallel=2)
controller.register("add", add)
runs = [controller.submit("add", index, 1, task_id=f"t{index}") for index in range(4)]
outcomes = controller.wait_all(timeout=_WINDOW)
assert len(outcomes) == 4
assert all(outcome.ok for outcome in outcomes)
assert sorted(outcome.value for outcome in outcomes) == [1, 2, 3, 4]
assert controller.active_count == 0
assert all(run.status is TaskStatus.SUCCEEDED for run in runs)
def test_max_parallel_must_be_positive() -> None:
with pytest.raises(ConfigurationError, match="max_parallel"):
Controller(max_parallel=0)
# ---------------------------------------------------------------------------
# 具名资源信号量与优先级（ConcurrencyPolicy）
#
# 这三个字段曾经只是"声明了、校验了、文档里写着支持"——调度器从不读它，于是
# ``limits`` 与 ``priority`` 设了等于没设，而且**不报任何错**。下面几条把它们钉住。
# ---------------------------------------------------------------------------
def test_named_resource_limit_caps_only_the_tagged_tasks() -> None:
"""``limits={"gpu": 1}`` 只约束带 ``gpu`` 标签的任务，不是给全局加一道闸。"""
controller = Controller(concurrency=ConcurrencyPolicy(limits={"gpu": 1}))
controller.register("slow", slow_loop)
first = controller.submit("slow", task_id="gpu-1", tags=["gpu"])
assert _wait_for_status(first, TaskStatus.RUNNING)
second = controller.submit("slow", task_id="gpu-2", tags=["gpu"])
assert second.status is TaskStatus.QUEUED, "gpu 名额已满，第二个必须排队"
# 不带该标签的任务不受它约束——具名信号量不是全局额度。
plain = controller.submit("slow", task_id="plain")
assert _wait_for_status(plain, TaskStatus.RUNNING)
assert controller.active_count == 2
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
def test_a_task_needs_room_in_every_limited_tag_it_carries() -> None:
"""多标签任务要**每个**受限期资源都还有余量才放行。"""
controller = Controller(concurrency=ConcurrencyPolicy(limits={"gpu": 1, "db": 1}))
controller.register("slow", slow_loop)
gpu = controller.submit("slow", task_id="gpu-holder", tags=["gpu"])
assert _wait_for_status(gpu, TaskStatus.RUNNING)
db = controller.submit("slow", task_id="db-holder", tags=["db"])
assert _wait_for_status(db, TaskStatus.RUNNING)
both = controller.submit("slow", task_id="needs-both", tags=["gpu", "db"])
assert both.status is TaskStatus.QUEUED, "两个资源都满了，带两个标签的任务更不能放行"
# 只空出一个：db 还占着，所以它依旧不能启动。等额度记账真的落下来再断言，
# 否则这句话可能在调度器还没醒的时候"通过"，什么都没证明。
assert controller.stop("gpu-holder") is True
controller.wait("gpu-holder", timeout=_WINDOW)
assert _wait_until(lambda: controller.active_count == 1), "额度应当已归还"
assert both.status is TaskStatus.QUEUED, "gpu 空了但 db 还占着，不能放行"
# 两个都空出来才轮到它。
assert controller.stop("db-holder") is True
controller.wait("db-holder", timeout=_WINDOW)
assert _wait_for_status(both, TaskStatus.RUNNING)
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
def test_tags_without_a_limit_do_not_throttle_anything() -> None:
"""没写进 ``limits`` 的标签不受约束（常见误解：以为打个标签就会被限流）。"""
controller = Controller(concurrency=ConcurrencyPolicy(limits={"gpu": 1}))
controller.register("slow", slow_loop)
first = controller.submit("slow", task_id="io-1", tags=["io"])
second = controller.submit("slow", task_id="io-2", tags=["io"])
assert _wait_for_status(first, TaskStatus.RUNNING)
assert _wait_for_status(second, TaskStatus.RUNNING)
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
def test_priority_picks_the_highest_first() -> None:
"""都能启动时先取优先级最高的；数值越大越优先。"""
controller = Controller(max_parallel=1)
controller.register("slow", slow_loop)
blocker = controller.submit("slow", task_id="blocker")
assert _wait_for_status(blocker, TaskStatus.RUNNING)
low = controller.submit("slow", task_id="low", priority=1)
high = controller.submit("slow", task_id="high", priority=5)
assert controller.stop("blocker") is True
controller.wait("blocker", timeout=_WINDOW)
assert _wait_for_status(high, TaskStatus.RUNNING), "高优先级的应当先被取走"
assert low.status is TaskStatus.QUEUED, "低优先级的不该插到前面"
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
def test_equal_priority_keeps_submission_order() -> None:
"""同优先级仍然先进先出——优先级机制不该打乱时间序。"""
controller = Controller(max_parallel=1)
controller.register("slow", slow_loop)
blocker = controller.submit("slow", task_id="blocker")
assert _wait_for_status(blocker, TaskStatus.RUNNING)
first_queued = controller.submit("slow", task_id="first")
second_queued = controller.submit("slow", task_id="second")
assert controller.stop("blocker") is True
controller.wait("blocker", timeout=_WINDOW)
assert _wait_for_status(first_queued, TaskStatus.RUNNING), "先提交的应当先跑"
assert second_queued.status is TaskStatus.QUEUED
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
def test_policy_default_priority_applies_to_submissions() -> None:
"""``ConcurrencyPolicy.priority`` 是 ``submit`` 不给 ``priority`` 时的默认值。"""
controller = Controller(concurrency=ConcurrencyPolicy(max_parallel=1, priority=7))
controller.register("slow", slow_loop)
blocker = controller.submit("slow", task_id="blocker")
assert _wait_for_status(blocker, TaskStatus.RUNNING)
assert blocker.priority == 7, "没给 priority 时应当取策略里的默认值"
explicit_low = controller.submit("slow", task_id="low", priority=1)
assert explicit_low.priority == 1, "显式给的优先级必须能覆盖默认值"
default_high = controller.submit("slow", task_id="high")
assert default_high.priority == 7
assert controller.stop("blocker") is True
controller.wait("blocker", timeout=_WINDOW)
assert _wait_for_status(default_high, TaskStatus.RUNNING)
assert explicit_low.status is TaskStatus.QUEUED
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
def test_concurrency_policy_max_parallel_is_honoured() -> None:
"""策略里的 ``max_parallel`` 与 ``Controller(max_parallel=...)`` 等价。"""
controller = Controller(concurrency=ConcurrencyPolicy(max_parallel=1))
controller.register("slow", slow_loop)
first = controller.submit("slow", task_id="t1")
assert _wait_for_status(first, TaskStatus.RUNNING)
second = controller.submit("slow", task_id="t2")
assert second.status is TaskStatus.QUEUED
assert controller.active_count == 1
controller.stop_all(reason="测试结束")
controller.wait_all(timeout=_WINDOW)
def test_policy_and_max_parallel_are_mutually_exclusive() -> None:
"""两份配置只能给一份：谁说了算必须明确，不能靠"后者覆盖前者"这种猜法。"""
with pytest.raises(ConfigurationError, match="只能给一个"):
Controller(concurrency=ConcurrencyPolicy(max_parallel=2), max_parallel=1)
def test_policy_rejects_a_non_positive_resource_limit() -> None:
with pytest.raises(ValueError, match="并发上限"):
ConcurrencyPolicy(limits={"gpu": 0})
def test_a_non_integer_priority_is_rejected_at_submit_time() -> None:
"""非整数优先级必须在提交时拦下。
放过去的话，它会在**派发线程**里抛 ``TypeError``——线程直接死掉，此后所有排队
任务都不再启动，而没有任何报错指向真正的原因。
"""
controller = Controller()
controller.register("add", add)
with pytest.raises(ConfigurationError, match="priority"):
controller.submit("add", 1, 2, task_id="t1", priority="high") # type: ignore[arg-type]
# ---------------------------------------------------------------------------
# 控制
# ---------------------------------------------------------------------------
def test_stop_routes_to_the_right_task() -> None:
controller = Controller()
controller.register("slow", slow_loop)
victim = controller.submit("slow", task_id="victim")
bystander = controller.submit("slow", task_id="bystander")
assert _wait_for_status(victim, TaskStatus.RUNNING)
assert _wait_for_status(bystander, TaskStatus.RUNNING)
assert controller.stop("victim", reason="发布窗口") is True
outcome = controller.wait("victim", timeout=_WINDOW)
assert outcome.status is TaskStatus.STOPPED
assert bystander.status is TaskStatus.RUNNING, "不该误伤别的任务"
controller.stop_all()
controller.wait_all(timeout=_WINDOW)
def test_stop_returns_false_for_unknown_task() -> None:
assert Controller().stop("nope") is False
def test_pause_and_resume_route_through_the_controller() -> None:
controller = Controller()
controller.register("slow", slow_loop)
run = controller.submit("slow", task_id="p1")
assert _wait_for_status(run, TaskStatus.RUNNING)
assert controller.pause("p1", timeout=_WINDOW) is True
assert run.status is TaskStatus.PAUSED
assert controller.resume("p1") is True
assert run.status is TaskStatus.RUNNING
controller.stop_all()
controller.wait_all(timeout=_WINDOW)
def test_stop_all_reports_how_many_were_affected() -> None:
controller = Controller()
controller.register("slow", slow_loop)
for index in range(3):
controller.submit("slow", task_id=f"t{index}")
assert _wait_until(lambda: controller.active_count == 3)
assert controller.stop_all(reason="批量收工") == 3
controller.wait_all(timeout=_WINDOW)
assert controller.stop_all() == 0, "已无未落终态的任务"
def test_submit_after_close_is_rejected() -> None:
controller = Controller()
controller.register("add", add)
controller.close(timeout=_WINDOW)
with pytest.raises(ConfigurationError, match="已关闭"):
controller.submit("add", 1, 1)
# ---------------------------------------------------------------------------
# 收尾
# ---------------------------------------------------------------------------
def test_close_does_not_execute_queued_task_bodies() -> None:
"""排队中的任务被取消时不该真的跑起来——否则"停止"只是把副作用推迟了。"""
marker.clear()
controller = Controller(max_parallel=1)
controller.register("slow", slow_loop)
controller.register("never", never_runs)
controller.submit("slow", task_id="t-block")
assert _wait_until(lambda: controller.active_count == 1)
controller.submit("never", task_id="t-queued")
assert controller.queued_count == 1
outcomes = controller.close(reason="收工", timeout=_WINDOW)
assert marker == [], "排队中的任务被启动了，任务体被执行"
statuses = {outcome.snapshot.task_id: outcome.status for outcome in outcomes}
assert statuses["t-queued"] is TaskStatus.STOPPED
assert statuses["t-block"] is TaskStatus.STOPPED
def test_prune_drops_oldest_terminal_records() -> None:
controller = Controller()
controller.register("add", add)
for index in range(5):
controller.submit("add", index, 1, task_id=f"t{index}")
controller.wait_all(timeout=_WINDOW)
assert len(controller.list_tasks()) == 5
assert controller.prune(keep_terminal=2) == 3
assert len(controller.list_tasks()) == 2
assert controller.prune(keep_terminal=2) == 0
# ---------------------------------------------------------------------------
# 渲染
# ---------------------------------------------------------------------------
def test_render_table_aligns_chinese_names_by_display_width() -> None:
"""中日韩全角字符占两列；不按显示宽度对齐，中文任务名会把整张表挤歪。"""
controller = Controller()
controller.register("add", add)
controller.submit("add", 1, 1, task_id="a1", label="抽取")
controller.submit("add", 1, 1, task_id="a2", label="a-much-longer-name")
controller.wait_all(timeout=_WINDOW)
table = render_table(controller.list_tasks())
lines = table.splitlines()
assert lines[0].startswith("TASK_ID")
assert set(lines[1]) <= {"-", " "}, "表头下方应当是分隔线"
widths = {len(line.rstrip()) for line in lines}
assert len(widths) > 1, "各行列宽应当随内容变化，而不是靠不齐"
def test_render_table_reports_failure_kind() -> None:
def boom(ctx: Any) -> None:
raise RuntimeError("炸了")
controller = Controller()
controller.register("boom", boom)
controller.submit("boom", task_id="b1")
controller.wait_all(timeout=_WINDOW)
table = render_table(controller.list_tasks())
assert "exception" in table
assert "failed" in table
# ---------------------------------------------------------------------------
# 分桶
# ---------------------------------------------------------------------------
def _two_keys_in_different_buckets(policy: BucketPolicy) -> tuple[str, str]:
"""找两个确实落在不同桶里的键。
硬编码两个键是不行的：它们有约 1/64 的概率撞进同一个桶，那会让"分桶隔离"
的用例变成一句空话，而且只在特定版本上偶发失败。
"""
first = "tenant-0"
first_bucket = policy.bucket_for(first)
for index in range(1, 500):
candidate = f"tenant-{index}"
if policy.bucket_for(candidate) != first_bucket:
return first, candidate
raise AssertionError("在 500 个键里找不到第二个桶，分桶策略有问题")
def test_bucket_is_derived_from_the_key_and_recorded() -> None:
controller = Controller()
controller.register("add", add)
run = controller.submit("add", 1, 1, task_id="a1", bucket_key="tenant-42")
controller.wait_all(timeout=_WINDOW)
snapshot = run.snapshot
assert snapshot.bucket == controller.bucket_policy.bucket_for("tenant-42")
assert snapshot.bucket_key == "tenant-42", "原始键必须留存，否则改桶数时无法迁移"
def test_same_key_lands_in_the_same_bucket() -> None:
controller = Controller()
controller.register("add", add)
for index in range(3):
controller.submit("add", index, 1, task_id=f"t{index}", bucket_key="tenant-9")
controller.wait_all(timeout=_WINDOW)
buckets = {snapshot.bucket for snapshot in controller.list_tasks()}
assert len(buckets) == 1
def test_task_without_bucket_key_is_unbucketed() -> None:
controller = Controller()
controller.register("add", add)
controller.submit("add", 1, 1, task_id="a1")
controller.wait_all(timeout=_WINDOW)
assert controller.get("a1").bucket is None # type: ignore[union-attr]
assert controller.bucket_counts() == {None: 1}
def test_per_bucket_limit_isolates_buckets_from_each_other() -> None:
"""一个桶把额度占满不应饿死其他桶——这是分桶在**本地**就有的价值。
若调度器直接取队首而不是扫描可运行者，一个积压的桶会阻塞所有桶，把分桶要解决
的问题又原样引入一遍。
"""
controller = Controller(max_parallel=10, max_parallel_per_bucket=1)
controller.register("slow", slow_loop)
hot, cold = _two_keys_in_different_buckets(controller.bucket_policy)
first = controller.submit("slow", task_id="a1", bucket_key=hot)
second = controller.submit("slow", task_id="a2", bucket_key=hot)
other = controller.submit("slow", task_id="b1", bucket_key=cold)
assert _wait_for_status(first, TaskStatus.RUNNING)
assert _wait_for_status(other, TaskStatus.RUNNING), "别的桶被积压的桶挡住了"
assert second.status is TaskStatus.QUEUED, "同桶超出额度应当排队"
assert controller.active_in_bucket(controller.bucket_policy.bucket_for(hot)) == 1
controller.stop_all()
controller.wait_all(timeout=_WINDOW)
def test_per_bucket_limit_still_respects_the_global_limit() -> None:
controller = Controller(max_parallel=1, max_parallel_per_bucket=5)
controller.register("slow", slow_loop)
hot, cold = _two_keys_in_different_buckets(controller.bucket_policy)
first = controller.submit("slow", task_id="a1", bucket_key=hot)
second = controller.submit("slow", task_id="b1", bucket_key=cold)
assert _wait_for_status(first, TaskStatus.RUNNING)
assert second.status is TaskStatus.QUEUED, "总额度才是硬上限"
controller.stop_all()
controller.wait_all(timeout=_WINDOW)
def test_per_bucket_limit_must_be_positive() -> None:
with pytest.raises(ConfigurationError, match="max_parallel_per_bucket"):
Controller(max_parallel_per_bucket=0)
def test_list_tasks_filters_by_bucket_and_bucket_set() -> None:
controller = Controller()
controller.register("add", add)
hot, cold = _two_keys_in_different_buckets(controller.bucket_policy)
controller.submit("add", 1, 1, task_id="a1", bucket_key=hot)
controller.submit("add", 1, 1, task_id="b1", bucket_key=cold)
controller.submit("add", 1, 1, task_id="c1")
controller.wait_all(timeout=_WINDOW)
hot_bucket = controller.bucket_policy.bucket_for(hot)
cold_bucket = controller.bucket_policy.bucket_for(cold)
assert [s.task_id for s in controller.list_tasks(bucket=hot_bucket)] == ["a1"]
assert [s.task_id for s in controller.list_tasks(buckets=[hot_bucket, cold_bucket])] == [
"a1",
"b1",
]
assert controller.list_tasks(bucket=9999) == []
def test_bucket_counts_summarises_each_bucket() -> None:
controller = Controller()
controller.register("add", add)
hot, cold = _two_keys_in_different_buckets(controller.bucket_policy)
controller.submit("add", 1, 1, task_id="a1", bucket_key=hot)
controller.submit("add", 1, 1, task_id="a2", bucket_key=hot)
controller.submit("add", 1, 1, task_id="b1", bucket_key=cold)
controller.wait_all(timeout=_WINDOW)
counts = controller.bucket_counts()
assert counts[controller.bucket_policy.bucket_for(hot)] == 2
assert counts[controller.bucket_policy.bucket_for(cold)] == 1
def test_render_table_shows_the_bucket_column() -> None:
controller = Controller()
controller.register("add", add)
controller.submit("add", 1, 1, task_id="a1", bucket_key="tenant-42")
controller.submit("add", 1, 1, task_id="a2")
controller.wait_all(timeout=_WINDOW)
table = render_table(controller.list_tasks())
lines = table.splitlines()
assert "BUCKET" in lines[0]
expected = str(controller.bucket_policy.bucket_for("tenant-42"))
assert expected in lines[2]
assert lines[3].split()[2] == "-", "未分桶的任务应当显示 -"
./tests\unit\test_inline_executor.py
"""内联执行器测试。
重点是两件事：它**确实**同步跑完（违反 launch 不阻塞的约定），以及它对无法
支持的暂停能力**如实返回 False** 而不是假装成功。
"""
from __future__ import annotations
import asyncio
from typing import Any
from chongming_task.models import FailureKind
from chongming_task.runtime.executors.base import ExecutionRequest
from chongming_task.runtime.executors.inline import InlineExecutor
def _request(target: Any, *args: Any) -> ExecutionRequest:
return ExecutionRequest(task_id="t-inline", name="t-inline", target=target, args=args)
def test_launch_runs_to_completion_before_returning() -> None:
"""这是它与 ``Executor`` 协议唯一的差异，也是它存在的意义。"""
order: list[str] = []
def work(ctx: Any) -> str:
del ctx
order.append("task")
return "done"
handle = InlineExecutor().launch(_request(work))
order.append("after-launch")
assert order == ["task", "after-launch"], "内联执行必须在 launch 返回前跑完"
assert handle.running is False
outcome = handle.wait()
assert outcome is not None
assert outcome.value == "done"
def test_arguments_are_forwarded_after_the_context() -> None:
def work(ctx: Any, left: int, right: int, *, scale: int = 1) -> int:
del ctx
return (left + right) * scale
handle = InlineExecutor().launch(
ExecutionRequest(
task_id="t-inline-args",
name="t-inline-args",
target=work,
args=(2, 3),
kwargs={"scale": 10},
)
)
outcome = handle.wait()
assert outcome is not None
assert outcome.value == 50
def test_failure_is_reported_like_any_other_executor() -> None:
def boom(ctx: Any) -> None:
del ctx
raise KeyError("missing")
outcome = InlineExecutor().launch(_request(boom)).wait()
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.EXCEPTION
assert "KeyError" in outcome.failure.message
assert outcome.failure.traceback is not None
def test_pause_reports_false_because_it_cannot_be_honoured() -> None:
"""如实返回 False，调用方据此回滚为 RUNNING，而不是以为暂停成功了。"""
handle = InlineExecutor().launch(_request(lambda ctx: "ok"))
assert handle.pause(timeout=0.01) is False
def test_resume_is_harmless_after_the_fact() -> None:
handle = InlineExecutor().launch(_request(lambda ctx: "ok"))
handle.resume()
assert handle.running is False
def test_stop_on_a_finished_inline_task_returns_its_real_outcome() -> None:
"""已结束的任务应当返回真实结果，而不是落进"兜底放弃"分支。"""
handle = InlineExecutor().launch(_request(lambda ctx: 7))
outcome = handle.stop(graceful_timeout=1.0)
assert outcome.ok
assert outcome.value == 7
assert outcome.abandoned is False
def test_context_is_injected() -> None:
seen: dict[str, Any] = {}
def work(ctx: Any) -> str:
seen["task_id"] = ctx.task_id
seen["token"] = ctx.token
return "ok"
handle = InlineExecutor().launch(_request(work))
assert seen["task_id"] == "t-inline"
assert seen["token"] is handle.token
def test_async_target_runs_to_completion_before_launch_returns() -> None:
"""``async def`` 任务在内联执行器下同样是同步跑完的（本线程起一个事件循环）。"""
async def work(ctx: Any) -> int:
await asyncio.sleep(0)
return 42
handle = InlineExecutor().launch(_request(work))
assert handle.running is False
outcome = handle.wait()
assert outcome is not None
assert outcome.ok
assert outcome.value == 42
./tests\unit\test_models.py
"""快照与正交维度模型的不变量。
重点覆盖那些"容易被合并成单一状态枚举"的维度：进度、失败原因、请求动作。
"""
from __future__ import annotations
import dataclasses
from datetime import UTC, datetime
import pytest
from chongming_task.models import (
AttemptRecord,
Failure,
FailureKind,
Progress,
RequestedAction,
TaskSnapshot,
TimeoutScope,
)
from chongming_task.states import TaskStatus
def test_progress_fraction_is_none_when_total_unknown() -> None:
"""total 未知必须返回 None 而非 0.0——后者会被渲染成"尚未开始"。"""
assert Progress(done=5).fraction is None
assert Progress(done=0, total=0).fraction is None
assert Progress(done=25, total=100).fraction == 0.25
def test_progress_rejects_done_beyond_total() -> None:
with pytest.raises(ValueError, match="不能超过"):
Progress(done=11, total=10)
def test_progress_rejects_negative_values() -> None:
with pytest.raises(ValueError, match="不能为负"):
Progress(done=-1)
with pytest.raises(ValueError, match="不能为负"):
Progress(total=-1)
def test_timeout_scope_maps_to_failure_kind() -> None:
"""超时是失败原因而非生命周期状态，两者的对应关系必须稳定。"""
assert FailureKind.for_timeout(TimeoutScope.START) is FailureKind.TIMEOUT_START
assert FailureKind.for_timeout(TimeoutScope.ATTEMPT) is FailureKind.TIMEOUT_ATTEMPT
assert FailureKind.for_timeout(TimeoutScope.IDLE) is FailureKind.TIMEOUT_IDLE
def test_attempt_record_numbering_starts_at_one() -> None:
now = datetime(2026, 1, 1, tzinfo=UTC)
with pytest.raises(ValueError, match="从 1 开始"):
AttemptRecord(task_id="t1", attempt=0, started_at=now)
def test_snapshot_requires_identity() -> None:
with pytest.raises(ValueError, match="task_id"):
TaskSnapshot(task_id="", name="x")
with pytest.raises(ValueError, match="name"):
TaskSnapshot(task_id="t1", name="")
def test_snapshot_rejects_attempt_beyond_max_attempts() -> None:
with pytest.raises(ValueError, match="max_attempts"):
TaskSnapshot(task_id="t1", name="x", attempt=3, max_attempts=2)
def test_snapshot_is_frozen() -> None:
"""快照不可变，才能安全跨线程/协程传递并直接作为持久化记录。"""
snap = TaskSnapshot(task_id="t1", name="x")
with pytest.raises(dataclasses.FrozenInstanceError):
snap.status = TaskStatus.RUNNING # type: ignore[misc]
def test_validate_flags_terminal_missing_ended_at() -> None:
snap = TaskSnapshot(task_id="t1", name="x", status=TaskStatus.SUCCEEDED)
assert any("ended_at" in problem for problem in snap.validate())
def test_validate_flags_leftover_requested_action_on_terminal() -> None:
"""终态残留请求会让 UI 永远显示"停止中"。"""
now = datetime(2026, 1, 1, tzinfo=UTC)
snap = TaskSnapshot(
task_id="t1",
name="x",
status=TaskStatus.STOPPED,
ended_at=now,
requested_action=RequestedAction.STOP,
)
assert any("requested_action" in problem for problem in snap.validate())
def test_retry_wait_may_carry_failure() -> None:
"""RETRY_WAIT 必须能携带上次失败原因，否则界面无法回答"它为什么在等"。"""
snap = TaskSnapshot(
task_id="t1",
name="x",
status=TaskStatus.RETRY_WAIT,
failure=Failure(kind=FailureKind.EXCEPTION, message="boom"),
)
assert snap.validate() == ()
def test_retry_wait_without_failure_is_inconsistent() -> None:
snap = TaskSnapshot(task_id="t1", name="x", status=TaskStatus.RETRY_WAIT)
assert any("RETRY_WAIT" in problem for problem in snap.validate())
def test_success_must_not_carry_failure() -> None:
now = datetime(2026, 1, 1, tzinfo=UTC)
snap = TaskSnapshot(
task_id="t1",
name="x",
status=TaskStatus.SUCCEEDED,
ended_at=now,
failure=Failure(kind=FailureKind.EXCEPTION, message="boom"),
)
assert any("SUCCEEDED" in problem for problem in snap.validate())
def test_consistent_terminal_snapshot_validates_clean() -> None:
now = datetime(2026, 1, 1, tzinfo=UTC)
snap = TaskSnapshot(
task_id="t1",
name="x",
status=TaskStatus.FAILED,
ended_at=now,
attempt=1,
max_attempts=1,
failure=Failure(kind=FailureKind.EXCEPTION, message="boom"),
)
assert snap.validate() == ()
assert snap.is_terminal
assert snap.is_revivable
def test_running_snapshot_requires_started_at_only_on_validate() -> None:
"""时间线完整性刻意不放在 __post_init__，以免逐步构造快照时处处受限。"""
snap = TaskSnapshot(task_id="t1", name="x", status=TaskStatus.RUNNING)
assert any("started_at" in problem for problem in snap.validate())
# 但构造本身不报错
assert snap.status is TaskStatus.RUNNING
./tests\unit\test_pipeline.py
"""DAG 管道测试。
分三层：
1. **定义期校验**——环、悬空依赖、未注册任务、quorum 都要在创建任何任务之前拒绝，
否则存储里会留下一堆看起来正常的半成品任务；
2. **推进**——屏障、扇出、扇入、三种汇合语义；
3. **收场**——fail_fast 连坐、下游跳过、整体状态推导。
约定：节点 id 一律用 ``node_id=`` 显式给出。第一个位置参数是**已注册的任务名**，
其后才是传给任务的位置参数——这两者很容易看混，显式写出来既清楚又自成文档。
"""
from __future__ import annotations
import threading
import time
from typing import Any
import pytest
from chongming_task.controller import Controller
from chongming_task.errors import ConfigurationError
from chongming_task.pipeline import (
Pipeline,
PipelineStatus,
render_pipeline,
)
from chongming_task.policies import AggregationPolicy, JoinMode
from chongming_task.states import TaskStatus
from chongming_task.stores.memory import InMemoryTaskStore
_WINDOW = 15.0
# ---------------------------------------------------------------------------
# 1. 定义期校验
# ---------------------------------------------------------------------------
def test_empty_pipeline_is_rejected() -> None:
with pytest.raises(ConfigurationError, match="没有任何节点"):
Pipeline("empty").validate()
def test_dangling_dependency_is_rejected() -> None:
pipeline = Pipeline("dangling")
pipeline.node("add", node_id="a", depends_on=["no-such-node"])
with pytest.raises(ConfigurationError, match="不存在的节点"):
pipeline.validate()
def test_self_dependency_is_rejected() -> None:
pipeline = Pipeline("self")
pipeline.node("add", node_id="a", depends_on=["a"])
with pytest.raises(ConfigurationError, match="依赖了自己"):
pipeline.validate()
def test_cycle_is_rejected() -> None:
"""有环时必须显式报错。
不检测的话调度器只会安静地什么都不做（所有节点都在等对方），表现为"管道卡住了"
——那种故障最难查。
"""
pipeline = Pipeline("cyclic")
first = pipeline.node("add", node_id="a")
second = pipeline.node("add", node_id="b", depends_on=[first])
# 人为制造环：a 反过来依赖 b。
object.__setattr__(pipeline.node_by_id("a"), "depends_on", ("b",))
assert second.depends_on == ("a",)
with pytest.raises(ConfigurationError, match="环形依赖"):
pipeline.validate()
def test_unregistered_task_is_rejected_before_creating_anything() -> None:
"""控制面只接受任务名。写错的 DAG 不能留下半成品任务。"""
store = InMemoryTaskStore()
controller = Controller(store=store)
controller.register("add", _add)
pipeline = Pipeline("bad")
pipeline.node("add", node_id="ok")
pipeline.node("nope", node_id="typo")
with pytest.raises(ConfigurationError):
controller.submit_pipeline(pipeline)
assert store.query() == [], "校验失败前不该创建任何任务"
def test_quorum_settlement_is_rejected_with_an_explanation() -> None:
"""诚实拒绝，而不是把 quorum 硬套到 DAG 上发明一个没人能预测的语义。"""
pipeline = Pipeline("q", aggregation=AggregationPolicy.of_quorum(2))
pipeline.node("add")
with pytest.raises(ConfigurationError, match="不支持 quorum"):
pipeline.validate()
def test_duplicate_node_id_is_rejected() -> None:
pipeline = Pipeline("dup")
pipeline.node("add", node_id="same")
with pytest.raises(ConfigurationError, match="节点 id 重复"):
pipeline.node("add", node_id="same")
def test_same_task_can_appear_many_times_with_auto_ids() -> None:
"""扇出靠"同名多次"表达，不该强迫调用方给每个节点起名。"""
pipeline = Pipeline("fanout")
pipeline.node("add")
pipeline.node("add")
pipeline.node("add")
assert [node.node_id for node in pipeline.nodes] == ["add", "add#2", "add#3"]
def test_positional_args_belong_to_the_task_not_the_node_id() -> None:
"""位置参数是传给任务的，不是节点名——这正是把参数叫 ``node_id`` 的原因。"""
pipeline = Pipeline("args")
node = pipeline.node("add", 2, 3, node_id="sum")
assert node.node_id == "sum"
assert node.args == (2, 3)
def test_topological_order_puts_upstream_first() -> None:
pipeline = Pipeline("order")
last = pipeline.node("add", node_id="last")
first = pipeline.node("add", node_id="first")
pipeline.node("add", node_id="middle", depends_on=[first, last])
order = [node.node_id for node in pipeline.topological_order()]
assert order.index("first") < order.index("middle")
assert order.index("last") < order.index("middle")
# ---------------------------------------------------------------------------
# 2. 推进
# ---------------------------------------------------------------------------
def test_a_linear_pipeline_runs_in_dependency_order() -> None:
controller = Controller(max_parallel=4)
controller.register("mark", _mark)
pipeline = Pipeline("linear")
extract = pipeline.node("mark", node_id="extract")
transform = pipeline.node("mark", node_id="transform", depends_on=[extract])
pipeline.node("mark", node_id="load", depends_on=[transform])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p1")
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
assert _finished_order(controller, "p1") == ["extract", "transform", "load"]
controller.close(timeout=_WINDOW)
def test_a_downstream_node_waits_in_blocked_until_its_upstream_finishes() -> None:
gate = threading.Event()
controller = Controller(max_parallel=4)
controller.register("wait_gate", _wait_gate)
controller.register("mark", _mark)
pipeline = Pipeline("blocking")
upstream = pipeline.node("wait_gate", gate, node_id="upstream")
pipeline.node("mark", node_id="downstream", depends_on=[upstream])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p2")
try:
assert _wait_until(lambda: run.node_status("upstream") is TaskStatus.RUNNING)
time.sleep(0.1)
assert run.node_status("downstream") is TaskStatus.BLOCKED, "上游没完成前必须停在 blocked"
assert run.status is PipelineStatus.RUNNING
gate.set()
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
finally:
gate.set()
controller.close(timeout=_WINDOW)
def test_fan_out_runs_siblings() -> None:
controller = Controller(max_parallel=8)
controller.register("mark", _mark)
pipeline = Pipeline("fanout")
root = pipeline.node("mark", node_id="root")
for index in range(4):
pipeline.node("mark", node_id=f"shard{index}", depends_on=[root])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p3")
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
assert set(run.task_ids) == {"root", "shard0", "shard1", "shard2", "shard3"}
controller.close(timeout=_WINDOW)
def test_fan_in_with_all_succeeded_waits_for_every_upstream() -> None:
controller = Controller(max_parallel=8)
controller.register("mark", _mark)
pipeline = Pipeline("fanin")
root = pipeline.node("mark", node_id="root")
left = pipeline.node("mark", node_id="left", depends_on=[root])
right = pipeline.node("mark", node_id="right", depends_on=[root])
pipeline.node("mark", node_id="join", depends_on=[left, right])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p4")
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
order = _finished_order(controller, "p4")
assert order.index("join") > order.index("left")
assert order.index("join") > order.index("right")
controller.close(timeout=_WINDOW)
def test_any_succeeded_join_runs_after_the_first_success() -> None:
"""先到先得：任一上游成功即可放行，不必等另一个失败的上游。"""
controller = Controller(max_parallel=8)
controller.register("mark", _mark)
controller.register("boom", _boom)
pipeline = Pipeline("any")
good = pipeline.node("mark", node_id="good")
bad = pipeline.node("boom", node_id="bad")
winner = pipeline.node(
"mark", node_id="winner", depends_on=[good, bad], join=JoinMode.ANY_SUCCEEDED
)
run = controller.submit_pipeline(pipeline, pipeline_run_id="p5")
outcome = run.wait(timeout=_WINDOW)
assert winner.join is JoinMode.ANY_SUCCEEDED
assert outcome.nodes["winner"] is TaskStatus.SUCCEEDED, "有一个上游成功就该放行"
assert outcome.status is PipelineStatus.FAILED, "另一个上游失败，整条管道仍失败"
controller.close(timeout=_WINDOW)
def test_any_succeeded_join_skips_when_nothing_succeeds() -> None:
controller = Controller(max_parallel=8)
controller.register("boom", _boom)
controller.register("mark", _mark)
pipeline = Pipeline("any-none")
bad = pipeline.node("boom", node_id="bad")
pipeline.node("mark", node_id="winner", depends_on=[bad], join=JoinMode.ANY_SUCCEEDED)
run = controller.submit_pipeline(pipeline, pipeline_run_id="p6")
outcome = run.wait(timeout=_WINDOW)
assert outcome.nodes["winner"] is TaskStatus.SKIPPED
assert outcome.status is PipelineStatus.FAILED
controller.close(timeout=_WINDOW)
def test_all_settled_join_runs_even_after_an_upstream_failure() -> None:
"""清理/通知型节点：不管上游成败都要跑。"""
controller = Controller(max_parallel=8)
controller.register("boom", _boom)
controller.register("mark", _mark)
pipeline = Pipeline("cleanup", aggregation=AggregationPolicy.all_settled())
failed = pipeline.node("boom", node_id="work")
cleanup = pipeline.node(
"mark", node_id="cleanup", depends_on=[failed], join=JoinMode.ALL_SETTLED
)
run = controller.submit_pipeline(pipeline, pipeline_run_id="p7")
outcome = run.wait(timeout=_WINDOW)
assert cleanup.join is JoinMode.ALL_SETTLED
assert outcome.nodes["cleanup"] is TaskStatus.SUCCEEDED, "ALL_SETTLED 下游必须照跑"
assert outcome.status is PipelineStatus.FAILED, "上游失败仍让整条管道失败"
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 3. 收场
# ---------------------------------------------------------------------------
def test_failed_upstream_skips_downstream_with_a_reason() -> None:
controller = Controller(max_parallel=8)
controller.register("boom", _boom)
controller.register("mark", _mark)
pipeline = Pipeline("skip")
bad = pipeline.node("boom", node_id="bad")
pipeline.node("mark", node_id="after", depends_on=[bad])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p8")
outcome = run.wait(timeout=_WINDOW)
assert outcome.nodes["after"] is TaskStatus.SKIPPED
snapshot = controller.get("p8:after")
assert snapshot is not None
assert snapshot.status is TaskStatus.SKIPPED
assert snapshot.failure is None, "跳过不是失败，不该带 failure"
controller.close(timeout=_WINDOW)
def test_fail_fast_cancels_independent_running_nodes() -> None:
"""真正需要连坐的场景：另一个分支已经在跑，必须被取消。
被连坐的必须落 ``CANCELLED`` 而不是 ``STOPPED``——"上游挂了"与"值班同学手动停的"
是完全不同的事，混成一个状态就再也分不出来了。
"""
gate = threading.Event()
controller = Controller(max_parallel=8)
controller.register("boom", _boom)
controller.register("wait_gate", _wait_gate)
pipeline = Pipeline("cascade", aggregation=AggregationPolicy.fail_fast())
pipeline.node("wait_gate", gate, node_id="independent")
pipeline.node("boom", node_id="bad")
run = controller.submit_pipeline(pipeline, pipeline_run_id="p9")
try:
assert _wait_until(lambda: run.node_status("independent") is TaskStatus.RUNNING)
outcome = run.wait(timeout=_WINDOW)
assert outcome.status is PipelineStatus.FAILED
assert outcome.nodes["independent"] is TaskStatus.CANCELLED, (
"被连坐的应当是 cancelled，而不是 stopped"
)
snapshot = controller.get("p9:independent")
assert snapshot is not None
assert snapshot.failure is not None
assert snapshot.failure.kind.value == "cancelled"
finally:
gate.set()
controller.close(timeout=_WINDOW)
def test_fail_fast_never_settles_while_a_cancelled_node_is_still_running() -> None:
"""``status`` 必须先看终态、再看失败标志。
顺序反了，等待者会被提前唤醒，看到一个"已经失败、但还有任务在跑"的矛盾状态。
"""
gate = threading.Event()
controller = Controller(max_parallel=8)
controller.register("boom", _boom)
controller.register("wait_gate", _wait_gate)
pipeline = Pipeline("ordering", aggregation=AggregationPolicy.fail_fast())
pipeline.node("wait_gate", gate, node_id="independent")
pipeline.node("boom", node_id="bad")
run = controller.submit_pipeline(pipeline, pipeline_run_id="p10")
try:
assert _wait_until(lambda: run.node_status("independent") is TaskStatus.RUNNING)
outcome = run.wait(timeout=_WINDOW)
for node_id in outcome.nodes:
assert outcome.nodes[node_id] in {
TaskStatus.SUCCEEDED,
TaskStatus.FAILED,
TaskStatus.SKIPPED,
TaskStatus.CANCELLED,
TaskStatus.STOPPED,
TaskStatus.LOST,
}, f"{node_id} 在管道给出结论时还没落终态"
finally:
gate.set()
controller.close(timeout=_WINDOW)
def test_all_settled_does_not_cancel_siblings() -> None:
gate = threading.Event()
controller = Controller(max_parallel=8)
controller.register("boom", _boom)
controller.register("wait_gate", _wait_gate)
pipeline = Pipeline("settled", aggregation=AggregationPolicy.all_settled())
pipeline.node("wait_gate", gate, node_id="independent")
pipeline.node("boom", node_id="bad")
run = controller.submit_pipeline(pipeline, pipeline_run_id="p11")
try:
assert _wait_until(lambda: run.node_status("independent") is TaskStatus.RUNNING)
time.sleep(0.2)
assert run.node_status("independent") is TaskStatus.RUNNING, "all_settled 不该连坐"
gate.set()
outcome = run.wait(timeout=_WINDOW)
assert outcome.nodes["independent"] is TaskStatus.SUCCEEDED
assert outcome.status is PipelineStatus.FAILED, "有人失败过，整体就是失败"
finally:
gate.set()
controller.close(timeout=_WINDOW)
def test_pipeline_tasks_carry_the_run_id() -> None:
controller = Controller(max_parallel=4)
controller.register("mark", _mark)
pipeline = Pipeline("tagged")
pipeline.node("mark", node_id="one")
pipeline.node("mark", node_id="two")
run = controller.submit_pipeline(pipeline, pipeline_run_id="tagged-1")
run.wait(timeout=_WINDOW)
for node_id, task_id in run.task_ids.items():
snapshot = controller.get(task_id)
assert snapshot is not None
assert snapshot.pipeline_run_id == "tagged-1"
assert task_id == f"tagged-1:{node_id}"
controller.close(timeout=_WINDOW)
def test_cancelling_the_whole_pipeline() -> None:
gate = threading.Event()
controller = Controller(max_parallel=8)
controller.register("wait_gate", _wait_gate)
pipeline = Pipeline("cancelme")
root = pipeline.node("wait_gate", gate, node_id="root")
pipeline.node("wait_gate", gate, node_id="child", depends_on=[root])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p12")
try:
assert _wait_until(lambda: run.node_status("root") is TaskStatus.RUNNING)
assert run.cancel(reason="运维取消") >= 1
outcome = run.wait(timeout=_WINDOW)
assert outcome.status is PipelineStatus.CANCELLED
assert outcome.nodes["root"] is TaskStatus.CANCELLED
assert outcome.nodes["child"] is TaskStatus.CANCELLED
finally:
gate.set()
controller.close(timeout=_WINDOW)
def test_wait_timeout_does_not_change_anything() -> None:
gate = threading.Event()
controller = Controller(max_parallel=4)
controller.register("wait_gate", _wait_gate)
pipeline = Pipeline("slow")
pipeline.node("wait_gate", gate, node_id="root")
run = controller.submit_pipeline(pipeline, pipeline_run_id="p13")
try:
outcome = run.wait(timeout=0.2)
assert outcome.status is PipelineStatus.RUNNING
assert run.node_status("root") in {TaskStatus.STARTING, TaskStatus.RUNNING}
assert not run.settled
finally:
gate.set()
controller.close(timeout=_WINDOW)
def test_pipeline_run_is_registered_on_the_controller() -> None:
controller = Controller(max_parallel=4)
controller.register("mark", _mark)
pipeline = Pipeline("reg")
pipeline.node("mark", node_id="only")
run = controller.submit_pipeline(pipeline, pipeline_run_id="reg-1")
run.wait(timeout=_WINDOW)
assert controller.get_pipeline("reg-1") is run
assert controller.get_pipeline("nope") is None
assert controller.pipelines() == (run,)
controller.close(timeout=_WINDOW)
def test_a_diamond_pipeline_works_end_to_end() -> None:
"""完整形状：一个源 → 两个分支 → 汇合。"""
controller = Controller(max_parallel=8)
controller.register("mark", _mark)
pipeline = Pipeline("diamond")
source = pipeline.node("mark", node_id="source")
left = pipeline.node("mark", node_id="left", depends_on=[source])
right = pipeline.node("mark", node_id="right", depends_on=[source])
pipeline.node("mark", node_id="sink", depends_on=[left, right])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p14")
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
order = _finished_order(controller, "p14")
assert order[0] == "source"
assert order[-1] == "sink"
assert "sink" in render_pipeline(run)
controller.close(timeout=_WINDOW)
def test_two_runs_of_the_same_pipeline_are_independent() -> None:
controller = Controller(max_parallel=4)
controller.register("mark", _mark)
pipeline = Pipeline("reusable")
pipeline.node("mark", node_id="only")
first = controller.submit_pipeline(pipeline, pipeline_run_id="r1")
second = controller.submit_pipeline(pipeline, pipeline_run_id="r2")
assert first.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
assert second.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
assert set(first.task_ids.values()).isdisjoint(second.task_ids.values())
controller.close(timeout=_WINDOW)
def test_pipeline_node_args_reach_the_task() -> None:
controller = Controller(max_parallel=4)
controller.register("add", _add)
pipeline = Pipeline("args")
pipeline.node("add", 20, 22, node_id="sum")
run = controller.submit_pipeline(pipeline, pipeline_run_id="p15")
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
snapshot = controller.get("p15:sum")
assert snapshot is not None
assert snapshot.outcome == {"value": 42}
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 与并发额度协同
# ---------------------------------------------------------------------------
def test_deferred_nodes_do_not_consume_concurrency_slots() -> None:
"""等待上游的节点不该占额度——否则一条长管道会把额度全锁死。"""
gate = threading.Event()
controller = Controller(max_parallel=2)
controller.register("wait_gate", _wait_gate)
controller.register("mark", _mark)
pipeline = Pipeline("slots")
root = pipeline.node("wait_gate", gate, node_id="root")
for index in range(5):
pipeline.node("mark", node_id=f"child{index}", depends_on=[root])
run = controller.submit_pipeline(pipeline, pipeline_run_id="p16")
try:
assert _wait_until(lambda: run.node_status("root") is TaskStatus.RUNNING)
time.sleep(0.1)
assert controller.active_count == 1, "只有 root 该占额度"
assert controller.queued_count == 0, "下游还在 blocked，不是 queued"
gate.set()
assert run.wait(timeout=_WINDOW).status is PipelineStatus.SUCCEEDED
finally:
gate.set()
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 工具
# ---------------------------------------------------------------------------
def _finished_order(controller: Controller, run_id: str) -> list[str]:
"""按结束时间排出某个管道 run 里各节点的完成顺序。"""
snapshots = [
item
for item in controller.list_tasks()
if item.pipeline_run_id == run_id and item.ended_at is not None
]
snapshots.sort(key=lambda item: item.ended_at or item.created_at) # type: ignore[arg-type,return-value]
return [item.name.split("/", 1)[-1] for item in snapshots]
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 _mark(ctx: Any) -> str:
del ctx
return "ok"
def _add(ctx: Any, left: int, right: int) -> int:
del ctx
return left + right
def _boom(ctx: Any) -> str:
del ctx
raise RuntimeError("故意失败")
def _wait_gate(ctx: Any, gate: threading.Event) -> str:
while not gate.wait(0.01):
ctx.checkpoint()
ctx.checkpoint()
return "released"
./tests\unit\test_postgres_lazy.py
"""``PostgresTaskStore`` 的惰性导入契约——**不需要数据库**。
钉住两条承诺：
1. 取到 ``chongming_task.PostgresTaskStore`` 之前，iss 不导入 psycopg——核心包保持
零第三方运行时依赖（``import chongming_task`` 不会因为没装驱动而变慢或失败）；
2. 没装可选依赖时，报错必须**可操作**：给出安装命令，而不是一个光秃秃的
``ModuleNotFoundError``。
真正的行为测试在 ``tests/integration/test_postgres_store.py``（需要 DSN）。
"""
from __future__ import annotations
import os
import subprocess
import sys
import pytest
import chongming_task
def _run(script: str) -> subprocess.CompletedProcess[str]:
environment = {**os.environ}
return subprocess.run(
[sys.executable, "-c", script],
capture_output=True,
text=True,
env=environment,
timeout=60,
)
def test_star_import_does_not_pull_the_driver() -> None:
"""``from chongming_task import *`` 不该要求安装数据库驱动。"""
assert "PostgresTaskStore" not in chongming_task.__all__
script = (
"import sys\n"
"from chongming_task import *\n"
"assert 'psycopg' not in sys.modules, '星号导入拉起了 psycopg'\n"
"assert 'PostgresTaskStore' not in dir(), '它不该出现在星号导入里'\n"
"print('ok')\n"
)
completed = _run(script)
assert completed.returncode == 0, completed.stderr
assert completed.stdout.strip() == "ok"
def test_lazy_attribute_does_not_import_the_driver_until_used() -> None:
script = (
"import sys\n"
"import chongming_task\n"
"assert 'psycopg' not in sys.modules, '导入核心包就拉起了 psycopg'\n"
"print('ok')\n"
)
completed = _run(script)
assert completed.returncode == 0, completed.stderr
assert completed.stdout.strip() == "ok"
def test_missing_driver_gives_an_actionable_error() -> None:
"""把 psycopg 屏蔽成"未安装"，取 ``PostgresTaskStore`` 必须报出安装指引。"""
script = (
"import sys\n"
"sys.modules['psycopg'] = None\n"
"sys.modules['psycopg_pool'] = None\n"
"import chongming_task\n"
"try:\n"
" chongming_task.PostgresTaskStore\n"
"except ImportError as exc:\n"
" text = str(exc)\n"
" assert 'chongming-task[postgres]' in text, text\n"
" print('ok')\n"
"else:\n"
" raise SystemExit('没有抛出 ImportError')\n"
)
completed = _run(script)
assert completed.returncode == 0, completed.stderr
assert completed.stdout.strip() == "ok"
def test_unknown_attribute_still_raises_attribute_error() -> None:
"""``__getattr__`` 只放行已知的可选导出，其余照常抛 AttributeError。"""
with pytest.raises(AttributeError, match="no attribute"):
_ = chongming_task.NoSuchThing # type: ignore[attr-defined]
def test_postgres_store_resolves_when_the_extra_is_installed() -> None:
"""装了可选依赖时必须能取到——并在 stores 包与顶层给出同一个类。"""
pytest.importorskip("psycopg")
from chongming_task.stores import PostgresTaskStore as FromStores
from chongming_task.stores.postgres import PostgresTaskStore as FromModule
assert chongming_task.PostgresTaskStore is FromModule
assert FromStores is FromModule
./tests\unit\test_process_executor.py
"""进程执行器测试。
这些用例全部用**真实子进程**（``spawn``）跑：执行器的全部价值就在于子进程能不能被
停住、能不能被强杀、父进程能不能看见它。用假的进程模型测不出这些。
三条约束写在前面，它们是这套设计的一部分而不是测试的毛病：
* 目标函数必须是**模块级**的——spawn 靠"模块 + 限定名"在子进程里定位它；
* 所以子进程会重新 import 本模块（模块顶层不能有副作用）；
* 每个用例给足超时（5 秒），子进程数量控制在小规模，测试整体仍然是秒级。
"""
from __future__ import annotations
import os
import threading
import time
from typing import Any
import pytest
from chongming_task.errors import ConfigurationError
from chongming_task.models import FailureKind
from chongming_task.runtime.executors.base import ExecutionRequest
from chongming_task.runtime.executors.process import ProcessExecutor
from chongming_task.testing import FakeClock
_WINDOW = 10.0
#: 子进程启动本身就要几百毫秒（新解释器 + import），协作路径的窗口要比线程宽。
_START = 15.0
def _request(target: Any, *args: Any, **kwargs: Any) -> ExecutionRequest:
return ExecutionRequest(
task_id="t-proc",
name="t-proc",
target=target,
args=args,
kwargs=kwargs,
)
# ---------------------------------------------------------------------------
# 子进程里跑的模块级函数（spawn 要求）
# ---------------------------------------------------------------------------
def _add(ctx: Any, left: int, right: int, *, scale: int = 1) -> int:
"""两数相加。"""
del ctx
return (left + right) * scale
def _boom(ctx: Any) -> None:
"""抛异常。"""
del ctx
raise ValueError("子进程里炸了")
def _describe(ctx: Any) -> dict[str, Any]:
"""把注入进来的上下文信息带回来，验证它是真的被注入了。"""
return {
"task_id": ctx.task_id,
"name": ctx.name,
"attempt": ctx.attempt,
"pid": os.getpid(),
"has_token": ctx.token is not None,
}
def _beater(ctx: Any) -> str:
"""持续 checkpoint，直到进程退出。"""
for _ in range(10_000):
ctx.checkpoint()
time.sleep(0.01)
return "never"
def _bounded_beater(ctx: Any) -> str:
"""有限次 checkpoint 后正常返回——用于验证"暂停没生效时不能把它卡住"。"""
for _ in range(5):
ctx.checkpoint()
time.sleep(0.1)
return "done"
def _late_checkpointer(ctx: Any) -> str:
"""先干一段**不 checkpoint** 的活，之后才正常推进。
用来测暂停回滚：暂停请求会在它到达第一个检查点之前就超时。
"""
time.sleep(0.5)
for _ in range(5):
ctx.checkpoint()
time.sleep(0.05)
return "done"
def _progress_reporter(ctx: Any) -> str:
"""报进度，父进程应当收到。"""
for index in range(1, 4):
ctx.progress(index, total=3, unit="批")
time.sleep(0.02)
return "ok"
def _ignore_everything(ctx: Any) -> str:
"""**不**调用 checkpoint 的死循环：只能靠强杀停下。"""
del ctx
end = time.monotonic() + 60
while time.monotonic() < end:
pass
return "never"
def _abrupt_exit(ctx: Any) -> str:
"""直接消失，连一句结论都不留。"""
del ctx
os._exit(3)
return "never" # pragma: no cover
def _unpicklable_result(ctx: Any) -> Any:
"""返回值不可 pickle（线程执行器能接住它，进程执行器不行）。"""
del ctx
return lambda: None
def _explode_on_unpickle() -> None: # pragma: no cover - 只在子进程里被调用
raise RuntimeError("反序列化时炸了")
class _CannotCrossTheBorder:
"""能 pickle 出去、却在子进程里反序列化失败的对象。
``__reduce__`` 只在**加载**时调用那个工厂函数，所以父进程的预检（只 dump）会放行
——正好用来测"子进程根本没跑起来"那条分支。
"""
def __reduce__(self) -> Any:
return (_explode_on_unpickle, ())
# ---------------------------------------------------------------------------
# 载荷校验：把"子进程里神秘失败"变成一句能照做的话
# ---------------------------------------------------------------------------
def test_a_lambda_target_is_rejected_with_an_actionable_error() -> None:
with pytest.raises(ConfigurationError, match="可 pickle"):
ProcessExecutor().launch(_request(lambda ctx: "nope"))
def test_a_closure_target_is_rejected() -> None:
captured = "closure"
def inner(ctx: Any) -> str:
del ctx
return captured
with pytest.raises(ConfigurationError, match="ThreadExecutor"):
ProcessExecutor().launch(_request(inner))
def test_an_unpicklable_argument_is_rejected() -> None:
with pytest.raises(ConfigurationError, match="参数"):
ProcessExecutor().launch(_request(_add, threading.Lock(), 1))
# ---------------------------------------------------------------------------
# 正常路径
# ---------------------------------------------------------------------------
def test_target_runs_in_a_child_process_and_returns_its_value() -> None:
handle = ProcessExecutor().launch(
ExecutionRequest(
task_id="t-args",
name="t-args",
target=_add,
args=(2, 3),
kwargs={"scale": 10},
)
)
outcome = handle.wait(timeout=_START)
assert outcome is not None
assert outcome.ok
assert outcome.value == 50
assert outcome.abandoned is False, "进程没有强杀这个词，结论是已知的"
def test_the_child_has_its_own_pid_and_gets_the_injected_context() -> None:
handle = ProcessExecutor().launch(_request(_describe))
outcome = handle.wait(timeout=_START)
assert outcome is not None and outcome.ok
payload = outcome.value
assert payload["task_id"] == "t-proc"
assert payload["name"] == "t-proc"
assert payload["has_token"] is True
assert payload["pid"] == handle.pid, "子进程里看到的 pid 就是句柄报出来的那个"
assert handle.pid != os.getpid(), "任务确实不在监督进程里跑"
def test_the_handle_knows_the_child_pid() -> None:
handle = ProcessExecutor().launch(_request(_add, 1, 1))
assert handle.pid is not None
assert handle.running is True
assert handle.wait(timeout=_START) is not None
assert handle.running is False, "有结论之后就不该再报'在运行'"
def test_a_task_failure_crosses_the_boundary_as_a_structured_failure() -> None:
outcome = ProcessExecutor().launch(_request(_boom)).wait(timeout=_START)
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.EXCEPTION
assert "子进程里炸了" in outcome.failure.message
assert outcome.failure.traceback, "子进程里的 traceback 必须带回来"
def test_an_unpicklable_result_is_reported_as_a_result_problem() -> None:
"""返回值过不去，报的是"产出接不住"，不是"任务出错了"——两者的修法完全不同。"""
outcome = ProcessExecutor().launch(_request(_unpicklable_result)).wait(timeout=_START)
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.RESULT_INVALID
assert outcome.failure.retryable is False, "同一个返回值重跑多少次都还是过不去"
# ---------------------------------------------------------------------------
# 心跳：idle 超时靠它，而它必须真的跨过进程边界
# ---------------------------------------------------------------------------
def test_heartbeats_cross_the_process_boundary() -> None:
"""把父进程的时钟推远，再等子进程的一次 heartbeat 把它拉回来。
用假时钟是为了**确定性地**区分"子进程真的报了心跳"与"父进程自己刚记了一次"。
"""
clock = FakeClock()
handle = ProcessExecutor(clock=clock).launch(_request(_beater))
try:
assert handle.wait(timeout=0.2) is None, "子进程不该这么快结束"
clock.advance(300)
assert handle.heartbeat.idle_seconds() >= 300
deadline = time.monotonic() + _START
while handle.heartbeat.idle_seconds() >= 300 and time.monotonic() < deadline:
time.sleep(0.02)
assert handle.heartbeat.idle_seconds() == 0, "子进程的 checkpoint 应当刷新心跳"
finally:
handle.stop(graceful_timeout=_WINDOW)
def test_progress_reports_are_forwarded_to_the_parent() -> None:
seen: list[Any] = []
request = ExecutionRequest(
task_id="t-progress",
name="t-progress",
target=_progress_reporter,
on_progress=seen.append,
)
outcome = ProcessExecutor().launch(request).wait(timeout=_START)
assert outcome is not None and outcome.ok
assert [item.done for item in seen] == [1, 2, 3]
assert seen[-1].unit == "批"
# ---------------------------------------------------------------------------
# 暂停：确认了才算，没确认就必须回滚
# ---------------------------------------------------------------------------
def test_pause_suspends_the_child_and_resume_lets_it_finish() -> None:
handle = ProcessExecutor().launch(_request(_bounded_beater))
assert handle.pause(timeout=_START) is True, "子进程应当在下一次 checkpoint 停下"
# 若没真的暂停，这个任务 0.5 秒内就该跑完了。
time.sleep(0.8)
assert handle.wait(timeout=0) is None, "被暂停的任务不该继续推进"
assert handle.running is True
handle.resume()
outcome = handle.wait(timeout=_START)
assert outcome is not None and outcome.ok
assert outcome.value == "done"
def test_pause_rolls_back_when_the_child_never_reaches_a_checkpoint() -> None:
"""暂停请求超时后必须**撤销**。
不撤销的话，调用方已经按"暂停没生效"把状态回滚成 RUNNING，而子进程会在稍后的
某个检查点自己停住：那是一个没人知道、也没人会恢复的挂起——比暂停失败更糟。
"""
handle = ProcessExecutor().launch(_request(_late_checkpointer))
assert handle.wait(timeout=0.1) is None, "它还在那段不 checkpoint 的活里"
assert handle.pause(timeout=0.2) is False, "窗口内它压根没到检查点"
outcome = handle.wait(timeout=_START)
assert outcome is not None, "回滚之后它必须继续跑完，而不是停在下一个检查点上"
assert outcome.ok, f"不该被自己撤销掉的暂停影响：{outcome.failure}"
# ---------------------------------------------------------------------------
# 停止：进程执行器存在的理由
# ---------------------------------------------------------------------------
def test_cooperative_stop_ends_the_task_without_killing_it() -> None:
handle = ProcessExecutor().launch(_request(_beater))
assert handle.wait(timeout=0.3) is None # 先让它跑起来
outcome = handle.stop(graceful_timeout=_WINDOW)
assert not outcome.ok
assert outcome.abandoned is False, "它是自己收尾的，不是被放弃"
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
assert handle.running is False
def test_a_task_that_never_checks_the_token_is_actually_killed() -> None:
"""**这是进程执行器唯一不可替代的能力。** 线程只能放弃等待并记 LOST；
子进程会真的被停掉，结论因此是已知的。"""
handle = ProcessExecutor().launch(_request(_ignore_everything))
pid = handle.pid
assert pid is not None
assert handle.wait(timeout=0.3) is None # 它连 checkpoint 都不调，只能强杀
outcome = handle.stop(graceful_timeout=0.3)
assert not outcome.ok
assert outcome.abandoned is False, "任务确实停了，不是'不知道停没停'"
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
assert handle.running is False, "载体已经死了，不能谎报还在跑"
assert handle.process is not None and handle.process.exitcode is not None
def test_stop_is_also_delivered_through_the_handle_token() -> None:
"""控制面有几条路径直接调 ``handle.token.request_stop()``，绕过了 ``stop()``。
令牌不转发的话子进程永远收不到那次停止——表现出来就是"停止按钮没反应"。
"""
handle = ProcessExecutor().launch(_request(_beater))
assert handle.wait(timeout=0.3) is None
handle.token.request_stop("走令牌这条路")
outcome = handle.wait(timeout=_WINDOW)
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
# ---------------------------------------------------------------------------
# 崩溃隔离：父进程必须活着，并且如实说清楚发生了什么
# ---------------------------------------------------------------------------
def test_an_abrupt_child_exit_is_reported_with_its_exit_code() -> None:
outcome = ProcessExecutor().launch(_request(_abrupt_exit)).wait(timeout=_START)
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.NONZERO_EXIT
assert outcome.failure.exit_code == 3, "退出码是唯一能解释'它为什么没了'的事实"
assert "副作用" in outcome.failure.message, "结果未知必须说清楚，别让人放心重跑"
def test_a_child_that_never_boots_is_reported_distinctly() -> None:
"""载荷在子进程里反序列化失败：任务**从未开始**，与"跑一半没了"是两回事。"""
handle = ProcessExecutor().launch(_request(_add, _CannotCrossTheBorder(), 1))
outcome = handle.wait(timeout=_START)
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.NONZERO_EXIT
assert "从未开始" in outcome.failure.message
./tests\unit\test_public_api.py
"""公开 API 的表面测试。
``chongming_task`` 的 ``__init__`` 是**嵌入其他应用**的入口，因此它本身的正确性
值得一条用例：漏导出某个名字、或者导出顺序导致循环导入，都会让集成方第一步就失败。
"""
from __future__ import annotations
import chongming_task
def test_version_is_exposed() -> None:
assert chongming_task.__version__
def test_all_has_no_duplicates() -> None:
"""重复导出会让"到底哪个生效"变得不确定。
顺序**不在这里**断言：``__all__`` 的排序由 ruff 的 ``RUF022``（isort 风格，
全大写常量优先）负责，那是 lint 的职责。写一条与 linter 规则不一致的断言，
只会让两边互相打架。
"""
names = chongming_task.__all__
assert len(names) == len(set(names))
def test_every_exported_name_resolves() -> None:
"""``__all__`` 里写了却取不到，等于给集成方埋了一个 ImportError。"""
missing = [name for name in chongming_task.__all__ if not hasattr(chongming_task, name)]
assert missing == []
def test_primary_integration_surface_is_importable() -> None:
"""这一段就是 README 里"嵌入你自己的应用"的最小示例。"""
from chongming_task import ( # noqa: F401 - 只为验证可导入
Controller,
InMemoryTaskStore,
SqliteTaskStore,
TaskStatus,
TaskStore,
run_task,
)
controller = Controller(max_parallel=2)
assert controller.registry.names() == ()
assert TaskStatus.PENDING.value == "pending"
assert callable(run_task)
def test_store_implementations_satisfy_the_protocol() -> None:
"""存储实现必须满足契约——第三方后端也照这条自证。"""
from chongming_task import InMemoryTaskStore, SqliteTaskStore, TaskStore
assert isinstance(InMemoryTaskStore(), TaskStore)
store = SqliteTaskStore(":memory:")
try:
assert isinstance(store, TaskStore)
finally:
# 不关掉连接会在解释器退出时触发 ResourceWarning；本项目的 pytest 配置
# 把 warning 升级为 error，因此泄漏是会被抓住的。
store.close()
def test_controller_accepts_a_store() -> None:
"""存储已经接线。
这条用例原先钉的是"尚未接线"——当 ``Controller(store=...)`` 真的开始工作时，
它会自己失败并提示删掉。现在它钉的是接线本身。
"""
from chongming_task import Controller, InMemoryTaskStore
store = InMemoryTaskStore()
controller = Controller(store=store, max_parallel=2)
assert controller.store is store
assert controller.persist_failures == ()
assert controller.owner
def test_private_helpers_are_not_exported() -> None:
"""以下划线开头的实现细节不应出现在公开表面上。"""
leaked = [name for name in chongming_task.__all__ if name.startswith("_")]
assert leaked == []
./tests\unit\test_run.py
"""单任务受控执行的测试。
覆盖三件用户最关心的事：**失败会怎样**（重试 / 超时 / 最终失败）、**我能不能干预**
（停止 / 暂停 / 恢复），以及**结果是否诚实**（不配合的任务落 LOST 而不是"已停止"）。
"""
from __future__ import annotations
import threading
import time
from typing import Any
import pytest
from chongming_task.errors import ConfigurationError
from chongming_task.models import FailureKind
from chongming_task.policies import FixedBackoff, RetryPolicy, Timeout
from chongming_task.run import TaskRun, run_task
from chongming_task.runtime.executors.inline import InlineExecutor
from chongming_task.states import TaskStatus
_WINDOW = 10.0
def _wait_for_status(run: TaskRun, status: TaskStatus, timeout: float = _WINDOW) -> bool:
"""等状态就位。
比 ``sleep`` 可靠：``start()`` 是异步的，"任务已经开始执行"与"状态机已迁移到
RUNNING"之间存在极短窗口。
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
if run.status is status:
return True
time.sleep(0.01)
return False
# ---------------------------------------------------------------------------
# 成功路径
# ---------------------------------------------------------------------------
def test_successful_run_returns_value_and_succeeds() -> None:
outcome = run_task(lambda ctx, a, b: a + b, 2, 3, task_id="t-ok", name="加法")
assert outcome.ok
assert outcome.status is TaskStatus.SUCCEEDED
assert outcome.value == 5
assert outcome.failure is None
assert outcome.snapshot.ended_at is not None
def test_transition_history_is_complete_and_ordered() -> None:
"""状态历史只追加，且每条都带触发源——这是可审计性的基础。"""
run = TaskRun(lambda ctx: 1, task_id="t-hist", name="demo")
run.start()
assert run.wait(timeout=_WINDOW).ok
assert [t.to_status for t in run.transitions] == [
TaskStatus.QUEUED,
TaskStatus.STARTING,
TaskStatus.RUNNING,
TaskStatus.SUCCEEDED,
]
assert all(t.trigger for t in run.transitions)
assert run.transitions[-1].task_id == "t-hist"
def test_exception_is_captured_with_traceback() -> None:
def boom(ctx: Any) -> None:
raise ValueError("炸了")
outcome = run_task(boom, task_id="t-boom")
assert outcome.status is TaskStatus.FAILED
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.EXCEPTION
assert outcome.failure.traceback is not None
assert "boom" in outcome.failure.traceback
# ---------------------------------------------------------------------------
# 重试
# ---------------------------------------------------------------------------
def test_retries_until_success_and_counts_attempts() -> None:
attempts: list[int] = []
def flaky(ctx: Any, fail_times: int) -> str:
attempts.append(1)
if len(attempts) <= fail_times:
raise RuntimeError("暂时性故障")
return "ok"
outcome = run_task(
flaky,
2,
retry=RetryPolicy(max_attempts=3, backoff=FixedBackoff(0.0)),
task_id="t-flaky",
)
assert outcome.ok
assert outcome.snapshot.attempt == 3
assert len(attempts) == 3
def test_retry_exhaustion_lands_failed() -> None:
def always_fails(ctx: Any) -> None:
raise RuntimeError("一直失败")
outcome = run_task(
always_fails,
retry=RetryPolicy(max_attempts=2, backoff=FixedBackoff(0.0)),
task_id="t-always",
)
assert outcome.status is TaskStatus.FAILED
assert outcome.snapshot.attempt == 2
assert outcome.failure is not None
def test_retry_wait_is_visible_in_history() -> None:
seen: list[int] = []
def flaky(ctx: Any) -> str:
if not seen:
seen.append(1)
raise RuntimeError("第一次失败")
return "ok"
run = TaskRun(
flaky,
retry=RetryPolicy(max_attempts=2, backoff=FixedBackoff(0.0)),
task_id="t-wait",
)
run.start()
assert run.wait(timeout=_WINDOW).ok
assert TaskStatus.RETRY_WAIT in [t.to_status for t in run.transitions]
# ---------------------------------------------------------------------------
# 超时
# ---------------------------------------------------------------------------
def test_per_attempt_timeout_stops_a_slow_task() -> None:
def slow(ctx: Any) -> str:
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.01)
return "永远跑不完"
outcome = run_task(slow, timeout=Timeout(per_attempt=0.3), task_id="t-slow")
assert outcome.status is TaskStatus.FAILED
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.TIMEOUT_ATTEMPT
def test_idle_timeout_is_skipped_while_paused() -> None:
"""被暂停的任务不推进心跳，若照常判 idle 超时，暂停就会被误判成卡死。"""
started = threading.Event()
def work(ctx: Any) -> str:
started.set()
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.02)
return "done"
run = TaskRun(
work,
timeout=Timeout(idle=0.2),
task_id="t-idle-paused",
graceful_stop=2.0,
)
run.start()
assert started.wait(_WINDOW)
assert _wait_for_status(run, TaskStatus.RUNNING)
assert run.pause(timeout=_WINDOW) is True
# 暂停期间超过 idle 阈值，但不应被判超时。
time.sleep(0.5)
assert run.status is TaskStatus.PAUSED
run.resume()
run.stop(reason="测试结束")
assert run.wait(timeout=_WINDOW).status is TaskStatus.STOPPED
# ---------------------------------------------------------------------------
# 停止
# ---------------------------------------------------------------------------
def test_stop_lands_stopped_and_is_auditable() -> None:
started = threading.Event()
def work(ctx: Any) -> str:
started.set()
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.005)
return "done"
run = TaskRun(work, task_id="t-stop")
run.start()
assert started.wait(_WINDOW)
assert _wait_for_status(run, TaskStatus.RUNNING)
run.stop(reason="发布窗口")
outcome = run.wait(timeout=_WINDOW)
assert outcome.status is TaskStatus.STOPPED
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
assert outcome.failure.retryable is False, "取消不是失败，不能重试"
stopping = [t for t in run.transitions if t.to_status is TaskStatus.STOPPING]
assert stopping, "应当能看到 STOPPING 这个中间态"
assert stopping[0].reason == "发布窗口"
def test_uncooperative_task_is_reported_as_lost_not_stopped() -> None:
"""线程停不掉时结果**未知**：任务可能仍在后台产生副作用。
记成"已停止"会让运维误判资源已释放，因此必须落 LOST。
"""
def stubborn(ctx: Any) -> str:
for _ in range(500):
time.sleep(0.01)
return "ignored"
run = TaskRun(stubborn, task_id="t-stubborn", graceful_stop=0.2)
run.start()
time.sleep(0.3)
run.stop(reason="不配合")
outcome = run.wait(timeout=_WINDOW)
assert outcome.status is TaskStatus.LOST
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.LOST
def test_stop_before_start_lands_stopped_without_burning_an_attempt() -> None:
"""都还没跑就要求停止，不该白白消耗一次尝试。"""
run = TaskRun(lambda ctx: "ok", task_id="t-pre-stop")
run.stop(reason="还没跑就不要了")
run.start()
outcome = run.wait(timeout=_WINDOW)
assert outcome.status is TaskStatus.STOPPED
assert outcome.snapshot.attempt == 0
assert all(t.to_status is not TaskStatus.RUNNING for t in run.transitions)
# ---------------------------------------------------------------------------
# 暂停与恢复
# ---------------------------------------------------------------------------
def test_pause_freezes_progress_and_resume_continues() -> None:
started = threading.Event()
progress: list[int] = []
def work(ctx: Any) -> str:
started.set()
for index in range(1_000_000):
ctx.checkpoint()
progress.append(index)
time.sleep(0.005)
return "done"
run = TaskRun(work, task_id="t-pause")
run.start()
assert started.wait(_WINDOW)
assert _wait_for_status(run, TaskStatus.RUNNING)
assert run.pause(timeout=_WINDOW) is True
assert run.status is TaskStatus.PAUSED
frozen = len(progress)
time.sleep(0.15)
assert len(progress) == frozen, "暂停生效后不应继续推进"
run.resume()
assert run.status is TaskStatus.RUNNING, "就地暂停的任务恢复后回到 RUNNING"
time.sleep(0.15)
assert len(progress) > frozen, "恢复后应继续推进"
run.stop()
assert run.wait(timeout=_WINDOW).status is TaskStatus.STOPPED
def test_pause_returns_false_when_task_never_checkpoints() -> None:
"""暂停没生效就必须如实回滚为 RUNNING，不能假装已暂停。"""
def work(ctx: Any) -> str:
for _ in range(500):
time.sleep(0.01)
return "ignored"
run = TaskRun(work, task_id="t-nopause", graceful_stop=0.2)
run.start()
assert _wait_for_status(run, TaskStatus.RUNNING)
assert run.pause(timeout=0.2) is False
assert run.status is TaskStatus.RUNNING
assert any(t.to_status is TaskStatus.PAUSING for t in run.transitions)
assert any(
t.to_status is TaskStatus.RUNNING and t.trigger == "runtime:pause_aborted"
for t in run.transitions
), "回滚应当留下 PAUSE_ABORTED 记录"
run.stop()
run.wait(timeout=_WINDOW)
def test_pause_on_a_finished_run_is_rejected() -> None:
run = TaskRun(lambda ctx: 1, task_id="t-pause-late")
run.start()
assert run.wait(timeout=_WINDOW).ok
assert run.pause(timeout=0.01) is False
assert run.status is TaskStatus.SUCCEEDED
# ---------------------------------------------------------------------------
# wait 的契约
# ---------------------------------------------------------------------------
def test_wait_timeout_raises_and_does_not_change_state() -> None:
"""把"我没等到"当成"任务失败"是最常见的误用，因此这里只抛异常。"""
release = threading.Event()
def work(ctx: Any) -> str:
release.wait(_WINDOW)
return "ok"
run = TaskRun(work, task_id="t-wait")
run.start()
with pytest.raises(TimeoutError, match="状态未被改变"):
run.wait(timeout=0.2)
assert run.status in (TaskStatus.STARTING, TaskStatus.RUNNING)
release.set()
assert run.wait(timeout=_WINDOW).ok
def test_wait_before_start_is_an_error() -> None:
run = TaskRun(lambda ctx: 1, task_id="t-not-started")
with pytest.raises(RuntimeError, match="start"):
run.wait()
def test_starting_twice_is_an_error() -> None:
run = TaskRun(lambda ctx: 1, task_id="t-twice")
run.start()
run.wait(timeout=_WINDOW)
with pytest.raises(RuntimeError, match="已经启动过"):
run.start()
# ---------------------------------------------------------------------------
# 执行器可替换
# ---------------------------------------------------------------------------
def test_executor_is_pluggable() -> None:
run = TaskRun(
lambda ctx: "inline",
task_id="t-inline",
executor=InlineExecutor(),
)
run.start()
outcome = run.wait(timeout=_WINDOW)
assert outcome.ok
assert outcome.value == "inline"
assert outcome.snapshot.executor == "inline"
class _RefusingExecutor:
"""拉不起来任何任务，只会抛错——进程执行器的"目标不可 pickle"就长这样。"""
name = "refusing"
def launch(self, request: Any) -> Any:
del request
raise ConfigurationError("目标不可 pickle：这类任务请改用线程执行器")
def test_a_launch_failure_is_recorded_as_the_tasks_own_failure() -> None:
"""执行器拉不起来时，必须如实记在**任务**上，而不是让异常逃进监督线程。
逃进去的后果有两条：任务先停在 ``starting``，最后由兜底说成"监督循环内部错误"
——那等于把一个任务侧的问题指成框架的 bug；而调用方拿到的诊断里也看不到真正
该改的东西。
"""
outcome = run_task(
lambda ctx: "never",
task_id="t-refused",
executor=_RefusingExecutor(),
)
assert outcome.status is TaskStatus.FAILED
assert outcome.failure is not None
assert "不可 pickle" in outcome.failure.message
assert "监督循环内部错误" not in outcome.failure.message, "别把任务侧的问题说成框架的"
assert outcome.failure.retryable is False, "配置类错误重试多少次都是同一个结果"
./tests\unit\test_states.py
"""状态表的结构性不变量。
这些测试只约束 :mod:`chongming_task.states` 的数据本身，不依赖状态机实现。
它们的作用是：让"转移表是唯一真相"成为可执行的断言，而不是文档里的承诺。
"""
from __future__ import annotations
import pytest
from chongming_task.states import (
NON_TERMINAL,
REVIVABLE,
TERMINAL,
TRANSITIONS,
TaskStatus,
allowed_from,
can_transition,
is_terminal,
)
def test_transition_table_is_exhaustive() -> None:
"""每个规范状态都必须有表项，否则状态机会在运行期 KeyError。"""
assert set(TRANSITIONS) == set(TaskStatus)
def test_alias_does_not_add_a_table_key() -> None:
"""``FINISHED`` 只是别名，不应在表中占据独立条目。"""
assert len(TRANSITIONS) == len(set(TaskStatus))
def test_terminal_set_is_derived_from_the_table() -> None:
"""终态必须由表推导，避免"表"与"集合"两处维护而不一致。"""
derived = frozenset(status for status, nxt in TRANSITIONS.items() if not nxt)
assert derived == TERMINAL
def test_terminal_and_non_terminal_partition_the_enum() -> None:
partition = TERMINAL | NON_TERMINAL
assert not (TERMINAL & NON_TERMINAL)
assert partition == set(TaskStatus)
@pytest.mark.parametrize("status", sorted(TERMINAL))
def test_terminal_states_have_no_outgoing_transitions(status: TaskStatus) -> None:
assert allowed_from(status) == frozenset()
for target in TaskStatus:
assert not can_transition(status, target)
@pytest.mark.parametrize("status", sorted(NON_TERMINAL))
def test_non_terminal_states_have_outgoing_transitions(status: TaskStatus) -> None:
"""非终态若无出边，任务会永久卡死。"""
assert allowed_from(status), f"{status} 必须有至少一条出边"
def test_no_self_loops() -> None:
"""自环会掩盖状态机缺陷，并让 status_version 无意义地增长。"""
for source, targets in TRANSITIONS.items():
assert source not in targets, f"{source} 存在自环"
def test_all_targets_are_known_statuses() -> None:
known = set(TaskStatus)
for targets in TRANSITIONS.values():
assert targets <= known
@pytest.mark.parametrize("start", sorted(NON_TERMINAL))
def test_every_non_terminal_state_can_reach_a_terminal_state(start: TaskStatus) -> None:
"""不存在只能内部互跳、永远收敛不了的死角。"""
seen: set[TaskStatus] = set()
stack: list[TaskStatus] = [start]
while stack:
current = stack.pop()
if current in TERMINAL:
return
if current in seen:
continue
seen.add(current)
stack.extend(allowed_from(current))
pytest.fail(f"{start} 无法到达任何终态")
# ---------------------------------------------------------------------------
# 以下用例锁定设计决策本身，防止后续重构悄悄改变语义
# ---------------------------------------------------------------------------
def test_paused_can_resume_either_in_place_or_back_to_the_queue() -> None:
"""恢复目标取决于暂停发生在哪一层：
* 就地暂停（经 ``PAUSING``）→ ``RUNNING``，槽位从未释放；
* 未启动即暂停 → ``QUEUED``，重新排队。
"""
assert can_transition(TaskStatus.PAUSED, TaskStatus.RUNNING)
assert can_transition(TaskStatus.PAUSED, TaskStatus.QUEUED)
def test_stopping_may_still_succeed() -> None:
"""停止请求发出后任务恰好自然结束，是合法结果，不能判为失败。"""
assert can_transition(TaskStatus.STOPPING, TaskStatus.SUCCEEDED)
def test_pending_cannot_jump_straight_to_running() -> None:
"""PENDING 必须经过 QUEUED/STARTING，不允许绕过并发槽与执行器启动。"""
assert not can_transition(TaskStatus.PENDING, TaskStatus.RUNNING)
def test_running_cannot_return_to_queued_directly() -> None:
"""RUNNING 回落 QUEUED 只能经由 PAUSED，否则会绕过暂停语义。"""
assert not can_transition(TaskStatus.RUNNING, TaskStatus.QUEUED)
def test_revivable_is_a_subset_of_terminal() -> None:
assert REVIVABLE <= TERMINAL
def test_success_and_skip_are_not_revivable() -> None:
"""成功重跑没有意义；SKIPPED 是策略决定的结果，不应被 retry 复活。"""
assert TaskStatus.SUCCEEDED not in REVIVABLE
assert TaskStatus.SKIPPED not in REVIVABLE
def test_legacy_finished_alias() -> None:
assert TaskStatus.FINISHED is TaskStatus.SUCCEEDED
assert "finished" not in {s.value for s in TaskStatus}
def test_helpers_agree_with_the_table() -> None:
assert is_terminal(TaskStatus.SUCCEEDED)
assert not is_terminal(TaskStatus.RUNNING)
assert can_transition(TaskStatus.RUNNING, TaskStatus.STOPPING)
assert not can_transition(TaskStatus.SUCCEEDED, TaskStatus.RUNNING)
./tests\unit\test_state_machine.py
"""状态机行为测试。
覆盖三类东西：
1. 转移表被**真正强制执行**，而不是只写在文档里；
2. 控制动词语义（stop / pause / resume / retry / reset / force）与设计文档一致；
3. 审计记录与版本号的不变量。
"""
from __future__ import annotations
import pytest
from chongming_task.core.state_machine import StateMachine, Trigger
from chongming_task.errors import AttemptsExhausted, InvalidTransition
from chongming_task.models import Failure, FailureKind, Progress, RequestedAction
from chongming_task.states import TaskStatus
from chongming_task.testing import FakeClock
@pytest.fixture
def clock() -> FakeClock:
return FakeClock()
@pytest.fixture
def machine(clock: FakeClock) -> StateMachine:
return StateMachine.create("t1", "抽取", max_attempts=3, clock=clock)
def _run_to_running(machine: StateMachine, executor: str = "inline") -> None:
machine.mark_queued()
machine.mark_starting(executor=executor)
machine.mark_running()
# ---------------------------------------------------------------------------
# 创建与转移表强制
# ---------------------------------------------------------------------------
def test_create_starts_pending(machine: StateMachine, clock: FakeClock) -> None:
assert machine.status is TaskStatus.PENDING
assert machine.snapshot.created_at == clock.now
assert machine.snapshot.attempt == 0
assert machine.history == ()
def test_illegal_transition_raises_and_leaves_state_untouched(
machine: StateMachine,
) -> None:
before = machine.snapshot
with pytest.raises(InvalidTransition):
machine.mark_running()
assert machine.snapshot is before
assert machine.history == ()
def test_invalid_transition_message_lists_allowed_targets(machine: StateMachine) -> None:
with pytest.raises(InvalidTransition) as info:
machine.mark_succeeded()
assert info.value.current is TaskStatus.PENDING
assert info.value.requested is TaskStatus.SUCCEEDED
assert "非法状态转移" in str(info.value)
assert "queued" in str(info.value)
def test_happy_path_pending_to_succeeded(machine: StateMachine, clock: FakeClock) -> None:
_run_to_running(machine, executor="thread")
record = machine.mark_succeeded(outcome={"rows": 10})
assert record.from_status is TaskStatus.RUNNING
assert record.to_status is TaskStatus.SUCCEEDED
assert machine.snapshot.outcome == {"rows": 10}
assert machine.snapshot.ended_at == clock.now
assert machine.snapshot.attempt == 1
assert machine.snapshot.validate() == ()
assert [t.to_status for t in machine.history] == [
TaskStatus.QUEUED,
TaskStatus.STARTING,
TaskStatus.RUNNING,
TaskStatus.SUCCEEDED,
]
def test_status_version_increases_by_one_per_transition(machine: StateMachine) -> None:
_run_to_running(machine)
assert [t.version for t in machine.history] == [1, 2, 3]
assert machine.snapshot.status_version == 3
def test_terminal_state_cannot_transition_at_all(machine: StateMachine) -> None:
_run_to_running(machine)
machine.mark_succeeded()
for action in (
machine.mark_queued,
machine.mark_blocked,
machine.mark_running,
machine.mark_cancelled,
machine.mark_lost,
machine.mark_skipped,
):
with pytest.raises(InvalidTransition):
action()
assert machine.status is TaskStatus.SUCCEEDED
# ---------------------------------------------------------------------------
# 高频旁路更新
# ---------------------------------------------------------------------------
def test_heartbeat_and_progress_do_not_touch_status_version(
machine: StateMachine, clock: FakeClock
) -> None:
"""心跳是高频写入；若参与乐观并发，会与状态写入持续冲突。"""
machine.mark_queued()
before = machine.snapshot.status_version
clock.advance(5)
machine.heartbeat()
machine.update_progress(Progress(done=1, total=4))
assert machine.snapshot.status_version == before
assert machine.snapshot.last_heartbeat_at == clock.now
assert machine.snapshot.progress is not None
assert machine.snapshot.progress.done == 1
def test_transitions_use_injected_clock(machine: StateMachine, clock: FakeClock) -> None:
machine.mark_queued()
queued_at = machine.snapshot.queued_at
clock.advance(30)
machine.mark_starting(executor="inline")
assert machine.history[-1].at == clock.now
assert queued_at != machine.history[-1].at
# ---------------------------------------------------------------------------
# 重试预算
# ---------------------------------------------------------------------------
def test_attempt_is_consumed_only_when_starting(machine: StateMachine) -> None:
"""预算只在"真正要跑"的那一刻消费。"""
machine.mark_queued()
assert machine.snapshot.attempt == 0
machine.mark_starting(executor="inline")
assert machine.snapshot.attempt == 1
def test_retry_revives_from_failed_and_keeps_history(clock: FakeClock) -> None:
machine = StateMachine.create("t1", "x", max_attempts=3, clock=clock)
_run_to_running(machine)
machine.mark_failed(Failure(kind=FailureKind.EXCEPTION, message="boom"))
record = machine.retry()
assert record.trigger == Trigger.RETRY_REQUESTED
assert machine.status is TaskStatus.QUEUED
assert machine.snapshot.failure is None
assert machine.snapshot.ended_at is None
assert machine.snapshot.attempt == 1 # 历史不被改写
machine.mark_starting(executor="inline")
assert machine.snapshot.attempt == 2
def test_retry_rejected_when_budget_exhausted(clock: FakeClock) -> None:
machine = StateMachine.create("t1", "x", max_attempts=1, clock=clock)
_run_to_running(machine)
machine.mark_failed(Failure(kind=FailureKind.EXCEPTION, message="boom"))
with pytest.raises(AttemptsExhausted):
machine.retry()
def test_forced_retry_raises_the_budget(clock: FakeClock) -> None:
"""force=True 是"我知道它已失败够多次了，再试一次"。"""
machine = StateMachine.create("t1", "x", max_attempts=1, clock=clock)
_run_to_running(machine)
machine.mark_failed(Failure(kind=FailureKind.EXCEPTION, message="boom"))
machine.retry(force=True)
assert machine.snapshot.max_attempts == 2
machine.mark_starting(executor="inline")
assert machine.snapshot.attempt == 2
def test_retry_rejected_from_non_revivable_terminal(clock: FakeClock) -> None:
"""SUCCEEDED / SKIPPED 不可复活：前者重跑无意义，后者是策略决定的结果。"""
machine = StateMachine.create("t1", "x", clock=clock)
_run_to_running(machine)
machine.mark_succeeded()
assert not machine.can_retry
with pytest.raises(InvalidTransition):
machine.retry()
def test_retry_is_not_the_path_for_retry_wait(clock: FakeClock) -> None:
"""RETRY_WAIT 已安排好重试，不应再被 retry() 接受，否则会双倍消费预算。"""
machine = StateMachine.create("t1", "x", max_attempts=3, clock=clock)
_run_to_running(machine)
machine.mark_retry_wait(Failure(kind=FailureKind.EXCEPTION, message="boom"))
with pytest.raises(InvalidTransition):
machine.retry()
# ---------------------------------------------------------------------------
# 停止语义
# ---------------------------------------------------------------------------
def test_stop_on_unstarted_task_lands_stopped_without_buffer(
machine: StateMachine,
) -> None:
"""还没跑起来的任务应当立即落 STOPPED；显示"停止中"是虚假信息。"""
record = machine.request_stop(actor="ops")
assert record.to_status is TaskStatus.STOPPED
assert machine.snapshot.requested_action is None
assert machine.snapshot.is_terminal
def test_stop_on_running_task_goes_through_stopping(machine: StateMachine) -> None:
_run_to_running(machine)
record = machine.request_stop(actor="ops", reason="发布窗口")
assert record.to_status is TaskStatus.STOPPING
assert machine.snapshot.requested_action is RequestedAction.STOP
assert not machine.snapshot.is_terminal
final = machine.mark_stopped()
assert final.to_status is TaskStatus.STOPPED
assert machine.snapshot.requested_action is None
assert machine.snapshot.validate() == ()
def test_forced_stop_uses_distinct_trigger(machine: StateMachine) -> None:
_run_to_running(machine, executor="thread")
machine.request_stop(actor="ops")
record = machine.mark_stopped(forced=True, reason="优雅窗口超时")
assert record.trigger == Trigger.STOP_FORCED
assert not record.is_override
def test_task_may_still_succeed_after_stop_requested(machine: StateMachine) -> None:
"""停止请求发出后任务恰好跑完，是合法结果，不能判为失败。"""
_run_to_running(machine)
machine.request_stop(actor="ops")
record = machine.mark_succeeded()
assert record.to_status is TaskStatus.SUCCEEDED
def test_stop_capability_is_false_on_terminal(machine: StateMachine) -> None:
_run_to_running(machine)
machine.mark_succeeded()
assert not machine.can_stop
with pytest.raises(InvalidTransition):
machine.request_stop(actor="ops")
# ---------------------------------------------------------------------------
# 暂停语义
# ---------------------------------------------------------------------------
def test_pause_on_unstarted_task_lands_paused_directly(machine: StateMachine) -> None:
"""排队中的任务本就静止，pause() 应立即落 PAUSED，不该显示"暂停中"。"""
machine.mark_queued()
record = machine.request_pause(actor="ops")
assert record.to_status is TaskStatus.PAUSED
assert machine.status is TaskStatus.PAUSED
assert machine.snapshot.requested_action is None
assert machine.can_resume
def test_pause_while_pausing_is_not_short_circuited(machine: StateMachine) -> None:
"""已在 PAUSING 的任务不能被"加急"成 PAUSED：它还没到达检查点。"""
_run_to_running(machine)
machine.request_pause()
assert machine.status is TaskStatus.PAUSING
with pytest.raises(InvalidTransition):
machine.request_pause()
assert machine.status is TaskStatus.PAUSING
def test_starting_always_records_started_at(machine: StateMachine) -> None:
"""任何到达 RUNNING 的路径都必须带有 started_at（含暂停撤销路径）。"""
machine.mark_queued()
machine.mark_starting(executor="inline")
assert machine.snapshot.started_at is not None
machine.request_pause()
machine.abort_pause(reason="未到达检查点")
assert machine.status is TaskStatus.RUNNING
assert machine.snapshot.validate() == ()
def test_pause_and_resume_round_trip(machine: StateMachine) -> None:
_run_to_running(machine)
machine.request_pause(actor="ops")
assert machine.snapshot.requested_action is RequestedAction.PAUSE
machine.mark_paused()
assert machine.status is TaskStatus.PAUSED
machine.resume()
assert machine.status is TaskStatus.RUNNING
assert machine.snapshot.requested_action is None
def test_inplace_resume_returns_to_running_not_queued(machine: StateMachine) -> None:
"""就地暂停的任务始终占着执行槽，调用栈还在内存里，恢复就是回到 RUNNING。
我们没有办法把任意的 Python 调用栈换出内存再换回来，所以"暂停后重新排队"
对就地暂停并不成立——那会导致任务状态与实际执行情况脱节。
"""
_run_to_running(machine)
machine.request_pause()
machine.mark_paused()
machine.resume()
assert machine.status is TaskStatus.RUNNING
assert machine.snapshot.queued_at is None or machine.snapshot.started_at is not None
def test_resume_of_a_never_started_pause_returns_to_queued(machine: StateMachine) -> None:
"""排队期间被暂停的任务从未启动过，恢复必须重新排队。"""
machine.mark_queued()
assert machine.request_pause().to_status is TaskStatus.PAUSED
assert machine.snapshot.attempt == 0
machine.resume()
assert machine.status is TaskStatus.QUEUED
def test_resume_of_a_retry_wait_pause_returns_to_queued(clock: FakeClock) -> None:
machine = StateMachine.create("t1", "x", max_attempts=3, clock=clock)
_run_to_running(machine)
machine.mark_retry_wait(Failure(kind=FailureKind.EXCEPTION, message="boom"))
assert machine.request_pause().to_status is TaskStatus.PAUSED
machine.resume()
assert machine.status is TaskStatus.QUEUED
def test_aborted_pause_reverts_to_running(machine: StateMachine) -> None:
"""暂停没能生效时必须如实回滚，而不是假装已暂停、让用户以为资源已释放。"""
_run_to_running(machine, executor="thread")
machine.request_pause()
record = machine.abort_pause(reason="任务未调用 checkpoint")
assert record.to_status is TaskStatus.RUNNING
assert record.trigger == Trigger.PAUSE_ABORTED
assert machine.snapshot.requested_action is None
# ---------------------------------------------------------------------------
# 人工干预
# ---------------------------------------------------------------------------
def test_force_transition_requires_audit_fields(machine: StateMachine) -> None:
with pytest.raises(ValueError, match="actor"):
machine.force_transition(TaskStatus.SUCCEEDED, actor=" ", reason="ok")
with pytest.raises(ValueError, match="reason"):
machine.force_transition(TaskStatus.SUCCEEDED, actor="ops", reason="")
def test_force_transition_bypasses_table_and_marks_override(machine: StateMachine) -> None:
"""PENDING -> SUCCEEDED 在转移表中非法，人工改写必须能通过，但要留痕。"""
record = machine.force_transition(
TaskStatus.SUCCEEDED, actor="ops@example", reason="已人工核验产物"
)
assert record.to_status is TaskStatus.SUCCEEDED
assert record.is_override
assert record.actor == "ops@example"
assert machine.snapshot.overridden
assert machine.snapshot.validate() == ()
def test_reset_returns_to_pending_and_clears_attempt_state(clock: FakeClock) -> None:
machine = StateMachine.create("t1", "x", max_attempts=2, clock=clock)
_run_to_running(machine)
machine.mark_failed(Failure(kind=FailureKind.EXCEPTION, message="boom"))
machine.reset(actor="ops", reason="重新跑一遍")
snapshot = machine.snapshot
assert snapshot.status is TaskStatus.PENDING
assert snapshot.attempt == 0
assert snapshot.failure is None
assert snapshot.ended_at is None
assert snapshot.started_at is None
assert not snapshot.overridden
# ---------------------------------------------------------------------------
# 能力预检（供控制面提前置灰按钮）
# ---------------------------------------------------------------------------
def test_capability_prechecks_match_actions(machine: StateMachine) -> None:
assert machine.can_stop
assert not machine.can_pause
assert not machine.can_resume
_run_to_running(machine)
assert machine.can_stop
assert machine.can_pause
assert not machine.can_resume
machine.request_pause()
machine.mark_paused()
assert machine.can_resume
assert machine.can_stop
def test_lost_is_terminal_but_revivable(machine: StateMachine) -> None:
_run_to_running(machine, executor="thread")
machine.mark_lost(reason="心跳超时")
assert machine.snapshot.is_terminal
assert machine.can_retry
./tests\unit\test_task_context.py
"""``@task`` 装饰器与 ``TaskContext`` 的能力测试。
三块：
1. **``@task``**——元数据从装饰器来；不产生全局状态（导入无副作用）；
2. **进度**——进快照、进存储、**不碰版本号**；
3. **上游数据**——管道节点之间真正能传值，且失败的上游不传 ``None``。
"""
from __future__ import annotations
import logging
import time
from typing import Any
import pytest
from chongming_task.controller import Controller, TaskRegistry
from chongming_task.decorators import get_task_spec, task
from chongming_task.errors import ConfigurationError
from chongming_task.pipeline import Pipeline
from chongming_task.policies import RetryPolicy, Timeout
from chongming_task.states import TaskStatus
from chongming_task.stores.memory import InMemoryTaskStore
_WINDOW = 15.0
# ---------------------------------------------------------------------------
# 1. @task
# ---------------------------------------------------------------------------
@task
def bare(ctx: Any) -> str:
"""裸装饰器。"""
del ctx
return "bare"
@task(tags={"io", "nightly"}, retry=RetryPolicy(max_attempts=4), timeout=Timeout(per_attempt=9.0))
def decorated(ctx: Any) -> str:
"""带元数据的任务。
第二行不该出现在描述里。
"""
del ctx
return "decorated"
@task(name="renamed")
def original_name(ctx: Any) -> str:
del ctx
return "renamed"
def undecorated(ctx: Any) -> str:
del ctx
return "nope"
def test_bare_decorator_works() -> None:
"""``@task`` 不带括号也太自然，不支持会让人反复踩空。"""
spec = get_task_spec(bare)
assert spec is not None
assert spec.name == "bare"
assert spec.description == "裸装饰器。"
def test_decorator_captures_metadata() -> None:
spec = get_task_spec(decorated)
assert spec is not None
assert spec.tags == frozenset({"io", "nightly"})
assert spec.retry is not None
assert spec.retry.max_attempts == 4
assert spec.timeout is not None
assert spec.timeout.per_attempt == 9.0
def test_description_defaults_to_the_first_docstring_line() -> None:
spec = get_task_spec(decorated)
assert spec is not None
assert spec.description == "带元数据的任务。", "只取第一行"
assert "第二行" not in spec.description
def test_name_override() -> None:
spec = get_task_spec(original_name)
assert spec is not None
assert spec.name == "renamed"
def test_undecorated_function_has_no_spec() -> None:
assert get_task_spec(undecorated) is None
def test_decorating_does_not_create_global_state() -> None:
"""规格只挂在函数对象上。
全局注册表意味着 **import 即注册**，而 import 可能发生在任何地方（包括测试收集
阶段），那种隐式行为极难排查。
"""
registry = TaskRegistry()
assert registry.names() == ()
assert get_task_spec(bare) is not None, "规格依然取得到，但没有登记进任何注册表"
def test_registry_registers_decorated_functions() -> None:
registry = TaskRegistry()
registration = registry.register_decorated(decorated)
assert registration.name == "decorated"
assert registration.tags == frozenset({"io", "nightly"})
assert registration.retry.max_attempts == 4
def test_registry_rejects_an_undecorated_function() -> None:
registry = TaskRegistry()
with pytest.raises(ConfigurationError, match="没有用 @task 装饰过"):
registry.register_decorated(undecorated)
def test_register_module_picks_up_every_decorated_function() -> None:
import sys
module = sys.modules[__name__]
registry = TaskRegistry()
registered = registry.register_module(module)
names = {item.name for item in registered}
assert {"bare", "decorated", "renamed"} <= names
assert "undecorated" not in names
def test_register_module_rejects_a_module_with_nothing_decorated() -> None:
import json
registry = TaskRegistry()
with pytest.raises(ConfigurationError, match="没有用 @task 装饰过的函数"):
registry.register_module(json)
def test_controller_registers_a_module_end_to_end() -> None:
import sys
controller = Controller(max_parallel=2)
controller.register_module(sys.modules[__name__])
run = controller.submit("bare")
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
assert outcome.ok
assert outcome.value == "bare"
controller.close(timeout=_WINDOW)
def test_decorated_retry_actually_applies() -> None:
"""装饰器上写的重试次数必须真的生效，而不只是躺在元数据里。"""
store = InMemoryTaskStore()
controller = Controller(store=store, max_parallel=2)
registry = TaskRegistry()
registry.register(
"flaky",
_flaky,
retry=RetryPolicy(max_attempts=3, backoff=_no_backoff()),
)
controller.register("flaky", _flaky, retry=RetryPolicy(max_attempts=3, backoff=_no_backoff()))
run = controller.submit("flaky")
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
assert outcome.ok, "第三次应成功"
snapshot = store.get(run.snapshot.task_id)
assert snapshot is not None
assert snapshot.attempt == 3
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 2. TaskContext：进度与日志
# ---------------------------------------------------------------------------
def test_progress_reaches_the_snapshot_without_bumping_the_version() -> None:
"""进度若参与乐观并发，会与状态写入持续冲突。"""
controller = Controller(max_parallel=2)
controller.register("reporter", _reporter)
run = controller.submit("reporter")
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
assert outcome.ok
snapshot = outcome.snapshot
assert snapshot.progress is not None
assert snapshot.progress.done == 3
assert snapshot.progress.total == 3
assert snapshot.progress.unit == "rows"
# 状态迁移只有 4 次（queued/starting/running/succeeded），进度不该挤进去。
assert snapshot.status_version == 4
def test_progress_is_persisted_and_visible_in_queries() -> None:
"""进度要经 touch 旁路落到存储，否则看板的 PROGRESS 列永远是空的。"""
store = InMemoryTaskStore()
controller = Controller(store=store, max_parallel=2, command_poll_interval=0.05)
controller.register("slow_reporter", _slow_reporter)
run = controller.submit("slow_reporter")
assert _wait_until(
lambda: (store.get(run.snapshot.task_id) or _stub()).progress is not None,
timeout=_WINDOW,
)
persisted = store.get(run.snapshot.task_id)
assert persisted is not None
assert persisted.progress is not None
assert persisted.progress.done >= 1
controller.stop(run.snapshot.task_id)
controller.wait_all(timeout=_WINDOW)
controller.close(timeout=_WINDOW)
def test_progress_on_an_unsupported_context_is_silently_dropped() -> None:
"""没有回写通道时静默丢弃，而不是抛异常。
直接构造上下文（不经过执行器）是测试与工具里的常见用法，不该因此报错。
"""
from chongming_task.runtime.cancellation import CancellationToken
from chongming_task.runtime.context import RuntimeContext
from chongming_task.runtime.heartbeat import Heartbeat
context = RuntimeContext(
task_id="t",
name="t",
attempt=1,
token=CancellationToken(),
heartbeat=Heartbeat(),
)
context.progress(1, total=2)
context.log("hello")
assert list(context.logs) == ["hello"]
def test_log_records_into_the_buffer_and_the_stdlib_logger(
caplog: pytest.LogCaptureFixture,
) -> None:
# caplog 默认只捕获 WARNING 及以上，而 INFO 才是我们最常用的级别。
caplog.set_level(logging.INFO, logger="chongming_task.task")
controller = Controller(max_parallel=2)
controller.register("talker", _talker)
run = controller.submit("talker")
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
assert outcome.ok
assert outcome.value == "done"
messages = [record.message for record in caplog.records]
assert "hello" in messages, "INFO 日志应当走到标准库 logging"
assert "careful" in messages
record = next(r for r in caplog.records if r.message == "hello")
assert record.task_id == run.snapshot.task_id
assert record.task_name == "talker"
assert record.attempt == 1
def test_task_context_is_the_runtime_context() -> None:
"""公开名字与内部名字是**同一个类**。
任务作者只面对一个 ``ctx``；两个类会立刻带来"我该用哪个"的困惑。
"""
from chongming_task import RuntimeContext, TaskContext
assert TaskContext is RuntimeContext
def test_checkpoint_still_works_alongside_progress() -> None:
controller = Controller(max_parallel=2)
controller.register("stopme", _stopme)
run = controller.submit("stopme")
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
controller.stop(run.snapshot.task_id, reason="收工")
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
assert outcome.status is TaskStatus.STOPPED
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 3. 上游数据（管道传值）
# ---------------------------------------------------------------------------
def test_downstream_receives_upstream_values() -> None:
controller = Controller(max_parallel=4)
controller.register("produce", _produce)
controller.register("consume", _consume)
pipeline = Pipeline("data")
source = pipeline.node("produce", node_id="source")
pipeline.node("consume", node_id="sink", depends_on=[source])
run = controller.submit_pipeline(pipeline, pipeline_run_id="d1")
outcome = run.wait(timeout=_WINDOW)
assert outcome.status.value == "succeeded"
assert outcome.nodes["sink"] is TaskStatus.SUCCEEDED
snapshot = controller.get("d1:sink")
assert snapshot is not None
assert snapshot.outcome == {"value": 42}, "下游应当拿到上游的返回值 42"
controller.close(timeout=_WINDOW)
def test_a_failed_upstream_contributes_no_value() -> None:
"""失败的上游不塞 ``None``：那会让下游以为"上游返回了 None"。"""
controller = Controller(max_parallel=4)
controller.register("boom", _boom)
controller.register("inspect", _inspect)
pipeline = Pipeline("missing")
bad = pipeline.node("boom", node_id="bad")
pipeline.node("inspect", node_id="probe", depends_on=[bad], join=_all_settled())
run = controller.submit_pipeline(pipeline, pipeline_run_id="d2")
run.wait(timeout=_WINDOW)
snapshot = controller.get("d2:probe")
assert snapshot is not None
assert snapshot.outcome == {"value": []}, "上游没有值，ctx.upstream 必须是空的"
controller.close(timeout=_WINDOW)
def test_input_can_be_read_with_a_default() -> None:
controller = Controller()
controller.register("use_default", _use_default)
run = controller.submit("use_default")
outcome = controller.wait(run.snapshot.task_id, timeout=_WINDOW)
assert outcome.value == "fallback"
controller.close(timeout=_WINDOW)
def test_inputs_cannot_be_injected_after_the_task_started() -> None:
controller = Controller(max_parallel=2)
controller.register("stopme", _stopme)
run = controller.submit("stopme")
assert _wait_until(lambda: run.status is TaskStatus.RUNNING)
with pytest.raises(RuntimeError, match="已经开始执行"):
run.set_inputs({"x": 1})
controller.stop(run.snapshot.task_id)
controller.wait_all(timeout=_WINDOW)
controller.close(timeout=_WINDOW)
# ---------------------------------------------------------------------------
# 任务体
# ---------------------------------------------------------------------------
def _stub() -> Any:
from chongming_task.models import TaskSnapshot
return TaskSnapshot(task_id="stub", name="stub")
def _no_backoff() -> Any:
from chongming_task.policies import FixedBackoff
return FixedBackoff(0.0)
def _flaky(ctx: Any) -> str:
state = getattr(_flaky, "_count", 0) + 1
_flaky._count = state # type: ignore[attr-defined]
if state < 3:
raise RuntimeError(f"第 {state} 次失败")
return "ok"
def _reporter(ctx: Any) -> str:
for index in range(1, 4):
ctx.checkpoint()
ctx.progress(index, total=3, unit="rows")
return "reported"
def _slow_reporter(ctx: Any) -> str:
for index in range(1, 200):
ctx.checkpoint()
ctx.progress(index, total=200, unit="rows")
time.sleep(0.01)
return "done"
def _talker(ctx: Any) -> str:
ctx.log("hello")
ctx.warning("careful")
return "done"
def _stopme(ctx: Any) -> str:
for _ in range(1_000_000):
ctx.checkpoint()
time.sleep(0.005)
return "never"
def _produce(ctx: Any) -> int:
del ctx
return 42
def _consume(ctx: Any) -> int:
value = ctx.input("source")
assert value == 42, f"预期 42，收到 {value!r}"
return int(value)
def _inspect(ctx: Any) -> list[str]:
return sorted(ctx.upstream)
def _use_default(ctx: Any) -> str:
return str(ctx.input("nobody", "fallback"))
def _boom(ctx: Any) -> str:
del ctx
raise RuntimeError("故意失败")
def _all_settled() -> Any:
from chongming_task.policies import JoinMode
return JoinMode.ALL_SETTLED
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
./tests\unit\test_thread_executor.py
"""线程执行器测试。
这些用例全部使用**真实线程与真实时间**：执行器的全部价值就在于线程能不能被
停下、能不能被暂停、放弃之后进程还能不能退出。用假的线程模型测不出这些。
为避免偶发失败，协作路径给足超时（5 秒），而"不配合"的用例刻意把任务做成
有限时长（约 5 秒后自行结束），这样即使被放弃，残留的守护线程也不会长期占用资源。
"""
from __future__ import annotations
import asyncio
import threading
import time
from typing import Any
import pytest
from chongming_task.models import FailureKind
from chongming_task.runtime.executors.base import ExecutionRequest
from chongming_task.runtime.executors.thread import ThreadExecutor
from chongming_task.testing import FakeClock
_COOPERATIVE_WINDOW = 5.0
_ABANDON_WINDOW = 0.2
def _request(task_id: str, target: Any, *args: Any) -> ExecutionRequest:
return ExecutionRequest(task_id=task_id, name=task_id, target=target, args=args)
# ---------------------------------------------------------------------------
# 上下文注入契约
# ---------------------------------------------------------------------------
def test_context_is_injected_with_the_handle_token_and_metadata() -> None:
"""令牌必须是句柄那一个，否则控制面会停错对象。
这也是"令牌由执行器注入、而非被任务闭包捕获"的直接体现——令牌每次尝试才诞生，
闭包捕获到的只可能是上一次尝试留下的那一个。
"""
seen: dict[str, Any] = {}
def work(ctx: Any) -> str:
seen["token"] = ctx.token
seen["task_id"] = ctx.task_id
seen["name"] = ctx.name
seen["attempt"] = ctx.attempt
return "ok"
handle = ThreadExecutor().launch(
ExecutionRequest(task_id="t-ctx", name="抽取", target=work, attempt=2)
)
assert handle.wait(timeout=_COOPERATIVE_WINDOW) is not None
assert seen["token"] is handle.token
assert seen["task_id"] == "t-ctx"
assert seen["name"] == "抽取"
assert seen["attempt"] == 2
def test_arguments_are_forwarded_after_the_context() -> None:
def work(ctx: Any, left: int, right: int, *, scale: int = 1) -> int:
del ctx
return (left + right) * scale
handle = ThreadExecutor().launch(
ExecutionRequest(
task_id="t-args",
name="t-args",
target=work,
args=(2, 3),
kwargs={"scale": 10},
)
)
outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
assert outcome is not None
assert outcome.value == 50
def test_checkpoint_refreshes_the_heartbeat() -> None:
"""checkpoint 就是"我还在推进"的信号，心跳绑在它上面是有意的。"""
clock = FakeClock()
observed: list[float] = []
def work(ctx: Any) -> str:
for _ in range(3):
clock.advance(30)
observed.append(ctx.heartbeat.idle_seconds())
ctx.checkpoint()
return "ok"
handle = ThreadExecutor(clock=clock).launch(_request("t-beat", work))
outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
assert outcome is not None
assert outcome.ok
assert observed == [30.0, 30.0, 30.0], "每次 checkpoint 都应把静默时长清零"
assert handle.heartbeat.idle_seconds() == 0.0
# ---------------------------------------------------------------------------
# async 任务
# ---------------------------------------------------------------------------
def test_async_target_is_actually_awaited() -> None:
"""``async def`` 任务必须真的被执行完，而不是把协程对象当返回值。
后者曾经是静默的：没有存储时它看起来"成功"了，返回值却是个 ``coroutine``；
有存储时才会以"返回值不可序列化"失败。
"""
async def work(ctx: Any) -> str:
await asyncio.sleep(0)
return "done"
handle = ThreadExecutor().launch(_request("t-async", work))
outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
assert outcome is not None
assert outcome.ok
assert outcome.value == "done"
def test_async_target_cooperates_with_stop() -> None:
"""协作式停止对 async 任务同样有效：``await ctx.acheckpoint()``。"""
async def work(ctx: Any) -> str:
for _ in range(10_000):
await ctx.acheckpoint()
await asyncio.sleep(0.005)
return "never"
handle = ThreadExecutor().launch(_request("t-async-stop", work))
outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
assert not outcome.ok, "停止请求之后不该报成功"
assert not outcome.abandoned, "它会到达检查点，应当优雅收尾而不是被放弃"
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
def test_async_target_failure_is_reported_across_the_loop() -> None:
"""协程里抛的异常照常转成结构化失败，不能被事件循环吞掉。"""
async def work(ctx: Any) -> str:
await asyncio.sleep(0)
raise ValueError("协程里炸了")
handle = ThreadExecutor().launch(_request("t-async-boom", work))
outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.EXCEPTION
assert "协程里炸了" in outcome.failure.message
# ---------------------------------------------------------------------------
# 正常路径
# ---------------------------------------------------------------------------
def test_launch_returns_immediately_without_waiting_for_the_task() -> None:
"""调度器契约：launch 必须立刻返回，否则调度循环会被单个任务卡住。"""
release = threading.Event()
def work(ctx: Any) -> str:
del ctx
release.wait(5)
return "done"
handle = ThreadExecutor().launch(_request("t-async-launch", work))
try:
assert handle.running is True
assert handle.wait(timeout=0.05) is None, "任务尚未结束，wait 应返回 None"
finally:
release.set()
outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
assert outcome is not None
assert outcome.ok
assert outcome.value == "done"
def test_successful_task_returns_its_value() -> None:
handle = ThreadExecutor().launch(_request("t-ok", lambda ctx: 6 * 7))
outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
assert outcome is not None
assert outcome.ok
assert outcome.value == 42
assert handle.running is False
def test_failing_task_reports_exception_with_traceback() -> None:
def boom(ctx: Any) -> None:
del ctx
raise ValueError("坏掉了")
handle = ThreadExecutor().launch(_request("t-boom", boom))
outcome = handle.wait(timeout=_COOPERATIVE_WINDOW)
assert outcome is not None
assert not outcome.ok
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.EXCEPTION
assert "ValueError" in outcome.failure.message
assert outcome.failure.traceback is not None
assert "boom" in outcome.failure.traceback
assert outcome.failure.retryable is True
# ---------------------------------------------------------------------------
# 协作式停止
# ---------------------------------------------------------------------------
def test_stop_completes_cooperatively_when_task_checks_the_token() -> None:
started = threading.Event()
steps: list[int] = []
def work(ctx: Any) -> str:
started.set()
for index in range(100_000):
ctx.checkpoint()
steps.append(index)
time.sleep(0.001)
return "ran to completion"
handle = ThreadExecutor().launch(_request("t-coop", work))
assert started.wait(_COOPERATIVE_WINDOW)
outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW, reason="发布窗口")
assert outcome.abandoned is False
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
assert outcome.failure.retryable is False, "取消不是失败，绝不能被重试"
assert handle.running is False
assert len(steps) < 100_000, "任务应当在收到停止请求后提前退出"
def test_cancellation_is_not_swallowed_by_a_broad_except_in_the_task_body() -> None:
"""这正是 ``TaskCancelled`` 继承 ``BaseException`` 的理由。
若它继承 ``Exception``，下面这种极常见的写法会把停止信号吃掉，任务永远停不掉。
"""
started = threading.Event()
swallowed: list[BaseException] = []
def work(ctx: Any) -> str:
started.set()
try:
while True:
ctx.checkpoint()
time.sleep(0.001)
except Exception as exc: # 用户代码里最常见的写法
swallowed.append(exc)
return "signal was swallowed"
handle = ThreadExecutor().launch(_request("t-swallow", work))
assert started.wait(_COOPERATIVE_WINDOW)
outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
assert swallowed == [], "取消信号被 except Exception 吞掉了"
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
def test_stop_is_idempotent() -> None:
def work(ctx: Any) -> str:
del ctx
return "quick"
handle = ThreadExecutor().launch(_request("t-idem", work))
first = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
second = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
assert first.ok
assert second.ok
assert handle.running is False
# ---------------------------------------------------------------------------
# 协作式暂停（D6 定案的语义，端到端验证）
# ---------------------------------------------------------------------------
def test_pause_suspends_a_running_task_and_resume_continues_it() -> None:
started = threading.Event()
progress: list[int] = []
def work(ctx: Any) -> str:
started.set()
for index in range(100_000):
ctx.checkpoint()
progress.append(index)
time.sleep(0.001)
return "ran to completion"
handle = ThreadExecutor().launch(_request("t-pause", work))
assert started.wait(_COOPERATIVE_WINDOW)
assert handle.pause(timeout=_COOPERATIVE_WINDOW) is True
frozen = len(progress)
assert frozen > 0, "暂停前任务应当已经推进过"
time.sleep(0.05)
assert len(progress) == frozen, "暂停生效后任务不应继续推进"
handle.resume()
time.sleep(0.05)
assert len(progress) > frozen, "恢复后任务应继续推进"
outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW)
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
def test_pause_reports_false_when_task_never_reaches_a_checkpoint() -> None:
"""这正是"暂停没生效必须回滚为 RUNNING"的依据：如实返回 False。"""
started = threading.Event()
def work(ctx: Any) -> str:
del ctx
started.set()
for _ in range(500):
time.sleep(0.01)
return "ignored"
handle = ThreadExecutor().launch(_request("t-nopause", work))
assert started.wait(_COOPERATIVE_WINDOW)
assert handle.pause(timeout=0.2) is False
handle.stop(graceful_timeout=_ABANDON_WINDOW)
def test_stop_wakes_a_paused_task() -> None:
"""暂停中的任务必须能被停止，否则暂停就成了不可撤销的挂起。"""
started = threading.Event()
def work(ctx: Any) -> str:
started.set()
for _ in range(100_000):
ctx.checkpoint()
time.sleep(0.001)
return "ran to completion"
handle = ThreadExecutor().launch(_request("t-pause-stop", work))
assert started.wait(_COOPERATIVE_WINDOW)
assert handle.pause(timeout=_COOPERATIVE_WINDOW) is True
outcome = handle.stop(graceful_timeout=_COOPERATIVE_WINDOW, reason="运维介入")
assert outcome.abandoned is False
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.CANCELLED
assert "运维介入" in outcome.failure.message
# ---------------------------------------------------------------------------
# 无法停止：放弃
# ---------------------------------------------------------------------------
def _uncooperative(ctx: Any) -> str:
"""刻意不检查令牌；约 5 秒后自行结束，避免残留线程长期占用资源。"""
del ctx
for _ in range(500):
time.sleep(0.01)
return "ignored the stop request"
def test_uncooperative_task_is_abandoned_and_reported_as_lost() -> None:
"""线程无法强杀：必须如实上报，不能假装已停止。"""
handle = ThreadExecutor().launch(_request("t-stubborn", _uncooperative))
outcome = handle.stop(graceful_timeout=_ABANDON_WINDOW)
assert outcome.abandoned is True
assert outcome.ok is False
assert outcome.failure is not None
assert outcome.failure.kind is FailureKind.LOST
assert "无法强制终止" in outcome.failure.message
def test_abandoned_task_still_reports_running() -> None:
"""它确实还在跑。谎报为 False 会让上层以为资源已经释放。"""
handle = ThreadExecutor().launch(_request("t-still-running", _uncooperative))
handle.stop(graceful_timeout=_ABANDON_WINDOW)
assert handle.running is True
def test_abandoned_outcome_is_sticky() -> None:
"""已记为 LOST 的结论不应被后台线程后来的返回悄悄改写，否则审计自相矛盾。"""
handle = ThreadExecutor().launch(_request("t-sticky", _uncooperative))
first = handle.stop(graceful_timeout=_ABANDON_WINDOW)
assert handle.wait(timeout=_ABANDON_WINDOW) is first
assert handle.stop(graceful_timeout=_ABANDON_WINDOW) is first
def test_worker_threads_are_daemon_so_an_abandoned_task_cannot_block_shutdown() -> None:
"""``ThreadPoolExecutor`` 的工作线程会在解释器退出时被 join，一个不响应取消
的任务会让整个进程无法退出。这里用守护线程正是为了避开这一点。"""
handle = ThreadExecutor().launch(_request("t-daemon", _uncooperative))
try:
assert handle.thread.daemon is True
assert handle.thread.name == "task-t-daemon"
finally:
handle.stop(graceful_timeout=_ABANDON_WINDOW)
@pytest.mark.parametrize("graceful_timeout", [0.0, 0.05])
def test_zero_or_tiny_graceful_window_still_yields_a_conclusion(
graceful_timeout: float,
) -> None:
"""优雅窗口再短也必须给出结论，不能抛异常或永远阻塞。"""
handle = ThreadExecutor().launch(_request(f"t-window-{graceful_timeout}", _uncooperative))
outcome = handle.stop(graceful_timeout=graceful_timeout)
assert outcome.abandoned is True
assert outcome.failure is not None
./tests\unit\test_timeouts.py
"""心跳与超时判定的测试。
全部基于 ``FakeClock``：不 sleep、不依赖真实时间，因此既快又不会有偶发失败。
"""
from __future__ import annotations
from datetime import UTC, datetime, timedelta
import pytest
from chongming_task.models import FailureKind, TimeoutScope
from chongming_task.policies import Timeout
from chongming_task.runtime.heartbeat import Heartbeat
from chongming_task.runtime.timeouts import (
Deadline,
is_expired,
next_deadline,
remaining_seconds,
)
from chongming_task.testing import FakeClock
_T0 = datetime(2026, 1, 1, tzinfo=UTC)
# ---------------------------------------------------------------------------
# 心跳
# ---------------------------------------------------------------------------
def test_heartbeat_starts_at_construction() -> None:
clock = FakeClock(_T0)
beat = Heartbeat(clock)
assert beat.last_beat == _T0
assert beat.idle_seconds() == 0.0
def test_heartbeat_advances_only_on_beat() -> None:
"""心跳由 checkpoint 驱动，不额外起线程；没 beat 就不该自己变新。"""
clock = FakeClock(_T0)
beat = Heartbeat(clock)
clock.advance(30)
assert beat.idle_seconds() == 30.0
beat.beat()
assert beat.last_beat == clock.now
assert beat.idle_seconds() == 0.0
def test_heartbeat_staleness_boundary_is_inclusive() -> None:
clock = FakeClock(_T0)
beat = Heartbeat(clock)
clock.advance(59.999)
assert not beat.is_stale(60)
clock.advance(0.001)
assert beat.is_stale(60)
def test_heartbeat_rejects_non_positive_timeout() -> None:
beat = Heartbeat(FakeClock(_T0))
with pytest.raises(ValueError, match="必须为正数"):
beat.is_stale(0)
# ---------------------------------------------------------------------------
# 超时判定
# ---------------------------------------------------------------------------
def test_no_timeout_configured_yields_none() -> None:
assert next_deadline(Timeout()) is None
def test_start_scope_applies_only_before_started() -> None:
timeout = Timeout(start=30)
pending = next_deadline(timeout, submitted_at=_T0)
assert pending is not None
assert pending.scope is TimeoutScope.START
assert pending.at == _T0 + timedelta(seconds=30)
# 一旦开始执行，start 超时即失效。
assert next_deadline(timeout, submitted_at=_T0, started_at=_T0) is None
def test_start_scope_needs_submitted_at() -> None:
assert next_deadline(Timeout(start=30)) is None
def test_attempt_scope_measures_from_started_at() -> None:
deadline = next_deadline(Timeout(per_attempt=600), started_at=_T0)
assert deadline is not None
assert deadline.scope is TimeoutScope.ATTEMPT
assert deadline.at == _T0 + timedelta(seconds=600)
def test_idle_scope_follows_the_latest_heartbeat() -> None:
"""心跳推进后，idle 截止点必须随之顺延——否则长任务会被误杀。"""
timeout = Timeout(idle=60)
first = next_deadline(timeout, started_at=_T0, last_beat_at=_T0)
assert first is not None
assert first.at == _T0 + timedelta(seconds=60)
later = _T0 + timedelta(seconds=45)
second = next_deadline(timeout, started_at=_T0, last_beat_at=later)
assert second is not None
assert second.at == later + timedelta(seconds=60)
def test_idle_scope_requires_the_task_to_have_started() -> None:
"""任务还没跑起来就谈不上"卡住"，否则排队时长会被误判成静默超时。"""
assert next_deadline(Timeout(idle=60), last_beat_at=_T0) is None
def test_before_start_only_the_start_scope_applies() -> None:
timeout = Timeout(start=10, per_attempt=600, idle=60)
deadline = next_deadline(timeout, submitted_at=_T0, last_beat_at=_T0)
assert deadline is not None
assert deadline.scope is TimeoutScope.START
assert deadline.at == _T0 + timedelta(seconds=10)
def test_after_start_attempt_and_idle_compete_and_earliest_wins() -> None:
"""开始执行后只剩两档竞争：宽松的一档不能掩盖已经到期的一档。"""
timeout = Timeout(per_attempt=600, idle=60)
fresh_beat = next_deadline(timeout, started_at=_T0, last_beat_at=_T0)
assert fresh_beat is not None
assert fresh_beat.scope is TimeoutScope.IDLE
# 心跳落后很久时，attempt 反而先到期。
late_beat = _T0 + timedelta(seconds=590)
stale_beat = next_deadline(timeout, started_at=_T0, last_beat_at=late_beat)
assert stale_beat is not None
assert stale_beat.scope is TimeoutScope.ATTEMPT
assert stale_beat.at == _T0 + timedelta(seconds=600)
def test_deadline_maps_to_failure_kind() -> None:
"""超时是失败原因，不是生命周期状态；映射必须稳定。"""
assert Deadline(TimeoutScope.START, _T0).kind is FailureKind.TIMEOUT_START
assert Deadline(TimeoutScope.ATTEMPT, _T0).kind is FailureKind.TIMEOUT_ATTEMPT
assert Deadline(TimeoutScope.IDLE, _T0).kind is FailureKind.TIMEOUT_IDLE
def test_is_expired_boundary_is_inclusive() -> None:
deadline = Deadline(TimeoutScope.ATTEMPT, _T0)
assert not is_expired(deadline, now=_T0 - timedelta(microseconds=1))
assert is_expired(deadline, now=_T0)
def test_remaining_seconds_goes_negative_after_expiry() -> None:
"""调度器要把负数当成"立刻到期"，而不是交给 sleep——sleep(-1) 会抛异常。"""
deadline = Deadline(TimeoutScope.ATTEMPT, _T0)
assert remaining_seconds(deadline, now=_T0 - timedelta(seconds=5)) == 5.0
assert remaining_seconds(deadline, now=_T0 + timedelta(seconds=5)) == -5.0