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。
380 lines
14 KiB
Python
380 lines
14 KiB
Python
"""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)
|