46ea12f1a1
- Added event retry settings with parsed delay schedules and jittered delay computation. - Extended event persistence to persist an `available_at` timestamp, honor it during dequeuing, and clear it when re-queuing. - Updated worker failure handling to queue bounded retries with backoff and transition to failed only when the retry budget is exhausted.
295 lines
9.9 KiB
Python
295 lines
9.9 KiB
Python
"""Cancellation primitives on WorkerPool.
|
|
|
|
These tests stay at the public-ish surface of `WorkerPool` — they exercise the
|
|
hook registration contextvar that workers use and verify the dispatcher marks
|
|
cancelled events as failed with the documented marker. They do NOT spin up a
|
|
real omp subprocess; that's covered by the integration smoke test.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from robomp.cancellation import (
|
|
clear_current_event,
|
|
register_cancel_hook,
|
|
set_current_event,
|
|
unregister_cancel_hook,
|
|
)
|
|
from robomp.config import Settings
|
|
from robomp.db import Database, EventRow
|
|
from robomp.queue import WorkerPool
|
|
from robomp.slot_pool import SlotPool
|
|
|
|
|
|
class _StubGitHub:
|
|
"""Sentinel; queue tests don't talk to GitHub."""
|
|
|
|
|
|
class _StubSandbox:
|
|
"""Sentinel; queue tests don't touch the workspace pool."""
|
|
|
|
natives_cache = None
|
|
|
|
|
|
class _StubGitTransport:
|
|
"""Sentinel; queue tests don't push."""
|
|
|
|
|
|
def _make_pool(settings: Settings, db: Database) -> WorkerPool:
|
|
return WorkerPool(
|
|
settings=settings,
|
|
db=db,
|
|
github=_StubGitHub(), # type: ignore[arg-type]
|
|
sandbox=_StubSandbox(), # type: ignore[arg-type]
|
|
git_transport=_StubGitTransport(), # type: ignore[arg-type]
|
|
slot_pool=SlotPool(),
|
|
)
|
|
|
|
|
|
def _row(delivery: str = "d1") -> EventRow:
|
|
return EventRow(
|
|
delivery_id=delivery,
|
|
event_type="issues",
|
|
repo="octo/widget",
|
|
issue_key="octo/widget#1",
|
|
payload={"action": "opened"},
|
|
received_at="2026-01-01T00:00:00Z",
|
|
state="running",
|
|
attempts=1,
|
|
last_error=None,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_fires_hook_armed_by_worker(settings: Settings, db: Database) -> None:
|
|
"""A worker that armed a hook gets it invoked when cancel_event runs."""
|
|
pool = _make_pool(settings, db)
|
|
row = _row()
|
|
|
|
fired = asyncio.Event()
|
|
|
|
async def fake_worker() -> None:
|
|
# Mimic _run_event entering its contextvar scope: the helpers below are
|
|
# what worker.py invokes from inside the asyncio.to_thread call.
|
|
token = set_current_event(pool, row.delivery_id)
|
|
try:
|
|
await asyncio.to_thread(register_cancel_hook, fired.set)
|
|
# Park until somebody fires the hook.
|
|
await fired.wait()
|
|
finally:
|
|
await asyncio.to_thread(unregister_cancel_hook)
|
|
clear_current_event(token)
|
|
|
|
worker = asyncio.create_task(fake_worker())
|
|
# Give the worker a tick to register.
|
|
for _ in range(20):
|
|
await asyncio.sleep(0)
|
|
if row.delivery_id in pool._cancel_hooks: # noqa: SLF001 — test inspecting state
|
|
break
|
|
assert row.delivery_id in pool._cancel_hooks # noqa: SLF001
|
|
|
|
assert await pool.cancel_event(row.delivery_id) is True
|
|
await asyncio.wait_for(worker, timeout=1.0)
|
|
assert row.delivery_id in pool._cancelled # noqa: SLF001
|
|
# Hook is consumed.
|
|
assert row.delivery_id not in pool._cancel_hooks # noqa: SLF001
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_before_arm_fires_immediately(settings: Settings, db: Database) -> None:
|
|
"""Cancelling before the worker arms must still terminate it on register."""
|
|
pool = _make_pool(settings, db)
|
|
row = _row("d2")
|
|
|
|
# Cancel is requested before any worker has armed a hook.
|
|
assert await pool.cancel_event(row.delivery_id) is False
|
|
assert row.delivery_id in pool._cancelled # noqa: SLF001
|
|
|
|
# When the worker eventually registers, the hook must fire synchronously.
|
|
calls: list[int] = []
|
|
token = set_current_event(pool, row.delivery_id)
|
|
try:
|
|
register_cancel_hook(lambda: calls.append(1))
|
|
finally:
|
|
clear_current_event(token)
|
|
|
|
assert calls == [1]
|
|
# Late-armed hook is NOT retained; cancel state is one-shot.
|
|
assert row.delivery_id not in pool._cancel_hooks # noqa: SLF001
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_dispatch_marks_cancelled_event_failed_with_marker(
|
|
settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""A dispatch that observed cancellation marks the row failed + 'cancelled by operator'."""
|
|
pool = _make_pool(settings, db)
|
|
|
|
db.record_event(
|
|
delivery_id="d3",
|
|
event_type="issues",
|
|
repo="octo/widget",
|
|
issue_key="octo/widget#1",
|
|
payload={"action": "opened"},
|
|
state="running",
|
|
)
|
|
row = _row("d3")
|
|
|
|
async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None:
|
|
# Simulate cancellation hitting mid-task and the omp subprocess raising.
|
|
await pool.cancel_event(r.delivery_id)
|
|
raise RuntimeError("subprocess died")
|
|
|
|
monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch)
|
|
await pool._run_event(row) # noqa: SLF001 — testing the dispatcher branch directly
|
|
|
|
stored = db.get_event("d3")
|
|
assert stored is not None
|
|
assert stored.state == "failed"
|
|
assert stored.last_error == "cancelled by operator"
|
|
# State is cleared for future events.
|
|
assert row.delivery_id not in pool._cancelled # noqa: SLF001
|
|
assert row.delivery_id not in pool._cancel_hooks # noqa: SLF001
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_non_cancelled_failure_keeps_real_traceback(
|
|
settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""A garden-variety dispatch failure still records the traceback path."""
|
|
pool = _make_pool(settings, db)
|
|
monkeypatch.setattr(settings, "event_max_retries", 0) # assert terminal failure, not retry
|
|
db.record_event(
|
|
delivery_id="d4",
|
|
event_type="issues",
|
|
repo="octo/widget",
|
|
issue_key="octo/widget#1",
|
|
payload={"action": "opened"},
|
|
state="running",
|
|
)
|
|
row = _row("d4")
|
|
|
|
async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None:
|
|
raise ValueError("boom 42")
|
|
|
|
monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch)
|
|
await pool._run_event(row) # noqa: SLF001
|
|
|
|
stored = db.get_event("d4")
|
|
assert stored is not None
|
|
assert stored.state == "failed"
|
|
assert stored.last_error is not None
|
|
assert "boom 42" in stored.last_error
|
|
assert "cancelled by operator" not in stored.last_error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_event_marks_failed_when_not_shutting_down(
|
|
settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""When `_shutting_down` is False, a dispatch failure still marks the row failed."""
|
|
pool = _make_pool(settings, db)
|
|
monkeypatch.setattr(settings, "event_max_retries", 0) # assert terminal failure, not retry
|
|
assert pool._shutting_down is False # noqa: SLF001
|
|
db.record_event(
|
|
delivery_id="d5",
|
|
event_type="issues",
|
|
repo="octo/widget",
|
|
issue_key="octo/widget#1",
|
|
payload={"action": "opened"},
|
|
state="running",
|
|
)
|
|
row = _row("d5")
|
|
|
|
async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None:
|
|
raise RuntimeError("regular failure")
|
|
|
|
monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch)
|
|
await pool._run_event(row) # noqa: SLF001
|
|
|
|
stored = db.get_event("d5")
|
|
assert stored is not None
|
|
assert stored.state == "failed"
|
|
assert stored.last_error is not None
|
|
assert "regular failure" in stored.last_error
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_unknown_delivery_returns_false(settings: Settings, db: Database) -> None:
|
|
"""Cancelling an unknown delivery is a no-op that returns False."""
|
|
pool = _make_pool(settings, db)
|
|
assert await pool.cancel_event("never-existed") is False
|
|
# The set still records the request — a later register would fire — but
|
|
# since no worker is armed, the cancel is harmless.
|
|
assert "never-existed" in pool._cancelled # noqa: SLF001
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_start_reaps_configured_slot_uids(
|
|
settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
calls: list[int] = []
|
|
monkeypatch.setattr("robomp.queue._reap_slot", lambda uid: calls.append(uid))
|
|
pool = WorkerPool(
|
|
settings=settings,
|
|
db=db,
|
|
github=_StubGitHub(), # type: ignore[arg-type]
|
|
sandbox=_StubSandbox(), # type: ignore[arg-type]
|
|
git_transport=_StubGitTransport(), # type: ignore[arg-type]
|
|
slot_pool=SlotPool([2001, 2002]),
|
|
)
|
|
|
|
await pool.start()
|
|
try:
|
|
assert sorted(calls) == [2001, 2002]
|
|
finally:
|
|
await pool.stop(drain_timeout=0.01, kill_timeout=0.01)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_event_reaps_slot_before_release(
|
|
settings: Settings, db: Database, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
slot_pool = SlotPool([2001])
|
|
pool = WorkerPool(
|
|
settings=settings,
|
|
db=db,
|
|
github=_StubGitHub(), # type: ignore[arg-type]
|
|
sandbox=_StubSandbox(), # type: ignore[arg-type]
|
|
git_transport=_StubGitTransport(), # type: ignore[arg-type]
|
|
slot_pool=slot_pool,
|
|
)
|
|
db.record_event(
|
|
delivery_id="d-slot",
|
|
event_type="issues",
|
|
repo="octo/widget",
|
|
issue_key="octo/widget#1",
|
|
payload={"action": "opened"},
|
|
state="running",
|
|
)
|
|
order: list[tuple[str, int | None]] = []
|
|
monkeypatch.setattr("robomp.queue._reap_slot", lambda uid: order.append(("reap", uid)))
|
|
release = slot_pool.release
|
|
|
|
def record_release(slot_uid: int | None) -> None:
|
|
order.append(("release", slot_uid))
|
|
release(slot_uid)
|
|
|
|
monkeypatch.setattr(slot_pool, "release", record_release)
|
|
|
|
async def fake_dispatch(self: WorkerPool, r: EventRow, *, slot_uid: int | None = None) -> None:
|
|
assert r.delivery_id == "d-slot"
|
|
assert slot_uid == 2001
|
|
|
|
monkeypatch.setattr(WorkerPool, "_dispatch", fake_dispatch)
|
|
|
|
await pool._run_event(_row("d-slot")) # noqa: SLF001
|
|
|
|
stored = db.get_event("d-slot")
|
|
assert stored is not None
|
|
assert stored.state == "done"
|
|
assert order == [("reap", 2001), ("release", 2001)]
|