"""节点与桶归属测试：**自治节点 + 统一平面**架构的划分单位。 这个架构里，每个节点有自己的任务定义与触发方式，共用同一套存储与控制层， 靠 ``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"