From 447413530190f1be48b27bc3cb61452b10698329 Mon Sep 17 00:00:00 2001
From: lda <lda@ldlda.com>
Date: Thu, 10 Sep 2026 03:57:40 +0700
Subject: [PATCH] test: cover admin dispatch registration race

---
 tests/wf_server/test_scheduler_integration.py | 131 ++++++++++++++++++
 1 file changed, 131 insertions(+)

diff --git a/tests/wf_server/test_scheduler_integration.py b/tests/wf_server/test_scheduler_integration.py
index 73de92f7..bf4c8466 100644
--- a/tests/wf_server/test_scheduler_integration.py
+++ b/tests/wf_server/test_scheduler_integration.py
@@ -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)