215 lines
7.8 KiB
Python
215 lines
7.8 KiB
Python
"""Ownership guards at scheduler mutation/recovery/dispatch entries (R4/F5).
|
|
|
|
The held lock must be proven at the entry points, not merely exist as a
|
|
class: ``poll``, schedule administration, and ``recover`` reject when the
|
|
caller cannot prove live ownership, before any write or side effect.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from tests.scheduling.controlled import (
|
|
DictDeployments,
|
|
ScriptedDispatcher,
|
|
fixture_environment,
|
|
)
|
|
from tests.scheduling.controlled import ScriptedDispatcher as SD
|
|
from wf_artifacts.runs.store import FileRunStore
|
|
from wf_scheduling import recovery as sched_recovery
|
|
from wf_scheduling.calendar import OneShotSource
|
|
from wf_scheduling.models import Schedule
|
|
from wf_scheduling.ownership import SchedulerOwnership, SecondOwnerError
|
|
from wf_scheduling.poll import Scheduler
|
|
from wf_scheduling.prepare import SchedulePreparer
|
|
from wf_scheduling.store import FileScheduleStore
|
|
|
|
|
|
def ts(y: int, mo: int, d: int, h: int = 0, mi: int = 0) -> datetime:
|
|
return datetime(y, mo, d, h, mi, tzinfo=UTC)
|
|
|
|
|
|
def _sched_model(sid: str, **kw: Any) -> Schedule:
|
|
now = ts(2026, 9, 8, 12, 0)
|
|
base: dict[str, Any] = {
|
|
"id": sid,
|
|
"deployment_id": "dep-1",
|
|
"trigger": {"kind": "cron", "expression": "0 * * * *", "timezone": "UTC"},
|
|
"input_bindings": [],
|
|
"created_at": now.isoformat(),
|
|
"updated_at": now.isoformat(),
|
|
}
|
|
base.update(kw)
|
|
return Schedule.model_validate(base)
|
|
|
|
|
|
def _scheduler(
|
|
tmp_path: Path,
|
|
ownership: SchedulerOwnership,
|
|
*,
|
|
script: dict | None = None,
|
|
) -> tuple[Scheduler, FileScheduleStore, FileRunStore]:
|
|
sched_store = FileScheduleStore(tmp_path / "sched")
|
|
run_store = FileRunStore(tmp_path / "runs")
|
|
sched = Scheduler(
|
|
schedule_store=sched_store,
|
|
run_store=run_store,
|
|
sources={},
|
|
capacity=4,
|
|
preparer=SchedulePreparer(
|
|
DictDeployments({"dep-1": {"rev": 1, "required": []}}),
|
|
fixture_environment,
|
|
),
|
|
dispatcher=ScriptedDispatcher(script),
|
|
ownership=ownership,
|
|
)
|
|
return sched, sched_store, run_store
|
|
|
|
|
|
def _due_setup(sched: Scheduler, store: FileScheduleStore, intended: datetime) -> None:
|
|
store.create_schedule(_sched_model("a"))
|
|
store.save_consumed("a", intended - timedelta(hours=1))
|
|
sched.sources["a"] = OneShotSource(intended)
|
|
|
|
|
|
def test_poll_without_ownership_rejects_before_writes(tmp_path: Path) -> None:
|
|
calls: list[str] = []
|
|
|
|
def spy(admission: Any, now: datetime) -> Any:
|
|
calls.append(admission.id)
|
|
raise AssertionError("dispatcher must not run without ownership")
|
|
|
|
ownership = SchedulerOwnership(tmp_path, owner="never-acquired")
|
|
sched, store, runs = _scheduler(tmp_path, ownership)
|
|
sched.dispatcher = ScriptedDispatcher({"*": spy})
|
|
intended = ts(2026, 9, 8, 12, 0)
|
|
_due_setup(sched, store, intended)
|
|
with pytest.raises(SecondOwnerError):
|
|
sched.poll(intended)
|
|
assert calls == []
|
|
assert runs.list_runs() == []
|
|
assert runs.list_admissions() == []
|
|
assert store.list_occurrences("a", limit=100)["total"] == 0
|
|
assert store.get_consumed("a") == intended - timedelta(hours=1)
|
|
|
|
|
|
def test_poll_with_released_ownership_rejects(tmp_path: Path) -> None:
|
|
ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
|
|
sched, store, runs = _scheduler(tmp_path, ownership)
|
|
intended = ts(2026, 9, 8, 12, 0)
|
|
_due_setup(sched, store, intended)
|
|
ownership.release()
|
|
assert not ownership.held
|
|
with pytest.raises(SecondOwnerError):
|
|
sched.poll(intended)
|
|
assert runs.list_runs() == []
|
|
assert store.list_occurrences("a", limit=100)["total"] == 0
|
|
|
|
|
|
def test_recover_without_ownership_rejects_before_writes(tmp_path: Path) -> None:
|
|
from tests.artifacts.test_run_store import artifact as _artifact
|
|
from tests.artifacts.test_run_store import deployment as _deployment
|
|
from wf_api.run_lifecycle import materialize_admitted_view, persist_admission
|
|
from wf_artifacts import PinnedRunEnvironment
|
|
|
|
sched_store = FileScheduleStore(tmp_path / "sched")
|
|
run_store = FileRunStore(tmp_path / "runs")
|
|
sched_store.create_schedule(_sched_model("a"))
|
|
env = PinnedRunEnvironment(
|
|
deployment=_deployment(), root_artifact=_artifact(), child_artifacts=[]
|
|
)
|
|
admission = persist_admission(
|
|
store=run_store,
|
|
run_id=run_store.allocate_run_id(),
|
|
environment=env,
|
|
resolved_input={},
|
|
max_steps=None,
|
|
scheduled_at=ts(2026, 9, 8, 12, 0),
|
|
schedule_id="a",
|
|
schedule_revision=1,
|
|
)
|
|
materialize_admitted_view(store=run_store, admission=admission)
|
|
ownership = SchedulerOwnership(tmp_path, owner="never-acquired")
|
|
with pytest.raises(SecondOwnerError):
|
|
sched_recovery.recover(
|
|
schedule_store=sched_store,
|
|
run_store=run_store,
|
|
now=ts(2026, 9, 8, 12, 0),
|
|
ownership=ownership,
|
|
)
|
|
# Rejected before any write: the admitted run is untouched.
|
|
assert run_store.get_run(admission.id).status.value == "admitted"
|
|
|
|
|
|
def test_held_ownership_allows_poll_and_recovery(tmp_path: Path) -> None:
|
|
ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
|
|
try:
|
|
sched, store, runs = _scheduler(tmp_path, ownership, script={"*": "hang"})
|
|
intended = ts(2026, 9, 8, 12, 0)
|
|
_due_setup(sched, store, intended)
|
|
assert sched.poll(intended) == {"a": f"admit:{runs.list_runs()[0].id}"}
|
|
diags = sched_recovery.recover(
|
|
schedule_store=store,
|
|
run_store=runs,
|
|
now=intended,
|
|
ownership=ownership,
|
|
)
|
|
assert diags != []
|
|
finally:
|
|
ownership.release()
|
|
|
|
|
|
def test_second_owner_cannot_acquire_for_poll(tmp_path: Path) -> None:
|
|
first = SchedulerOwnership(tmp_path, owner="first").acquire()
|
|
try:
|
|
second = SchedulerOwnership(tmp_path, owner="second")
|
|
with pytest.raises(SecondOwnerError):
|
|
second.acquire()
|
|
assert not second.held
|
|
# A scheduler proving the loser's (unheld) ownership is rejected.
|
|
sched, store, _ = _scheduler(tmp_path, second)
|
|
intended = ts(2026, 9, 8, 12, 0)
|
|
_due_setup(sched, store, intended)
|
|
with pytest.raises(SecondOwnerError):
|
|
sched.poll(intended)
|
|
finally:
|
|
first.release()
|
|
|
|
|
|
def test_admin_mutations_require_ownership(tmp_path: Path) -> None:
|
|
ownership = SchedulerOwnership(tmp_path, owner="test")
|
|
sched, store, _ = _scheduler(tmp_path, ownership)
|
|
store.create_schedule(_sched_model("a", paused=True))
|
|
with pytest.raises(SecondOwnerError):
|
|
sched.resume_schedule("a", ts(2026, 9, 8, 12, 0))
|
|
with pytest.raises(SecondOwnerError):
|
|
sched.edit_schedule("a", ts(2026, 9, 8, 12, 0))
|
|
assert store.get_schedule("a").paused is True
|
|
assert store.get_schedule("a").revision == 1
|
|
|
|
|
|
def test_settle_requires_ownership(tmp_path: Path) -> None:
|
|
ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
|
|
try:
|
|
sched, store, runs = _scheduler(tmp_path, ownership, script={"*": "hang"})
|
|
intended = ts(2026, 9, 8, 12, 0)
|
|
_due_setup(sched, store, intended)
|
|
sched.poll(intended)
|
|
run_id = runs.list_runs()[0].id
|
|
assert runs.is_executing(run_id)
|
|
ownership.release()
|
|
admission = runs.get_admission(run_id)
|
|
state = cast(SD, sched.dispatcher).finish(admission, "complete")
|
|
with pytest.raises(SecondOwnerError):
|
|
sched.record_stopped_execution(
|
|
run_id, state, intended + timedelta(minutes=1)
|
|
)
|
|
assert runs.get_run(run_id).status.value == "admitted"
|
|
assert runs.is_executing(run_id)
|
|
finally:
|
|
ownership.release()
|