sched: fence scheduled resumes at shutdown so none outlive ownership

tests/wf_server/test_scheduler_integration.py

@@ -43,7 +43,10 @@ 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.resume_gate import (
+    ScheduledCapacityBusyError,
+    ScheduledResumeShutdownError,
+)
 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
@@ -1264,6 +1267,280 @@ async def test_service_stop_drains_inflight_scheduled_resume(tmp_path: Path) ->
         await service.stop()
 
 
+async def test_scheduled_resume_completes_within_grace_with_history(
+    tmp_path: Path,
+) -> None:
+    """A resume finishing inside grace stops cleanly with history (shutdown).
+
+    Operator-driven service (no background ticks): the in-flight resume
+    completes during the drain, stop reports no cancellation, and the
+    resumed completion lands in occurrence history exactly once.
+    """
+    _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=1, drain_grace_s=10.0, auto_tick=False)
+    try:
+        await service.start()
+        async with asyncio.timeout(20):
+            while not _run_ids(root):
+                await service.poll_once(datetime.now(UTC))
+                await asyncio.sleep(0.02)
+        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)
+        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()
+        resumed = await asyncio.wait_for(resume_task, timeout=20.0)
+        report = await asyncio.wait_for(stop_task, timeout=20.0)
+        assert resumed["status"] == "completed"
+        assert report.cancelled == 0
+        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)
+        assert await _kinds(root, "asker") == ["admitted", "interrupted", "completed"]
+    finally:
+        _gate_open.set()
+        await service.stop()
+
+
+async def test_shutdown_timeout_fences_inflight_resume_and_restart_keeps_decision(
+    tmp_path: Path,
+) -> None:
+    """A resume outliving grace cannot overwrite the next owner (shutdown).
+
+    Stop cancels and joins the gated resume before releasing ownership:
+    the task is done-cancelled at stop return, durable state keeps the
+    exact crash shape (interrupted, executing mark, ACTIVE attempt), and
+    after restart recovery fails the run closed the old task cannot flip
+    the decision back to completed or clear the ACTIVE attempt.
+    """
+    _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=1, drain_grace_s=0.05, auto_tick=False)
+    try:
+        await service.start()
+        async with asyncio.timeout(20):
+            while not _run_ids(root):
+                await service.poll_once(datetime.now(UTC))
+                await asyncio.sleep(0.02)
+        ask_id = _run_ids(root)[0]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        )
+        interrupted_checkpoint = FileRunStore(root).get_run(ask_id).latest_checkpoint_id
+        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))
+        # The gate stays closed: the drain must fence, not wait forever.
+        report = await service.stop()
+        assert report.cancelled == 1
+        assert resume_task.done()
+        assert resume_task.cancelled()
+        # Cancellation truthfulness: crash-shaped durable state, no result.
+        assert FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        assert FileRunStore(root).is_executing(ask_id)
+        attempt = FileRunStore(root).get_resume_attempt(ask_id)
+        assert attempt is not None and attempt.state == "ACTIVE"
+        assert await _kinds(root, "asker") == ["admitted", "interrupted"]
+
+        revived = _scheduler(server, capacity=1, auto_tick=False)
+        try:
+            await revived.start()
+            # Restart recovery fails the ambiguous resume closed.
+            assert FileRunStore(root).get_run(ask_id).status.value == "failed"
+            attempt2 = FileRunStore(root).get_resume_attempt(ask_id)
+            assert attempt2 is not None and attempt2.state == "ACTIVE"
+            assert not FileRunStore(root).is_executing(ask_id)
+            assert (
+                FileRunStore(root).get_run(ask_id).latest_checkpoint_id
+                == interrupted_checkpoint
+            )
+            assert len(_entries(root, "asker", "failed")) == 1
+            # The fenced old task is already dead: opening the gate and
+            # polling under the new owner changes nothing.
+            _gate_open.set()
+            await revived.poll_once(datetime.now(UTC))
+            assert FileRunStore(root).get_run(ask_id).status.value == "failed"
+            attempt3 = FileRunStore(root).get_resume_attempt(ask_id)
+            assert attempt3 is not None and attempt3.state == "ACTIVE"
+            assert _entries(root, "asker", "completed") == []
+            assert len(_entries(root, "asker", "failed")) == 1
+        finally:
+            _gate_open.set()
+            await revived.stop()
+    finally:
+        _gate_open.set()
+        await service.stop()
+
+
+async def test_scheduled_resume_requested_during_drain_rejected(
+    tmp_path: Path,
+) -> None:
+    """New scheduled resumes are rejected once shutdown begins (shutdown).
+
+    With a free slot available the rejection still fires (drain, not
+    capacity): no dispatch, no ACTIVE attempt, no executing mark. Manual
+    resumes bypass the gate and complete mid-drain.
+    """
+    _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"])
+    await _seed(
+        server,
+        "manual_ask",
+        "manual_ask.default",
+        _single_interrupt_plan("manual_ask"),
+        ["submitted"],
+    )
+    store = FileScheduleStore(root)
+    store.create_schedule(_one_shot("blocker", "gated.default", datetime.now(UTC)))
+    store.create_schedule(_one_shot("asker", "ask1.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)) == 2)
+        ids = {
+            FileRunStore(root).get_admission(rid).schedule_id or "": rid
+            for rid in _run_ids(root)
+        }
+        gate_id = ids["blocker"]
+        ask_id = ids["asker"]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        )
+        await _wait_for(lambda: FileRunStore(root).is_executing(gate_id))
+        asked = await server.api.run_deployment(
+            deployment_id="manual_ask.default", workflow_input={}
+        )
+        assert asked["status"] == "interrupted"
+        manual_id = asked["run_id"]
+        assert isinstance(manual_id, str)
+        # Hold the drain on the gated execution; one slot stays free, so
+        # only the shutdown check can reject the resume below.
+        stop_task = asyncio.create_task(service.stop())
+        await _wait_for(lambda: service._stopping)
+        assert not stop_task.done()
+        with pytest.raises(ScheduledResumeShutdownError):
+            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"
+        manual_resumed = await server.api.resume_run(
+            run_id=manual_id, resume_payload={}, resume_outcome="submitted"
+        )
+        assert manual_resumed["status"] == "completed"
+        _gate_open.set()
+        report = await asyncio.wait_for(stop_task, timeout=20.0)
+        assert report.cancelled == 0
+        assert FileRunStore(root).get_run(gate_id).status.value == "completed"
+    finally:
+        _gate_open.set()
+        await service.stop()
+
+
+async def test_shutdown_timeout_spares_healthy_sibling(tmp_path: Path) -> None:
+    """A fenced resume leaves a healthy sibling's outcome/history intact."""
+    _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"],
+    )
+    await _seed(server, "const", "const.default", _constant_plan("const"), ["ok"])
+    store = FileScheduleStore(root)
+    store.create_schedule(_one_shot("asker", "ask_gate.default", datetime.now(UTC)))
+    store.create_schedule(_one_shot("sib", "const.default", datetime.now(UTC)))
+    service = _scheduler(server, capacity=1, drain_grace_s=0.05, auto_tick=False)
+    try:
+        await service.start()
+        async with asyncio.timeout(20):
+            while len(_run_ids(root)) < 2:
+                await service.poll_once(datetime.now(UTC))
+                await asyncio.sleep(0.02)
+        ids = {
+            FileRunStore(root).get_admission(rid).schedule_id or "": rid
+            for rid in _run_ids(root)
+        }
+        ask_id = ids["asker"]
+        sib_id = ids["sib"]
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(ask_id).status.value == "interrupted"
+        )
+        await _wait_for(
+            lambda: FileRunStore(root).get_run(sib_id).status.value == "completed"
+        )
+        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)
+        report = await service.stop()
+        assert report.cancelled == 1
+        assert resume_task.cancelled()
+        assert FileRunStore(root).get_run(sib_id).status.value == "completed"
+        assert await _kinds(root, "sib") == ["admitted", "completed"]
+
+        revived = _scheduler(server, capacity=1, auto_tick=False)
+        try:
+            await revived.start()
+            assert FileRunStore(root).get_run(ask_id).status.value == "failed"
+            assert FileRunStore(root).get_run(sib_id).status.value == "completed"
+            assert await _kinds(root, "sib") == ["admitted", "completed"]
+            assert _entries(root, "asker", "completed") == []
+        finally:
+            _gate_open.set()
+            await revived.stop()
+    finally:
+        _gate_open.set()
+        await service.stop()
+
+
 async def _kinds(root: Path, schedule_id: str) -> list[str]:
     page = FileScheduleStore(root).list_occurrences(schedule_id, limit=100)
     rows = cast(list[dict[str, Any]], page["occurrences"])