chakcy
f4de43c743
单任务受控执行（15 态状态机 + 审计轨迹）、DAG 管道、Controller 门面。
执行器三个：线程（默认，含 async def 任务）、内联（同步 / 测试）、进程
（spawn 子进程：硬性停止、绕开 GIL、崩溃隔离）。
调度：并发额度 + 每桶限流 + 具名资源信号量（按标签）+ 优先级排队。
持久层：TaskStore 契约 + 内存 / SQLite / Postgres 实现，契约测试随包发布。
跨进程与多节点：写穿 / 读穿、指令通道与租约、owned_buckets 划界、reconcile。
CLI、12 个可直接运行的示例、docs/ 下 5 份文档。
门禁：ruff / ruff format / mypy(strict) 干净；
pytest 563 passed（44 skipped 是需要真实 Postgres 的那批）；
带一台可连的 Postgres 时 607 passed。
679 lines
26 KiB
Python
679 lines
26 KiB
Python
"""任务管理门面（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] == "-", "未分桶的任务应当显示 -"
|