From 94baa894237bc18b3542104189e519298c39dd59 Mon Sep 17 00:00:00 2001 From: Dennis Thiessen Date: Sat, 8 Aug 2026 14:15:05 +0200 Subject: [PATCH] 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 --- app/scheduler.py | 33 ++++++-- app/services/job_run_store.py | 59 +++++++++----- tests/unit/test_job_run_persistence.py | 105 +++++++++++++++++++++++++ 3 files changed, 173 insertions(+), 24 deletions(-) diff --git a/app/scheduler.py b/app/scheduler.py index 5868985..6986747 100644 --- a/app/scheduler.py +++ b/app/scheduler.py @@ -353,11 +353,34 @@ def _schedule_persist(job_name: str) -> None: task.add_done_callback(_persist_tasks.discard) -async def flush_job_run_persists(timeout: float = 5.0) -> None: - """Await in-flight last-run writes. Called from the app's shutdown path.""" - if not _persist_tasks: - return - await asyncio.wait(set(_persist_tasks), timeout=timeout) +async def flush_job_run_persists(timeout: float = 5.0, settle: float = 0.05) -> None: + """Drain last-run writes, including ones queued while we are draining. + + ``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 + 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]: diff --git a/app/services/job_run_store.py b/app/services/job_run_store.py index 5a59534..532bdc6 100644 --- a/app/services/job_run_store.py +++ b/app/services/job_run_store.py @@ -11,6 +11,8 @@ from collections.abc import Iterable from datetime import datetime, timezone 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 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]: - """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( - 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()} -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. - Select-then-update-or-insert rather than a dialect-specific upsert, matching - ``settings_store.upsert_setting`` — tests run on SQLite, prod on Postgres. - """ - existing = await db.execute( - select(JobRunState).where(JobRunState.job_name == job_name) - ) - row = existing.scalar_one_or_none() + Atomic, and newer-wins. Select-then-insert loses races that really happen + 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 + step together would both see no row and both insert, and the loser's + IntegrityError is swallowed by the caller, so the run silently vanishes. + 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) - status = str(runtime.get("status") or "completed") message = runtime.get("message") + now = datetime.now(timezone.utc) values = { - "status": status, + "job_name": job_name, + "status": str(runtime.get("status") or "completed"), "started_at": _as_datetime(runtime.get("started_at")), "finished_at": finished_at, "processed": runtime.get("processed"), "total": runtime.get("total"), "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: - row = JobRunState(job_name=job_name, **values) - db.add(row) - else: - for field, value in values.items(): - setattr(row, field, value) - return row + statement = _insert_for(db)(JobRunState).values(**values) + await db.execute( + statement.on_conflict_do_update( + index_elements=[JobRunState.job_name], + set_={key: statement.excluded[key] for key in values if key != "job_name"}, + where=JobRunState.finished_at < statement.excluded.finished_at, + ) + ) diff --git a/tests/unit/test_job_run_persistence.py b/tests/unit/test_job_run_persistence.py index 4af6f2b..ae1f99c 100644 --- a/tests/unit/test_job_run_persistence.py +++ b/tests/unit/test_job_run_persistence.py @@ -164,3 +164,108 @@ class TestListJobsSplitsLiveFromPersisted: assert job["runtime_status"] == "idle" assert job["running"] is False 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()