sched: fence scheduled resumes at shutdown so none outlive ownership
This commit is contained in:
@@ -96,8 +96,11 @@ retrying it. Corrupt or contradictory records fail closed with
|
|||||||
diagnostics and block the schedule rather than clearing overlap.
|
diagnostics and block the schedule rather than clearing overlap.
|
||||||
|
|
||||||
On shutdown the server stops admission first and drains active tasks
|
On shutdown the server stops admission first and drains active tasks
|
||||||
within the grace period; anything still running keeps its executing
|
within the grace period; new scheduled resumes are rejected for the
|
||||||
mark, and the next startup recovery abandons it truthfully.
|
duration of the drain, and a resume still running past the deadline is
|
||||||
|
cancelled and joined before ownership is released, keeping its marks
|
||||||
|
for the next startup recovery. Anything else still running keeps its
|
||||||
|
executing mark, and the next startup recovery abandons it truthfully.
|
||||||
|
|
||||||
## Occurrence inspection
|
## Occurrence inspection
|
||||||
|
|
||||||
@@ -194,8 +197,9 @@ await schedules.update_schedule(
|
|||||||
- Manual runs bypass scheduler capacity by design; capacity governs
|
- Manual runs bypass scheduler capacity by design; capacity governs
|
||||||
scheduled dispatch plus scheduled resumes. Resuming a scheduled
|
scheduled dispatch plus scheduled resumes. Resuming a scheduled
|
||||||
interrupted run acquires a server execution slot first (rejected
|
interrupted run acquires a server execution slot first (rejected
|
||||||
while saturated, without dispatching) and releases it when the
|
while saturated or draining, without dispatching) and releases it when the
|
||||||
resumed result is persisted.
|
resumed result is persisted; a drain-cancelled resume keeps its marks
|
||||||
|
for recovery instead.
|
||||||
- A set `max_steps` budget cannot be cleared back to unset through
|
- A set `max_steps` budget cannot be cleared back to unset through
|
||||||
update (recreate the schedule for an unbounded budget).
|
update (recreate the schedule for an unbounded budget).
|
||||||
- MCP-backed servers reject scheduler enablement for now.
|
- MCP-backed servers reject scheduler enablement for now.
|
||||||
|
|||||||
@@ -310,7 +310,13 @@ server capacity default is deployment configuration, with deterministic tests
|
|||||||
using a small injected limit.
|
using a small injected limit.
|
||||||
|
|
||||||
On shutdown stop admission first and drain active tasks within a configured
|
On shutdown stop admission first and drain active tasks within a configured
|
||||||
grace period. Record cancellation/failure when possible; abrupt termination
|
grace period. New scheduled resumes are rejected once shutdown begins
|
||||||
|
(no dispatch, no ungated fallback); in-flight scheduled resumes share
|
||||||
|
the grace window, then are cancelled and joined before ownership is
|
||||||
|
released, so no old execution persists after a new owner takes over.
|
||||||
|
A cancelled resume keeps its executing mark and ACTIVE attempt, and
|
||||||
|
startup recovery fails it closed exactly like a crash mid-resume.
|
||||||
|
Record cancellation/failure when possible; abrupt termination
|
||||||
uses startup recovery. Paused/deleted schedule definitions must not prevent
|
uses startup recovery. Paused/deleted schedule definitions must not prevent
|
||||||
run completion or resume from updating retained occurrence history.
|
run completion or resume from updating retained occurrence history.
|
||||||
|
|
||||||
@@ -321,9 +327,11 @@ canonical ownership, recovers without executing, then ticks calendar
|
|||||||
polling without blocking on long workflows: each dispatch spawns exactly
|
polling without blocking on long workflows: each dispatch spawns exactly
|
||||||
one bounded execution task behind the async-completion seam, and the
|
one bounded execution task behind the async-completion seam, and the
|
||||||
scheduler's own capacity gate is the execution-slot bound (an executing
|
scheduler's own capacity gate is the execution-slot bound (an executing
|
||||||
run keeps its slot until it stops). Shutdown stops admission, drains
|
run keeps its slot until it stops). Shutdown stops admission, rejects
|
||||||
within `drain_grace_s`, leaves unfinished work under its executing mark
|
new scheduled resumes for the duration of the drain, cancels and joins
|
||||||
for startup recovery to abandon truthfully, and releases ownership last.
|
unfinished scheduled resumes after `drain_grace_s` (cancelled work keeps
|
||||||
|
its executing mark and ACTIVE attempt for startup recovery to abandon
|
||||||
|
truthfully), and releases ownership last.
|
||||||
A failed startup releases the lock and raises.
|
A failed startup releases the lock and raises.
|
||||||
|
|
||||||
Administration (`WorkflowApi` schedules methods, `workflow.schedules.*`
|
Administration (`WorkflowApi` schedules methods, `workflow.schedules.*`
|
||||||
@@ -344,7 +352,9 @@ resumed through the run API acquires a server execution slot through the
|
|||||||
scheduler's own accounting before dispatch — rejection leaves no resume
|
scheduler's own accounting before dispatch — rejection leaves no resume
|
||||||
attempt behind — holds the durable executing mark for the re-execution
|
attempt behind — holds the durable executing mark for the re-execution
|
||||||
(visible to capacity and drain like any live execution), and releases
|
(visible to capacity and drain like any live execution), and releases
|
||||||
the slot when its stopped result is persisted. A resumed scheduled run
|
the slot when its stopped result is persisted. Shutdown drain rejects
|
||||||
|
further scheduled resumes; a resume cancelled by the drain keeps its
|
||||||
|
marks for recovery instead of releasing them. A resumed scheduled run
|
||||||
reconciles its terminal history live through the same idempotent
|
reconciles its terminal history live through the same idempotent
|
||||||
recording as dispatch; restart recovery still repairs torn boundaries.
|
recording as dispatch; restart recovery still repairs torn boundaries.
|
||||||
|
|
||||||
|
|||||||
+33
-6
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
from dataclasses import asdict
|
from dataclasses import asdict
|
||||||
from typing import Any, Protocol
|
from typing import Any, Protocol
|
||||||
|
|
||||||
@@ -198,14 +199,20 @@ class WorkflowRunApi:
|
|||||||
# Shared execution slot for schedule-owned resumes: the resume
|
# Shared execution slot for schedule-owned resumes: the resume
|
||||||
# gate holds the scheduler's own capacity accounting (no second
|
# gate holds the scheduler's own capacity accounting (no second
|
||||||
# semaphore) behind the durable executing mark. Acquisition runs
|
# semaphore) behind the durable executing mark. Acquisition runs
|
||||||
# BEFORE the ACTIVE attempt mark, so a busy rejection leaves no
|
# BEFORE the ACTIVE attempt mark, so a busy or draining rejection
|
||||||
# fake ACTIVE attempt for work that never dispatched. Manual runs
|
# leaves no fake ACTIVE attempt for work that never dispatched.
|
||||||
# (no schedule admission, or no live scheduler) skip the gate.
|
# The caller's task is bound atomically with acquisition, so the
|
||||||
|
# shutdown drain can cancel and join it after the grace deadline.
|
||||||
|
# Manual runs (no schedule admission, or no live scheduler) skip
|
||||||
|
# the gate.
|
||||||
gate = self.resume_slot_gate
|
gate = self.resume_slot_gate
|
||||||
slot = gate.acquire(run_id) if gate is not None else None
|
if gate is not None:
|
||||||
|
slot = gate.acquire(run_id, owner_task=asyncio.current_task())
|
||||||
|
else:
|
||||||
|
slot = None
|
||||||
slot_held = slot is not None
|
slot_held = slot is not None
|
||||||
try:
|
try:
|
||||||
return await self._resume_scheduled_or_manual(
|
result = await self._resume_scheduled_or_manual(
|
||||||
run_id=run_id,
|
run_id=run_id,
|
||||||
resume_payload=resume_payload,
|
resume_payload=resume_payload,
|
||||||
resume_outcome=resume_outcome,
|
resume_outcome=resume_outcome,
|
||||||
@@ -213,9 +220,29 @@ class WorkflowRunApi:
|
|||||||
trace_values=trace_values,
|
trace_values=trace_values,
|
||||||
store=store,
|
store=store,
|
||||||
)
|
)
|
||||||
finally:
|
except asyncio.CancelledError:
|
||||||
|
# Shutdown drain cancelled the execution mid-flight (or the
|
||||||
|
# caller went away): fence, don't release. The durable
|
||||||
|
# executing mark and the ACTIVE attempt stay exactly as a
|
||||||
|
# crash mid-resume would leave them, so restart recovery
|
||||||
|
# fails the run closed instead of presenting a half-resumed
|
||||||
|
# run as safe to retry. Cancellation is never swallowed.
|
||||||
|
if slot_held and gate is not None:
|
||||||
|
gate.fence(run_id)
|
||||||
|
raise
|
||||||
|
except BaseException:
|
||||||
|
# Pre-persist errors, torn persists, and validation failures:
|
||||||
|
# release is truthful here because either nothing executed
|
||||||
|
# (the marks were only ever ours) or the durable ACTIVE
|
||||||
|
# attempt / stopped result already captures the ambiguity for
|
||||||
|
# recovery to reconcile by attempt identity.
|
||||||
if slot_held and gate is not None:
|
if slot_held and gate is not None:
|
||||||
gate.release(run_id)
|
gate.release(run_id)
|
||||||
|
raise
|
||||||
|
else:
|
||||||
|
if slot_held and gate is not None:
|
||||||
|
gate.release(run_id)
|
||||||
|
return result
|
||||||
|
|
||||||
async def _resume_scheduled_or_manual(
|
async def _resume_scheduled_or_manual(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -224,7 +224,11 @@ class SchedulerService:
|
|||||||
acquire the same slot through the resume gate (same load count, same
|
acquire the same slot through the resume gate (same load count, same
|
||||||
durable executing mark, tracked for drain); genuinely manual runs —
|
durable executing mark, tracked for drain); genuinely manual runs —
|
||||||
those with no schedule admission — bypass the scheduler entirely
|
those with no schedule admission — bypass the scheduler entirely
|
||||||
(unchanged API behavior).
|
(unchanged API behavior, including while draining). Once shutdown
|
||||||
|
begins the gate rejects new scheduled resumes instead of dispatching
|
||||||
|
them; in-flight resumes get the grace window, then are cancelled and
|
||||||
|
joined before ownership is released, so no old execution can persist
|
||||||
|
after a new owner takes over.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
schedule_store: Any
|
schedule_store: Any
|
||||||
@@ -254,6 +258,9 @@ class SchedulerService:
|
|||||||
_live_resumes: set[str] = field(
|
_live_resumes: set[str] = field(
|
||||||
default_factory=set, init=False, repr=False, compare=False
|
default_factory=set, init=False, repr=False, compare=False
|
||||||
)
|
)
|
||||||
|
_resume_tasks: dict[str, asyncio.Task[Any]] = field(
|
||||||
|
default_factory=dict, init=False, repr=False, compare=False
|
||||||
|
)
|
||||||
_stopping: bool = field(default=False, init=False, repr=False, compare=False)
|
_stopping: bool = field(default=False, init=False, repr=False, compare=False)
|
||||||
_started: bool = field(default=False, init=False, repr=False, compare=False)
|
_started: bool = field(default=False, init=False, repr=False, compare=False)
|
||||||
_tick_count: int = field(default=0, init=False, repr=False, compare=False)
|
_tick_count: int = field(default=0, init=False, repr=False, compare=False)
|
||||||
@@ -346,15 +353,24 @@ class SchedulerService:
|
|||||||
|
|
||||||
No new tick starts after ``stop`` begins (an already-running poll
|
No new tick starts after ``stop`` begins (an already-running poll
|
||||||
is joined, and any dispatch it performed is covered by the
|
is joined, and any dispatch it performed is covered by the
|
||||||
drain). Live executions get the configured grace period; tasks
|
drain), and the resume gate rejects new scheduled resumes from
|
||||||
still running afterwards are cancelled and left under their
|
that point on. Live executions and in-flight scheduled resumes
|
||||||
durable executing mark for startup recovery to abandon
|
get the configured grace period; anything still running
|
||||||
truthfully. Ownership is released only after the drain finishes.
|
afterwards is cancelled — executions are left under their durable
|
||||||
Safe to call when not started (returns a zero report).
|
executing mark, and resumes keep their executing mark plus ACTIVE
|
||||||
|
attempt via the gate fence — for startup recovery to abandon
|
||||||
|
truthfully. Every tracked resume task is joined before ownership
|
||||||
|
is released, so no old execution can mutate the stores after a
|
||||||
|
new owner takes over. Ownership is released only after the drain
|
||||||
|
finishes. Safe to call when not started (returns a zero report).
|
||||||
"""
|
"""
|
||||||
if not self._started:
|
if not self._started:
|
||||||
return DrainReport()
|
return DrainReport()
|
||||||
self._stopping = True
|
with self._lock:
|
||||||
|
# Under the service lock, atomically with gate acquisition:
|
||||||
|
# every resume that passed the stopping check is already
|
||||||
|
# tracked, and every later one is rejected.
|
||||||
|
self._stopping = True
|
||||||
if self._tick_task is not None:
|
if self._tick_task is not None:
|
||||||
self._tick_task.cancel()
|
self._tick_task.cancel()
|
||||||
try:
|
try:
|
||||||
@@ -374,6 +390,26 @@ class SchedulerService:
|
|||||||
grace_deadline = time.monotonic() + self.config.drain_grace_s
|
grace_deadline = time.monotonic() + self.config.drain_grace_s
|
||||||
while self._live_resumes and time.monotonic() < grace_deadline:
|
while self._live_resumes and time.monotonic() < grace_deadline:
|
||||||
await asyncio.sleep(0.02)
|
await asyncio.sleep(0.02)
|
||||||
|
# Fence: cancel and join every still-tracked resume task BEFORE
|
||||||
|
# releasing ownership. Joining first is what keeps a new owner
|
||||||
|
# safe: a joined task cannot persist afterwards. Cancelled tasks
|
||||||
|
# unwind through the gate fence (in-memory accounting dropped,
|
||||||
|
# durable marks kept), so recovery sees the exact crash shape.
|
||||||
|
current = asyncio.current_task()
|
||||||
|
with self._lock:
|
||||||
|
resume_pending = [
|
||||||
|
task
|
||||||
|
for task in self._resume_tasks.values()
|
||||||
|
if task is not current and not task.done()
|
||||||
|
]
|
||||||
|
for task in resume_pending:
|
||||||
|
task.cancel()
|
||||||
|
if resume_pending:
|
||||||
|
await asyncio.gather(*resume_pending, return_exceptions=True)
|
||||||
|
with self._lock:
|
||||||
|
self._resume_tasks.clear()
|
||||||
|
self._live_resumes.clear()
|
||||||
|
resume_cancelled = len([t for t in resume_pending if t.cancelled()])
|
||||||
pending = list(self._executions)
|
pending = list(self._executions)
|
||||||
if pending:
|
if pending:
|
||||||
remaining = max(0.0, grace_deadline - time.monotonic())
|
remaining = max(0.0, grace_deadline - time.monotonic())
|
||||||
@@ -387,7 +423,7 @@ class SchedulerService:
|
|||||||
# observe the drained set deterministically.
|
# observe the drained set deterministically.
|
||||||
for task in pending:
|
for task in pending:
|
||||||
self._executions.discard(task)
|
self._executions.discard(task)
|
||||||
cancelled = len([t for t in pending if t.cancelled()])
|
cancelled = len([t for t in pending if t.cancelled()]) + resume_cancelled
|
||||||
with self._lock:
|
with self._lock:
|
||||||
report = DrainReport(
|
report = DrainReport(
|
||||||
settled=self._settled,
|
settled=self._settled,
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ rejection leaves no fake ACTIVE attempt behind.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
from wf_scheduling.ownership import SecondOwnerError
|
from wf_scheduling.ownership import SecondOwnerError
|
||||||
@@ -27,6 +28,15 @@ class ScheduledCapacityBusyError(ValueError):
|
|||||||
"""A scheduled resume found every server execution slot occupied."""
|
"""A scheduled resume found every server execution slot occupied."""
|
||||||
|
|
||||||
|
|
||||||
|
class ScheduledResumeShutdownError(ValueError):
|
||||||
|
"""A scheduled resume was requested while the scheduler is draining.
|
||||||
|
|
||||||
|
Raised instead of acquiring once shutdown begins: a draining service
|
||||||
|
must neither dispatch new scheduled work nor fall back to ungated
|
||||||
|
execution, which would bypass capacity and escape drain tracking.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
class SchedulerResumeGate:
|
class SchedulerResumeGate:
|
||||||
"""Execution-slot gate binding scheduled resumes to scheduler capacity.
|
"""Execution-slot gate binding scheduled resumes to scheduler capacity.
|
||||||
|
|
||||||
@@ -39,16 +49,29 @@ class SchedulerResumeGate:
|
|||||||
def __init__(self, service: Any) -> None:
|
def __init__(self, service: Any) -> None:
|
||||||
self._service = service
|
self._service = service
|
||||||
|
|
||||||
def acquire(self, run_id: str) -> Any | None:
|
def acquire(
|
||||||
|
self, run_id: str, *, owner_task: asyncio.Task[Any] | None = None
|
||||||
|
) -> Any | None:
|
||||||
"""Take the shared execution slot for a scheduled resume.
|
"""Take the shared execution slot for a scheduled resume.
|
||||||
|
|
||||||
Returns the schedule admission when the slot is held (the caller
|
Returns the schedule admission when the slot is held (the caller
|
||||||
must :meth:`release` it), or ``None`` when this run is genuinely
|
must :meth:`release` it, or :meth:`fence` it when cancelled), or
|
||||||
manual, not resumable, or the scheduler is not live — all keep
|
``None`` when this run is genuinely manual, not resumable, or the
|
||||||
the legacy resume path unchanged. Raises
|
scheduler is not live — all keep the legacy resume path
|
||||||
:class:`ScheduledCapacityBusyError` when every slot is occupied;
|
unchanged. Raises :class:`ScheduledCapacityBusyError` when every
|
||||||
nothing is marked in that case, in particular no ACTIVE resume
|
slot is occupied, and :class:`ScheduledResumeShutdownError` once
|
||||||
attempt.
|
shutdown begins; nothing is marked in either case, in particular
|
||||||
|
no ACTIVE resume attempt. The shutdown check sits after the
|
||||||
|
manual/not-resumable early returns, so manual resumes stay
|
||||||
|
independent of the drain.
|
||||||
|
|
||||||
|
``owner_task`` binds the caller's execution lifetime for the
|
||||||
|
shutdown drain: the service cancels and joins tracked tasks after
|
||||||
|
the grace deadline, before releasing ownership. Registration
|
||||||
|
happens here, under the service lock, atomically with the
|
||||||
|
stopping check — a resume that passes the check is always
|
||||||
|
visible to the drain, and a resume that loses the race is
|
||||||
|
rejected instead of running unfenced.
|
||||||
"""
|
"""
|
||||||
service = self._service
|
service = self._service
|
||||||
with service._lock:
|
with service._lock:
|
||||||
@@ -72,6 +95,11 @@ class SchedulerResumeGate:
|
|||||||
scheduler = service._scheduler
|
scheduler = service._scheduler
|
||||||
if scheduler is None:
|
if scheduler is None:
|
||||||
return None
|
return None
|
||||||
|
if service._stopping:
|
||||||
|
raise ScheduledResumeShutdownError(
|
||||||
|
f"scheduler is draining for shutdown; "
|
||||||
|
f"scheduled resume of {run_id!r} rejected without dispatch"
|
||||||
|
)
|
||||||
if scheduler._task_load() >= service.config.capacity:
|
if scheduler._task_load() >= service.config.capacity:
|
||||||
raise ScheduledCapacityBusyError(
|
raise ScheduledCapacityBusyError(
|
||||||
f"server execution capacity is saturated "
|
f"server execution capacity is saturated "
|
||||||
@@ -83,14 +111,40 @@ class SchedulerResumeGate:
|
|||||||
# and the poll loop counts the resume as live load.
|
# and the poll loop counts the resume as live load.
|
||||||
service.run_store.mark_executing(run_id)
|
service.run_store.mark_executing(run_id)
|
||||||
service._live_resumes.add(run_id)
|
service._live_resumes.add(run_id)
|
||||||
|
if owner_task is not None:
|
||||||
|
service._resume_tasks[run_id] = owner_task
|
||||||
return admission
|
return admission
|
||||||
|
|
||||||
def release(self, run_id: str) -> None:
|
def release(self, run_id: str) -> None:
|
||||||
"""Free the slot held by :meth:`acquire` after the stopped persist."""
|
"""Free the slot held by :meth:`acquire` after the stopped persist.
|
||||||
|
|
||||||
|
Correct for every non-cancellation outcome: on success the
|
||||||
|
stopped result plus the DONE attempt already capture the outcome;
|
||||||
|
on pre-persist errors the marks were only ever ours (the durable
|
||||||
|
ACTIVE attempt still records the ambiguity for recovery); on torn
|
||||||
|
persists the durable stopped result plus ACTIVE attempt let
|
||||||
|
recovery reconcile by attempt identity.
|
||||||
|
"""
|
||||||
service = self._service
|
service = self._service
|
||||||
with service._lock:
|
with service._lock:
|
||||||
service.run_store.clear_executing(run_id)
|
service.run_store.clear_executing(run_id)
|
||||||
service._live_resumes.discard(run_id)
|
service._live_resumes.discard(run_id)
|
||||||
|
service._resume_tasks.pop(run_id, None)
|
||||||
|
|
||||||
|
def fence(self, run_id: str) -> None:
|
||||||
|
"""Forget a cancelled execution; keep its durable marks for recovery.
|
||||||
|
|
||||||
|
The shutdown drain cancels a resume that outlives the grace
|
||||||
|
deadline and joins it before releasing ownership, so no old
|
||||||
|
execution can persist afterwards. The durable executing mark and
|
||||||
|
the ACTIVE attempt stay exactly as a crash mid-resume would leave
|
||||||
|
them: restart recovery fails the run closed (ambiguous attempt),
|
||||||
|
never silently resumable, never replayed.
|
||||||
|
"""
|
||||||
|
service = self._service
|
||||||
|
with service._lock:
|
||||||
|
service._live_resumes.discard(run_id)
|
||||||
|
service._resume_tasks.pop(run_id, None)
|
||||||
|
|
||||||
def note_resumed_result(
|
def note_resumed_result(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -43,7 +43,10 @@ from wf_platform import (
|
|||||||
)
|
)
|
||||||
from wf_scheduling.lifecycle import SchedulerServiceConfig, SchedulerStartupError
|
from wf_scheduling.lifecycle import SchedulerServiceConfig, SchedulerStartupError
|
||||||
from wf_scheduling.models import Schedule
|
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_scheduling.store import FileScheduleStore, ScheduleExistsError
|
||||||
from wf_server import WorkflowServer, build_local_static_workflow_server
|
from wf_server import WorkflowServer, build_local_static_workflow_server
|
||||||
from wf_server.scheduling import build_scheduler_service
|
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()
|
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]:
|
async def _kinds(root: Path, schedule_id: str) -> list[str]:
|
||||||
page = FileScheduleStore(root).list_occurrences(schedule_id, limit=100)
|
page = FileScheduleStore(root).list_occurrences(schedule_id, limit=100)
|
||||||
rows = cast(list[dict[str, Any]], page["occurrences"])
|
rows = cast(list[dict[str, Any]], page["occurrences"])
|
||||||
|
|||||||
Reference in New Issue
Block a user