sched: gate scheduled resumes on the shared execution slot, with drain tracking (B2)

2026-09-09 17:23:28 +07:00
parent fd2f85ef96
commit ada0239a1c
8 changed files with 459 additions and 15 deletions
@@ -43,6 +43,7 @@ from wf_platform import (
 )
 from wf_scheduling.lifecycle import SchedulerServiceConfig, SchedulerStartupError
 from wf_scheduling.models import Schedule
+from wf_scheduling.resume_gate import ScheduledCapacityBusyError
 from wf_scheduling.store import FileScheduleStore, ScheduleExistsError
 from wf_server import WorkflowServer, build_local_static_workflow_server
 from wf_server.scheduling import build_scheduler_service
@@ -984,3 +985,278 @@ async def test_paused_deleted_run_completion(
    finally:
        _gate_open.set()
        await deleter.stop()
+
+
+def _interrupt_then_fail_plan(name: str) -> RawWorkflowPlan:
+    """Approval interrupt whose resumed continuation raises (resume failure)."""
+    return RawWorkflowPlan.model_validate(
+        {
+            "name": name,
+            "input_schema": {"type": "object", "properties": {}},
+            "state_schema": {"type": "object", "properties": {}},
+            "output_schema": {"type": "object", "properties": {}},
+            "outcomes": ["submitted"],
+            "start": "ask",
+            "nodes": [
+                {"id": "ask", "type": "interrupt", "kind": "approval"},
+                {
+                    "id": "fail",
+                    "type": "node",
+                    "node": "wf.std.runtime_error",
+                    "input": [
+                        {
+                            "value": "boom on resume",
+                            "target": {"root": "local", "parts": ["message"]},
+                        }
+                    ],
+                    "output": [],
+                },
+                {"id": "end_submitted", "type": "end", "outcome": "submitted"},
+            ],
+            "edges": [
+                {"from": "ask", "outcome": "submitted", "to": "fail"},
+                {"from": "fail", "outcome": "ok", "to": "end_submitted"},
+            ],
+            "output": [],
+        }
+    )
+
+
+def _interrupt_then_gate_plan(name: str) -> RawWorkflowPlan:
+    """Approval interrupt whose resumed continuation blocks on the test gate."""
+    return RawWorkflowPlan.model_validate(
+        {
+            "name": name,
+            "input_schema": {"type": "object", "properties": {}},
+            "state_schema": {
+                "type": "object",
+                "properties": {"result": {"type": "string"}},
+            },
+            "output_schema": {
+                "type": "object",
+                "properties": {"result": {"type": "string"}},
+            },
+            "outcomes": ["submitted"],
+            "start": "ask",
+            "nodes": [
+                {"id": "ask", "type": "interrupt", "kind": "approval"},
+                {
+                    "id": "gate",
+                    "type": "node",
+                    "node": "test.gate.wait",
+                    "input": [
+                        {
+                            "value": "hello",
+                            "target": {"root": "local", "parts": ["value"]},
+                        }
+                    ],
+                    "output": [
+                        {
+                            "source": {"root": "local", "parts": ["value"]},
+                            "target": {"root": "state", "parts": ["result"]},
+                        }
+                    ],
+                },
+                {"id": "end_submitted", "type": "end", "outcome": "submitted"},
+            ],
+            "edges": [
+                {"from": "ask", "outcome": "submitted", "to": "gate"},
+                {"from": "gate", "outcome": "ok", "to": "end_submitted"},
+            ],
+            "output": [],
+        }
+    )
+
+
+async def test_scheduled_resume_rejected_when_capacity_saturated(
+    tmp_path: Path,
+) -> None:
+    """A scheduled resume acquires the shared execution slot (B2).
+
+    With the only slot occupied by a gated scheduled run, resuming the
+    interrupted scheduled run is rejected BEFORE any dispatch: no ACTIVE
+    attempt, no executing mark, both runs untouched. Freeing the slot
+    lets the same resume through, releasing the slot afterwards.
+    """
+    _gate_open.clear()
+    root = tmp_path / "store"
+    server = build_local_static_workflow_server(root, extra_sources=_gate_sources())
+    await _seed(server, "ask1", "ask1.default", _single_interrupt_plan("ask1"), ["submitted"])
+    await _seed(server, "gated", "gated.default", _gate_plan("gated"), ["ok"])
+    store = FileScheduleStore(root)
+    store.create_schedule(_one_shot("asker", "ask1.default", datetime.now(UTC)))
+    service = _scheduler(server, capacity=1)
+    try:
+        await service.start()
+        await _wait_for(lambda: len(_run_ids(root)) == 1)
+        ask_id = _run_ids(root)[0]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        )
+        store.create_schedule(_one_shot("blocker", "gated.default", datetime.now(UTC)))
+        await _wait_for(lambda: len(_run_ids(root)) == 2)
+        gate_id = [rid for rid in _run_ids(root) if rid != ask_id][0]
+        await _wait_for(lambda: FileRunStore(root).is_executing(gate_id))
+
+        with pytest.raises(ScheduledCapacityBusyError):
+            await server.api.resume_run(
+                run_id=ask_id, resume_payload={}, resume_outcome="submitted"
+            )
+        assert FileRunStore(root).get_resume_attempt(ask_id) is None
+        assert not FileRunStore(root).is_executing(ask_id)
+        assert FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        assert FileRunStore(root).is_executing(gate_id)
+        assert len(_entries(root, "asker", "completed")) == 0
+
+        _gate_open.set()
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(gate_id).status.value == "completed"
+        )
+        resumed = await server.api.resume_run(
+            run_id=ask_id, resume_payload={}, resume_outcome="submitted"
+        )
+        assert resumed["status"] == "completed"
+        attempt = FileRunStore(root).get_resume_attempt(ask_id)
+        assert attempt is not None and attempt.state == "DONE"
+        assert not FileRunStore(root).is_executing(ask_id)
+    finally:
+        _gate_open.set()
+        await service.stop()
+
+
+async def test_scheduled_resume_releases_slot_on_reinterruption(
+    tmp_path: Path,
+) -> None:
+    """A scheduled resume that interrupts again frees its slot (B2).
+
+    The re-interruption stays durably resumable (attempt DONE, new
+    checkpoint) and a second resume of the same run completes.
+    """
+    root = tmp_path / "store"
+    server = build_local_static_workflow_server(root)
+    await _seed(
+        server,
+        "ask_twice",
+        "ask_twice.default",
+        _double_interrupt_plan("ask_twice"),
+        ["submitted"],
+    )
+    store = FileScheduleStore(root)
+    store.create_schedule(_one_shot("asker", "ask_twice.default", datetime.now(UTC)))
+    service = _scheduler(server, capacity=1)
+    try:
+        await service.start()
+        await _wait_for(lambda: len(_run_ids(root)) == 1)
+        ask_id = _run_ids(root)[0]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        )
+        first = await server.api.resume_run(
+            run_id=ask_id, resume_payload={}, resume_outcome="submitted"
+        )
+        assert first["status"] == "interrupted"
+        attempt = FileRunStore(root).get_resume_attempt(ask_id)
+        assert attempt is not None and attempt.state == "DONE"
+        assert not FileRunStore(root).is_executing(ask_id)
+        second = await server.api.resume_run(
+            run_id=ask_id, resume_payload={}, resume_outcome="submitted"
+        )
+        assert second["status"] == "completed"
+        assert not FileRunStore(root).is_executing(ask_id)
+    finally:
+        await service.stop()
+
+
+async def test_scheduled_resume_releases_slot_on_failure(tmp_path: Path) -> None:
+    """A scheduled resume that fails releases its slot (B2).
+
+    The failed stopped result is durably persisted with its attempt
+    cleared, and the freed slot admits new scheduled work.
+    """
+    root = tmp_path / "store"
+    server = build_local_static_workflow_server(root)
+    await _seed(
+        server,
+        "ask_fail",
+        "ask_fail.default",
+        _interrupt_then_fail_plan("ask_fail"),
+        ["submitted"],
+    )
+    await _seed(
+        server,
+        "const",
+        "const.default",
+        _constant_plan("const"),
+        ["ok"],
+    )
+    store = FileScheduleStore(root)
+    store.create_schedule(_one_shot("asker", "ask_fail.default", datetime.now(UTC)))
+    service = _scheduler(server, capacity=1)
+    try:
+        await service.start()
+        await _wait_for(lambda: len(_run_ids(root)) == 1)
+        ask_id = _run_ids(root)[0]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        )
+        resumed = await server.api.resume_run(
+            run_id=ask_id, resume_payload={}, resume_outcome="submitted"
+        )
+        assert resumed["status"] == "failed"
+        attempt = FileRunStore(root).get_resume_attempt(ask_id)
+        assert attempt is not None and attempt.state == "DONE"
+        assert not FileRunStore(root).is_executing(ask_id)
+        # The freed slot admits new scheduled work.
+        store.create_schedule(_one_shot("after", "const.default", datetime.now(UTC)))
+        await _wait_for(lambda: len(_run_ids(root)) == 2)
+        after_id = [rid for rid in _run_ids(root) if rid != ask_id][0]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(after_id).status.value == "completed"
+        )
+    finally:
+        await service.stop()
+
+
+async def test_service_stop_drains_inflight_scheduled_resume(tmp_path: Path) -> None:
+    """Shutdown waits for an in-flight scheduled resume (B2 drain)."""
+    _gate_open.clear()
+    root = tmp_path / "store"
+    server = build_local_static_workflow_server(root, extra_sources=_gate_sources())
+    await _seed(
+        server,
+        "ask_gate",
+        "ask_gate.default",
+        _interrupt_then_gate_plan("ask_gate"),
+        ["submitted"],
+    )
+    store = FileScheduleStore(root)
+    store.create_schedule(_one_shot("asker", "ask_gate.default", datetime.now(UTC)))
+    service = _scheduler(server, capacity=2, drain_grace_s=10.0)
+    try:
+        await service.start()
+        await _wait_for(lambda: len(_run_ids(root)) == 1)
+        ask_id = _run_ids(root)[0]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        )
+        resume_task = asyncio.create_task(
+            server.api.resume_run(
+                run_id=ask_id, resume_payload={}, resume_outcome="submitted"
+            )
+        )
+        await _wait_for(lambda: ask_id in service._live_resumes)
+        await _wait_for(lambda: FileRunStore(root).is_executing(ask_id))
+        stop_task = asyncio.create_task(service.stop())
+        await asyncio.sleep(0.3)
+        # The gate is still closed: the drain must still be waiting.
+        assert not stop_task.done()
+        assert not resume_task.done()
+        _gate_open.set()
+        report = await asyncio.wait_for(stop_task, timeout=20.0)
+        resumed = await asyncio.wait_for(resume_task, timeout=20.0)
+        assert resumed["status"] == "completed"
+        assert not FileRunStore(root).is_executing(ask_id)
+        assert report.cancelled == 0
+    finally:
+        _gate_open.set()
+        await service.stop()