fix(jobs): make last-run writes atomic and drain them properly on shutdown
Two review findings on 22bee28, both reproduced before fixing.
[1] Concurrent writes could lose or rewind a row. record_finish did
select-then-insert-or-update, and pipelines are separate scheduler jobs that can
overlap while sharing step ids -- data_collector belongs to all four. Reproduced
both halves: two sessions that SELECT before either INSERTs make the second
commit raise IntegrityError, which _persist_job_run swallows, so the run
silently vanishes; and a later write carrying an OLDER finished_at rewound the
row from 12:00 back to 09:00, dragging the status with it, so the panel would
report a stale outcome as the latest.
Now a single atomic INSERT ... ON CONFLICT (job_name) DO UPDATE, guarded by
WHERE job_run_state.finished_at < excluded.finished_at so an older completion
can never overwrite a newer one. Dialect-specific because prod is Postgres and
tests are SQLite; both support it (SQLite >= 3.24, ours is 3.45). updated_at is
set explicitly, as the model's onupdate hook does not fire for a core upsert,
and get_map now uses populate_existing since core writes leave any
previously-loaded ORM instance stale in the identity map.
[2] Shutdown could dispose the engine underneath a pending write.
scheduler.shutdown(wait=False) returns before APScheduler dispatches its
completion events, and those events are what create persist tasks -- so a single
snapshot of the task set missed writes still to be queued. flush_job_run_persists
now settles briefly for pending callbacks, then drains in a loop until the set
stays empty, with the deadline still bounding total shutdown time. Left
shutdown(wait=False) alone deliberately: waiting would block a deploy restart
behind a long-running scan.
Six regression tests: interleaved first writes, older-never-rewinds,
newer-still-wins, a task queued mid-drain, prompt return when idle, and giving
up rather than hanging shutdown.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
+27
-4
@@ -353,11 +353,34 @@ def _schedule_persist(job_name: str) -> None:
|
|||||||
task.add_done_callback(_persist_tasks.discard)
|
task.add_done_callback(_persist_tasks.discard)
|
||||||
|
|
||||||
|
|
||||||
async def flush_job_run_persists(timeout: float = 5.0) -> None:
|
async def flush_job_run_persists(timeout: float = 5.0, settle: float = 0.05) -> None:
|
||||||
"""Await in-flight last-run writes. Called from the app's shutdown path."""
|
"""Drain last-run writes, including ones queued while we are draining.
|
||||||
if not _persist_tasks:
|
|
||||||
|
``scheduler.shutdown(wait=False)`` returns before APScheduler has dispatched
|
||||||
|
its job-completion events, and those events are what create persist tasks. A
|
||||||
|
single snapshot of the set therefore misses writes still to be queued, and
|
||||||
|
``engine.dispose()`` could then close the pool underneath them. So: give the
|
||||||
|
loop a moment for pending callbacks to land, then keep draining until the
|
||||||
|
set stays empty or the deadline passes.
|
||||||
|
"""
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
deadline = loop.time() + timeout
|
||||||
|
# Bounded settle so callbacks dispatched by shutdown get to queue their work
|
||||||
|
# before the first emptiness check decides there is nothing to wait for.
|
||||||
|
await asyncio.sleep(min(settle, timeout))
|
||||||
|
while True:
|
||||||
|
pending = {task for task in _persist_tasks if not task.done()}
|
||||||
|
if not pending:
|
||||||
return
|
return
|
||||||
await asyncio.wait(set(_persist_tasks), timeout=timeout)
|
remaining = deadline - loop.time()
|
||||||
|
if remaining <= 0:
|
||||||
|
logger.warning(
|
||||||
|
"Timed out draining %d last-run write(s); some may be lost", len(pending)
|
||||||
|
)
|
||||||
|
return
|
||||||
|
await asyncio.wait(pending, timeout=remaining)
|
||||||
|
# Loop rather than return: a completion callback may have queued another.
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
|
|
||||||
def get_job_runtime_snapshot(job_name: str | None = None) -> dict[str, dict[str, object]] | dict[str, object]:
|
def get_job_runtime_snapshot(job_name: str | None = None) -> dict[str, dict[str, object]] | dict[str, object]:
|
||||||
|
|||||||
@@ -11,6 +11,8 @@ from collections.abc import Iterable
|
|||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
from sqlalchemy import select
|
from sqlalchemy import select
|
||||||
|
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||||
|
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
from app.models.job_run_state import JobRunState
|
from app.models.job_run_state import JobRunState
|
||||||
@@ -31,40 +33,59 @@ def _as_datetime(value: object) -> datetime | None:
|
|||||||
|
|
||||||
|
|
||||||
async def get_map(db: AsyncSession, job_names: Iterable[str]) -> dict[str, JobRunState]:
|
async def get_map(db: AsyncSession, job_names: Iterable[str]) -> dict[str, JobRunState]:
|
||||||
"""Return {job_name: row} for the given jobs that have ever finished."""
|
"""Return {job_name: row} for the given jobs that have ever finished.
|
||||||
|
|
||||||
|
``populate_existing`` because rows are written by core upserts, which leave
|
||||||
|
any previously-loaded ORM instance in the identity map stale.
|
||||||
|
"""
|
||||||
result = await db.execute(
|
result = await db.execute(
|
||||||
select(JobRunState).where(JobRunState.job_name.in_(list(job_names)))
|
select(JobRunState)
|
||||||
|
.where(JobRunState.job_name.in_(list(job_names)))
|
||||||
|
.execution_options(populate_existing=True)
|
||||||
)
|
)
|
||||||
return {row.job_name: row for row in result.scalars().all()}
|
return {row.job_name: row for row in result.scalars().all()}
|
||||||
|
|
||||||
|
|
||||||
async def record_finish(db: AsyncSession, job_name: str, runtime: dict) -> JobRunState:
|
def _insert_for(db: AsyncSession):
|
||||||
|
"""ON CONFLICT is dialect-specific; prod is Postgres, tests are SQLite."""
|
||||||
|
dialect = db.get_bind().dialect.name
|
||||||
|
return pg_insert if dialect == "postgresql" else sqlite_insert
|
||||||
|
|
||||||
|
|
||||||
|
async def record_finish(db: AsyncSession, job_name: str, runtime: dict) -> None:
|
||||||
"""Upsert the last-run row from a scheduler runtime snapshot.
|
"""Upsert the last-run row from a scheduler runtime snapshot.
|
||||||
|
|
||||||
Select-then-update-or-insert rather than a dialect-specific upsert, matching
|
Atomic, and newer-wins. Select-then-insert loses races that really happen
|
||||||
``settings_store.upsert_setting`` — tests run on SQLite, prod on Postgres.
|
here: pipelines are separate scheduler jobs that can overlap, and they share
|
||||||
"""
|
step ids -- data_collector belongs to all four. Two of them finishing that
|
||||||
existing = await db.execute(
|
step together would both see no row and both insert, and the loser's
|
||||||
select(JobRunState).where(JobRunState.job_name == job_name)
|
IntegrityError is swallowed by the caller, so the run silently vanishes.
|
||||||
)
|
|
||||||
row = existing.scalar_one_or_none()
|
|
||||||
|
|
||||||
|
The ``where`` guard is the other half: without it a slower pipeline
|
||||||
|
finishing an *older* run last would rewind finished_at and the status with
|
||||||
|
it, so the panel would report a stale outcome as the latest one.
|
||||||
|
"""
|
||||||
finished_at = _as_datetime(runtime.get("finished_at")) or datetime.now(timezone.utc)
|
finished_at = _as_datetime(runtime.get("finished_at")) or datetime.now(timezone.utc)
|
||||||
status = str(runtime.get("status") or "completed")
|
|
||||||
message = runtime.get("message")
|
message = runtime.get("message")
|
||||||
|
now = datetime.now(timezone.utc)
|
||||||
values = {
|
values = {
|
||||||
"status": status,
|
"job_name": job_name,
|
||||||
|
"status": str(runtime.get("status") or "completed"),
|
||||||
"started_at": _as_datetime(runtime.get("started_at")),
|
"started_at": _as_datetime(runtime.get("started_at")),
|
||||||
"finished_at": finished_at,
|
"finished_at": finished_at,
|
||||||
"processed": runtime.get("processed"),
|
"processed": runtime.get("processed"),
|
||||||
"total": runtime.get("total"),
|
"total": runtime.get("total"),
|
||||||
"message": str(message)[:4000] if message else None,
|
"message": str(message)[:4000] if message else None,
|
||||||
|
# Set explicitly: the model's onupdate hook does not fire for a core
|
||||||
|
# INSERT ... ON CONFLICT DO UPDATE.
|
||||||
|
"updated_at": now,
|
||||||
}
|
}
|
||||||
|
|
||||||
if row is None:
|
statement = _insert_for(db)(JobRunState).values(**values)
|
||||||
row = JobRunState(job_name=job_name, **values)
|
await db.execute(
|
||||||
db.add(row)
|
statement.on_conflict_do_update(
|
||||||
else:
|
index_elements=[JobRunState.job_name],
|
||||||
for field, value in values.items():
|
set_={key: statement.excluded[key] for key in values if key != "job_name"},
|
||||||
setattr(row, field, value)
|
where=JobRunState.finished_at < statement.excluded.finished_at,
|
||||||
return row
|
)
|
||||||
|
)
|
||||||
|
|||||||
@@ -164,3 +164,108 @@ class TestListJobsSplitsLiveFromPersisted:
|
|||||||
assert job["runtime_status"] == "idle"
|
assert job["runtime_status"] == "idle"
|
||||||
assert job["running"] is False
|
assert job["running"] is False
|
||||||
scheduler.remove_all_jobs()
|
scheduler.remove_all_jobs()
|
||||||
|
|
||||||
|
|
||||||
|
class TestConcurrentWrites:
|
||||||
|
"""Pipelines are separate scheduler jobs that can overlap, and they share
|
||||||
|
step ids — data_collector belongs to all four."""
|
||||||
|
|
||||||
|
async def test_interleaved_first_writes_do_not_collide(self):
|
||||||
|
"""Both sessions SELECT before either INSERTs: select-then-insert lost
|
||||||
|
this race with an IntegrityError, and the caller swallows it."""
|
||||||
|
async with _test_session_factory() as a, _test_session_factory() as b:
|
||||||
|
await job_run_store.record_finish(
|
||||||
|
a, "data_collector",
|
||||||
|
{"status": "completed", "finished_at": "2026-08-08T10:00:00+00:00"},
|
||||||
|
)
|
||||||
|
await job_run_store.record_finish(
|
||||||
|
b, "data_collector",
|
||||||
|
{"status": "completed", "finished_at": "2026-08-08T10:00:01+00:00"},
|
||||||
|
)
|
||||||
|
await a.commit()
|
||||||
|
await b.commit() # must not raise
|
||||||
|
|
||||||
|
rows = await _committed_rows()
|
||||||
|
assert rows["data_collector"].status == "completed"
|
||||||
|
|
||||||
|
async def test_an_older_finish_never_rewinds_the_row(self):
|
||||||
|
"""A slower pipeline finishing an older run last must not overwrite a
|
||||||
|
newer outcome with a stale one."""
|
||||||
|
async with _test_session_factory() as s:
|
||||||
|
await job_run_store.record_finish(
|
||||||
|
s, "alerts",
|
||||||
|
{"status": "completed", "finished_at": "2026-08-08T12:00:00+00:00"},
|
||||||
|
)
|
||||||
|
await s.commit()
|
||||||
|
await job_run_store.record_finish(
|
||||||
|
s, "alerts",
|
||||||
|
{"status": "error", "finished_at": "2026-08-08T09:00:00+00:00", "message": "stale"},
|
||||||
|
)
|
||||||
|
await s.commit()
|
||||||
|
|
||||||
|
row = (await _committed_rows())["alerts"]
|
||||||
|
assert row.finished_at.isoformat().startswith("2026-08-08T12:00")
|
||||||
|
assert row.status == "completed"
|
||||||
|
assert row.message is None
|
||||||
|
|
||||||
|
async def test_a_newer_finish_still_wins(self):
|
||||||
|
async with _test_session_factory() as s:
|
||||||
|
await job_run_store.record_finish(
|
||||||
|
s, "rr_scanner",
|
||||||
|
{"status": "completed", "finished_at": "2026-08-08T09:00:00+00:00"},
|
||||||
|
)
|
||||||
|
await s.commit()
|
||||||
|
await job_run_store.record_finish(
|
||||||
|
s, "rr_scanner",
|
||||||
|
{"status": "error", "finished_at": "2026-08-08T12:00:00+00:00", "message": "boom"},
|
||||||
|
)
|
||||||
|
await s.commit()
|
||||||
|
|
||||||
|
row = (await _committed_rows())["rr_scanner"]
|
||||||
|
assert row.status == "error"
|
||||||
|
assert row.message == "boom"
|
||||||
|
|
||||||
|
|
||||||
|
class TestShutdownDrain:
|
||||||
|
async def test_drains_a_task_queued_after_the_flush_starts(self):
|
||||||
|
"""scheduler.shutdown(wait=False) returns before APScheduler dispatches
|
||||||
|
its completion events, so writes can appear mid-drain. Snapshotting the
|
||||||
|
task set once would miss them and dispose the engine underneath."""
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
done: list[str] = []
|
||||||
|
|
||||||
|
async def slow_first():
|
||||||
|
await asyncio.sleep(0.02)
|
||||||
|
done.append("first")
|
||||||
|
# Queued only once the first write is already finishing.
|
||||||
|
sched._persist_tasks.add(asyncio.get_running_loop().create_task(late()))
|
||||||
|
|
||||||
|
async def late():
|
||||||
|
await asyncio.sleep(0.02)
|
||||||
|
done.append("late")
|
||||||
|
|
||||||
|
sched._persist_tasks.clear()
|
||||||
|
sched._persist_tasks.add(asyncio.get_running_loop().create_task(slow_first()))
|
||||||
|
|
||||||
|
await sched.flush_job_run_persists(timeout=2.0)
|
||||||
|
|
||||||
|
assert done == ["first", "late"]
|
||||||
|
sched._persist_tasks.clear()
|
||||||
|
|
||||||
|
async def test_returns_promptly_when_there_is_nothing_to_drain(self):
|
||||||
|
sched._persist_tasks.clear()
|
||||||
|
await sched.flush_job_run_persists(timeout=2.0)
|
||||||
|
|
||||||
|
async def test_gives_up_rather_than_hanging_shutdown(self):
|
||||||
|
import asyncio
|
||||||
|
|
||||||
|
async def never():
|
||||||
|
await asyncio.sleep(30)
|
||||||
|
|
||||||
|
sched._persist_tasks.clear()
|
||||||
|
task = asyncio.get_running_loop().create_task(never())
|
||||||
|
sched._persist_tasks.add(task)
|
||||||
|
await sched.flush_job_run_persists(timeout=0.15) # returns, does not hang
|
||||||
|
task.cancel()
|
||||||
|
sched._persist_tasks.clear()
|
||||||
|
|||||||
Reference in New Issue
Block a user