test: cover admin dispatch registration race

This commit is contained in:
lda
2026-09-10 04:02:27 +07:00 Verified
parent 8f5523bd2e
commit 4474135301
@@ -488,6 +488,137 @@ async def test_scheduled_completion(tmp_path: Path) -> None:
assert report.settled >= 1
@pytest.mark.parametrize("mutation", ["pause", "update", "delete", "shutdown"])
def test_admin_mutation_does_not_deadlock_dispatch_registration(
tmp_path: Path, mutation: str
) -> None:
"""A store mutation must not block registration on the event loop."""
child = f"""
import asyncio
import sys
import threading
import time
from pathlib import Path
sys.path.insert(0, {str(_REPO_ROOT)!r})
from datetime import UTC, datetime, timedelta
from tests.wf_server.test_scheduler_integration import (
_constant_plan,
_one_shot,
_scheduler,
_seed,
)
import wf_scheduling.lifecycle as lifecycle
from wf_scheduling.ownership import SchedulerOwnership
from wf_server import build_local_static_workflow_server
async def main() -> None:
root = Path(sys.argv[1])
server = build_local_static_workflow_server(root, schedules=True)
await _seed(
server, "constant", "constant.default", _constant_plan("constant"), ["ok"]
)
due = _one_shot("once", "constant.default", datetime.now(UTC))
await server.api.schedules.create_schedule(
schedule_id=due.id,
deployment_id=due.deployment_id,
trigger=due.trigger.model_dump(mode="json"),
)
# API creation starts the watermark at creation time; move it behind the
# already-created one-shot so the deterministic poll observes it as due.
server.api.schedules._schedule_store().save_consumed(
"once", due.trigger.at - timedelta(seconds=1)
)
service = _scheduler(server, capacity=1, auto_tick=False)
await service.start()
registration_entered = threading.Event()
release_registration = threading.Event()
original = asyncio.run_coroutine_threadsafe
def blocked_registration(coro, loop):
registration_entered.set()
if not release_registration.wait(5):
raise AssertionError("registration release was not signalled")
return original(coro, loop)
lifecycle.asyncio.run_coroutine_threadsafe = blocked_registration
poll_task = asyncio.create_task(
service.poll_once(due.trigger.at + timedelta(seconds=1))
)
if not await asyncio.to_thread(registration_entered.wait, 5):
raise AssertionError("poll did not reach the registration gate")
if sys.argv[2] in ("pause", "shutdown"):
admin_task = asyncio.create_task(
server.api.schedules.pause_schedule(schedule_id="once")
)
elif sys.argv[2] == "update":
admin_task = asyncio.create_task(
server.api.schedules.update_schedule(
schedule_id="once", expected_revision=1, overlap="parallel"
)
)
else:
admin_task = asyncio.create_task(
server.api.schedules.delete_schedule(schedule_id="once")
)
threading.Thread(
target=lambda: (time.sleep(0.2), release_registration.set()),
daemon=True,
).start()
try:
if sys.argv[2] == "shutdown":
shutdown_task = asyncio.create_task(service.stop())
result, paused, report = await asyncio.wait_for(
asyncio.gather(poll_task, admin_task, shutdown_task), 4
)
assert report is not None
assert service.running is False
probe = SchedulerOwnership(root, owner="shutdown-probe").acquire()
probe.release()
else:
result, paused = await asyncio.wait_for(
asyncio.gather(poll_task, admin_task), 4
)
assert result["once"].startswith("admit:run-")
run_id = result["once"].split(":", 1)[1]
admission = server.stores.run_store.get_admission(run_id)
assert admission.schedule_revision == 1
if sys.argv[2] in ("pause", "shutdown"):
assert paused["paused"] is True
elif sys.argv[2] == "update":
assert paused["revision"] == 2
else:
assert paused["deleted"] is True
print("admin-registration-complete", flush=True)
finally:
release_registration.set()
await service.stop()
asyncio.run(main())
"""
script = tmp_path / "admin_registration_deadlock.py"
script.write_text(child, encoding="utf-8")
process = subprocess.Popen(
[sys.executable, str(script), str(tmp_path / "store"), mutation],
cwd=_REPO_ROOT,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
)
try:
output, _ = process.communicate(timeout=7)
except subprocess.TimeoutExpired:
process.kill()
output, _ = process.communicate(timeout=5)
pytest.fail(f"admin/registration deadlocked; child output: {output}")
assert process.returncode == 0, output
assert "admin-registration-complete" in output
async def test_scheduled_failure(tmp_path: Path) -> None:
root = tmp_path / "store"
server = build_local_static_workflow_server(root)