fix: join offloaded scheduler settlement

This commit is contained in:
lda
2026-09-09 21:29:01 +07:00 Verified
parent 384f610301
commit c5087fe037
2 changed files with 137 additions and 10 deletions
+47 -10
View File
@@ -146,9 +146,12 @@ class RuntimeDispatcher:
a settlement failure such as a torn write — goes through the a settlement failure such as a torn write — goes through the
service-provided ``abandon`` hook, which runs the same per-run service-provided ``abandon`` hook, which runs the same per-run
reconciliation as startup recovery but scoped to the broken run id: reconciliation as startup recovery but scoped to the broken run id:
sibling executions are never touched. Cancellation (shutdown drain) is sibling executions are never touched. Cancellation before the runtime
re-raised unsettled: the durable executing mark is left for startup returns is re-raised unsettled: the durable executing mark is left for
recovery, which abandons it truthfully. startup recovery, which abandons it truthfully. Once execution has
returned, settlement or scoped recovery is joined to completion even if
its task is cancelled, because a synchronous store transition cannot be
allowed to outlive the lifecycle owner.
""" """
def __init__( def __init__(
@@ -195,14 +198,35 @@ class RuntimeDispatcher:
except asyncio.CancelledError: except asyncio.CancelledError:
raise raise
except Exception as exc: except Exception as exc:
self._abandon(admission.id, exc) await self._run_blocking_joined(self._abandon, admission.id, exc)
return return
try: try:
self._settle(admission.id, state) await self._run_blocking_joined(self._settle, admission.id, state)
except asyncio.CancelledError: except asyncio.CancelledError:
raise raise
except Exception as exc: except Exception as exc:
self._abandon(admission.id, exc) await self._run_blocking_joined(self._abandon, admission.id, exc)
async def _run_blocking_joined(
self,
operation: Callable[..., Any],
*args: Any,
) -> Any:
"""Run a sync lifecycle operation off-loop and join it if cancelled.
``asyncio.to_thread`` alone lets its worker continue after the
awaiting task is cancelled. The outer execution remains in
``SchedulerService._executions``; swallowing cancellation until this
worker finishes reuses that existing lifecycle tracking and prevents
persistence from racing shutdown ownership release.
"""
worker = asyncio.create_task(asyncio.to_thread(operation, *args))
while True:
try:
return await asyncio.shield(worker)
except asyncio.CancelledError:
if worker.done():
return await worker
@dataclass(slots=True) @dataclass(slots=True)
@@ -534,11 +558,24 @@ class SchedulerService:
"""Register one execution task; dispatch always spawns exactly one.""" """Register one execution task; dispatch always spawns exactly one."""
if self._loop is None: # pragma: no cover - start() sets this first if self._loop is None: # pragma: no cover - start() sets this first
raise SchedulerStartupError("scheduler service is not started") raise SchedulerStartupError("scheduler service is not started")
# Submitted from the poll worker thread: wrap explicitly against # Submitted from the poll worker thread: create and track the actual
# the server loop, which has no current-loop binding out there. # loop task before returning. Tracking a wrap_future proxy would let
tracked = asyncio.wrap_future( # cancellation mark the proxy done while the underlying coroutine
asyncio.run_coroutine_threadsafe(coro, self._loop), loop=self._loop # (and any joined settlement worker) was still unwinding.
registration = asyncio.run_coroutine_threadsafe(
self._register_execution(coro), self._loop
) )
try:
return registration.result()
except BaseException:
coro.close()
raise
async def _register_execution(
self, coro: Coroutine[Any, Any, None]
) -> asyncio.Task[Any]:
"""Create one execution task on the service event loop."""
tracked = asyncio.create_task(coro)
self._executions.add(tracked) self._executions.add(tracked)
tracked.add_done_callback(self._executions.discard) tracked.add_done_callback(self._executions.discard)
return tracked return tracked
+90
View File
@@ -11,6 +11,7 @@ dedicated integration tests.
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import threading
from datetime import UTC, datetime, timedelta from datetime import UTC, datetime, timedelta
from pathlib import Path from pathlib import Path
from typing import Any, cast from typing import Any, cast
@@ -238,6 +239,95 @@ async def test_scheduled_failure_records_failed_history(tmp_path: Path) -> None:
await service.stop() await service.stop()
async def test_settlement_does_not_block_the_event_loop(
tmp_path: Path, monkeypatch: Any
) -> None:
"""Synchronous settlement persistence must not stall other coroutines."""
intended = ts(2026, 9, 8, 12, 0)
service = _service(tmp_path, ScriptedRuntime("complete"))
settlement_started = threading.Event()
settlement_finished = threading.Event()
release = threading.Event()
original = Scheduler.record_stopped_execution
def _blocking_settle(self: Any, run_id: str, state: Any, now: datetime) -> None:
settlement_started.set()
release.wait(timeout=2)
original(self, run_id, state, now)
settlement_finished.set()
monkeypatch.setattr(Scheduler, "record_stopped_execution", _blocking_settle)
release_timer = threading.Timer(1.0, release.set)
release_timer.start()
try:
await service.start()
service.schedule_store.create_schedule(_sched_model("a", intended))
await service.poll_once(intended + timedelta(seconds=1))
await asyncio.sleep(0.05)
assert settlement_started.is_set()
assert not settlement_finished.is_set()
assert service.live_executions == 1
release.set()
await _wait_for(lambda: service.live_executions == 0)
finally:
release.set()
release_timer.cancel()
await service.stop()
async def test_shutdown_joins_cancelled_settlement_before_releasing_ownership(
tmp_path: Path, monkeypatch: Any
) -> None:
"""Shutdown waits for an offloaded settlement after cancelling its task."""
intended = ts(2026, 9, 8, 12, 0)
service = _service(
tmp_path,
ScriptedRuntime("complete"),
SchedulerServiceConfig(
poll_interval_s=0.01,
drain_grace_s=0,
auto_tick=False,
),
)
settlement_started = threading.Event()
release = threading.Event()
original = Scheduler.record_stopped_execution
def _blocking_settle(self: Any, run_id: str, state: Any, now: datetime) -> None:
settlement_started.set()
release.wait(timeout=2)
original(self, run_id, state, now)
monkeypatch.setattr(Scheduler, "record_stopped_execution", _blocking_settle)
release_timer = threading.Timer(1.0, release.set)
release_timer.start()
try:
await service.start()
service.schedule_store.create_schedule(_sched_model("a", intended))
await service.poll_once(intended + timedelta(seconds=1))
for _ in range(50):
if settlement_started.is_set():
break
await asyncio.sleep(0.01)
else:
pytest.fail("settlement did not start")
stop_task = asyncio.create_task(service.stop())
await asyncio.sleep(0.05)
assert not stop_task.done()
release.set()
report = await stop_task
assert report.settled == 1
assert report.cancelled == 0
assert service.run_store.get_run(
_only_run_id(service.run_store)
).status.value == ("completed")
finally:
release.set()
release_timer.cancel()
await service.stop()
async def test_scheduled_interrupt_stays_resumable(tmp_path: Path) -> None: async def test_scheduled_interrupt_stays_resumable(tmp_path: Path) -> None:
intended = ts(2026, 9, 8, 12, 0) intended = ts(2026, 9, 8, 12, 0)
service = _service(tmp_path, ScriptedRuntime("interrupt")) service = _service(tmp_path, ScriptedRuntime("interrupt"))