From e9027f2b30f37bb3052714eb08fcf479542f4fc0 Mon Sep 17 00:00:00 2001 From: Simon Pinfold Date: Sat, 3 Oct 2026 01:20:37 -0700 Subject: [PATCH] fix(assets): write the prune and offline marking in short batches so saves aren't locked out (#16696) * fix(assets): batch the prune's and the offline marking's writes The startup prune, POST /api/assets/prune and the fast scan's marking step each held the SQLite write lock for their whole loop, so foreground output registration failed with "database is locked" during a large one. They now write in short batches, wait while a prompt runs between batches, and the prune endpoint runs off the event loop. * fix(assets): start the queued scan after a standalone prune, and recheck listing rows after a pause A prompt that ends while POST /api/assets/prune runs queues its output rescan; the prune now starts it when it finishes, as a scan does. The output-listing rescan takes its batch gate before reading the live rows, so a pause during the walk makes the marking re-stat what it retires. A cancel that arrives after the last batch no longer reports a finished prune as cancelled. * refactor(assets): drop the pause rechecks and the cancellable standalone prune Batching the writes is what keeps the lock short; the layers on top of it guarded edge cases that heal on the next scan. Batches now just commit, sleep about as long as they held the lock, and between batches honour the scan's pause/cancel checkpoint. The standalone prune is batched but not pausable, so it needs no cancel status or pending-scan handling, and the API contract is unchanged apart from running off the event loop. * fix(assets): start the scan queued behind a standalone prune; skip the last batch's yield POST /api/assets/prune now runs off the event loop, so a prompt can finish while it runs and queue its output rescan; the prune starts it when it ends, as a scan does. The batch loop checks for a stop before every batch and no longer sleeps after the last one. * test(assets): compare the set-mark paths in their stored, absolute form create_content stores os.path.abspath(path), which carries a drive letter on Windows, so the expected list must be built the same way. * fix(assets): a seed request during an API prune waits for it instead of 409 The prune now runs off the event loop, so POST /api/assets/seed can arrive while it holds the seeder; start() fails and the route answered 409, which a client reads as "a scan is already coming". A prune emits no scan events, so the refresh was lost. The route now waits the prune out and starts the scan, as it effectively did when the prune blocked the loop. * fix(assets): a cancel or shutdown stops a standalone prune between batches The API prune runs on a worker thread that interpreter exit joins, so a shutdown that only flagged it left Ctrl-C waiting for the whole prune. It now stops at the next batch once cancelled, and shutdown waits for that. A seed request also retries start() once after any failure, covering a prune that ends between the failed start and the check. * fix(assets): report a cancelled API prune as cancelled, not completed A cancel now stops a standalone prune between batches, so its response can carry a partial count; say so with status "cancelled" rather than presenting it as a finished prune. * fix(assets): a cancelled standalone prune leaves a queued scan queued Shutdown cancels the prune; starting the scan a prompt had queued from the prune's finalizer would run it on into teardown after shutdown returned. It now stays queued for the next scan's finalizer. * test(assets): assert the cancelled prune's outcome in the test thread pytest.raises inside the worker thread only produced a warning when the exception was missing, so the test could not fail on it. * fix(assets): wait for a prune on the loop, and close shutdown gaps around it A seed request during an API prune now polls on the event loop instead of holding an executor thread for the prune's length, and retries while a prune holds the seeder. Shutdown marks the seeder so a prune that has not started yet does not, both of its waits share one deadline, and the prune's idle flag is set even if its cleanup raises. --- app/assets/api/routes.py | 23 +- app/assets/database/queries/__init__.py | 2 + app/assets/database/queries/records.py | 27 + app/assets/scanner.py | 138 ++-- app/assets/seeder.py | 120 ++- openapi.yaml | 1 + .../services/test_batched_marking.py | 722 ++++++++++++++++++ .../assets_test/services/test_scanner_b.py | 12 +- tests-unit/seeder_test/test_seeder.py | 20 +- 9 files changed, 979 insertions(+), 86 deletions(-) create mode 100644 tests-unit/assets_test/services/test_batched_marking.py diff --git a/app/assets/api/routes.py b/app/assets/api/routes.py index b6fbfe131..c7fc7022b 100644 --- a/app/assets/api/routes.py +++ b/app/assets/api/routes.py @@ -43,7 +43,7 @@ from app.assets.database.queries.records import ( get_preview_file_paths_by_ids, list_records_page, ) -from app.assets.seeder import ScanInProgressError, asset_seeder +from app.assets.seeder import PruneCancelledError, ScanInProgressError, asset_seeder from app.assets.services import ( DependencyMissingError, HashMismatchError, @@ -1030,6 +1030,9 @@ async def get_tags_refine(request: web.Request) -> web.Response: return web.json_response(payload.model_dump(mode="json", exclude_none=True), status=200) +_PRUNE_POLL_SECONDS = 0.25 + + @ROUTES.post("/api/assets/seed") @_require_assets_feature_enabled async def seed_assets(request: web.Request) -> web.Response: @@ -1059,6 +1062,18 @@ async def seed_assets(request: web.Request) -> web.Response: started = asset_seeder.start( roots=valid_roots, compute_hashes=mode.hashing_enabled() ) + # A prune from POST /api/assets/prune is not a scan and emits no scan events, so + # wait it out rather than answer 409, which a client takes as "a scan is coming". + # Polled on the loop: no executor thread is held for the prune's length. The + # retry also covers a prune that ended, or another that began, after start() failed. + for _ in range(2): + if started: + break + while asset_seeder.standalone_prune_running(): + await asyncio.sleep(_PRUNE_POLL_SECONDS) + started = asset_seeder.start( + roots=valid_roots, compute_hashes=mode.hashing_enabled() + ) if not started: return web.json_response({"status": "already_running"}, status=409) @@ -1125,17 +1140,21 @@ async def mark_missing_assets(request: web.Request) -> web.Response: Returns: 200 OK with count of marked assets + 200 OK with status "cancelled" and the count marked before a cancel stopped it 409 Conflict if a scan is currently running 500 Internal Server Error with PRUNE_FAILED if the marking failed, so a prune that did not run is never reported as a completed one """ try: - marked = asset_seeder.mark_missing_outside_prefixes() + # Off the event loop: a large prune takes seconds to minutes. + marked = await asyncio.to_thread(asset_seeder.mark_missing_outside_prefixes) except ScanInProgressError: return web.json_response( {"status": "scan_running", "marked": 0}, status=409, ) + except PruneCancelledError as cancelled: + return web.json_response({"status": "cancelled", "marked": cancelled.marked}, status=200) if marked is None: return _build_error_response( 500, diff --git a/app/assets/database/queries/__init__.py b/app/assets/database/queries/__init__.py index 14dff90a8..7293e6f40 100644 --- a/app/assets/database/queries/__init__.py +++ b/app/assets/database/queries/__init__.py @@ -16,6 +16,7 @@ from app.assets.database.queries.records import ( is_live_path_conflict, list_records_page, mark_content_missing, + mark_contents_missing, rename_record, unset_content_missing, update_record_access_time, @@ -31,6 +32,7 @@ __all__ = [ "is_live_path_conflict", "list_records_page", "mark_content_missing", + "mark_contents_missing", "rename_record", "unset_content_missing", "update_record_access_time", diff --git a/app/assets/database/queries/records.py b/app/assets/database/queries/records.py index bb12e2315..19dec16cf 100644 --- a/app/assets/database/queries/records.py +++ b/app/assets/database/queries/records.py @@ -325,6 +325,33 @@ def mark_content_missing(session: Session, content_id: str) -> None: session.flush() +def mark_contents_missing(session: Session, content_ids: Sequence[str]) -> list[str]: + """mark_content_missing for many rows in a few statements; returns the ids it marked. + A row that is gone or already missing is skipped.""" + if not content_ids: + return [] + # "= 0", not "IS 0": SQLite only uses the partial live-path index for "= 0". This + # lookup is by primary key either way; the form matches the other live-row lookups. + live = list(session.scalars(sa.select(AssetContent.id).where(AssetContent.id.in_(content_ids), AssetContent.is_missing == sa.false()))) + if not live: + return [] + session.execute(sa.update(AssetContent).where(AssetContent.id.in_(live)).values(is_missing=True)) + ensure_tag(session, "missing") + unlinked = sa.select(Asset.id, sa.literal("missing"), sa.literal("automatic"), sa.literal(get_utc_now(), sa.DateTime())).where( + Asset.content_id.in_(live), + ~sa.exists().where(AssetTag.asset_id == Asset.id, AssetTag.tag_name == "missing"), + ) + try: + with session.begin_nested(): + session.execute(sa.insert(AssetTag).from_select(["asset_id", "tag_name", "origin", "added_at"], unlinked)) + except IntegrityError: + # A concurrent writer linked one of them first; settle each link the race-safe way. + for record_id in session.scalars(sa.select(Asset.id).where(Asset.content_id.in_(live))): + ensure_tag_link(session, asset_id=record_id, tag_name="missing", origin="automatic") + session.flush() + return live + + def unset_content_missing(session: Session, content_id: str) -> None: content = session.get(AssetContent, content_id) if content is None: diff --git a/app/assets/scanner.py b/app/assets/scanner.py index 9ed26336c..193f515c4 100644 --- a/app/assets/scanner.py +++ b/app/assets/scanner.py @@ -11,6 +11,7 @@ so the rows it never reached are selected again when the scan resumes. import enum import logging import os +import time from dataclasses import dataclass from itertools import islice from pathlib import Path @@ -25,7 +26,7 @@ from app.assets.event_log import emit, error_kind, error_type from app.assets.database.queries import ( create_content_reporting_insert, is_live_path_conflict, - mark_content_missing, + mark_contents_missing, create_record, ) from app.assets.database.models import Asset, AssetContent @@ -134,6 +135,46 @@ class _ReferenceObservation(NamedTuple): stat_result: os.stat_result | None +# Before each write batch: blocks while the scan is paused, returns True to stop. +ShouldStop = Callable[[], bool] + + +def _never_stop() -> bool: + return False + + +# The marking steps write in batches of WRITE_BATCH_ROWS, each its own transaction, so +# the write lock is never held for a whole pass and a foreground write (an output being +# registered) waits at most one batch. After each commit the thread sleeps about as +# long as it held the lock: a writer in SQLite's busy handler polls with growing sleeps +# (up to 100 ms), so taking the lock straight back would keep winning it. +WRITE_BATCH_ROWS = 256 +WRITE_YIELD_MIN_SECONDS = 0.02 +WRITE_YIELD_MAX_SECONDS = 0.1 + + +def _write_in_batches( + items: list, + write_batch: Callable[[Session, list], int], + should_stop: ShouldStop, + committed: list[int], +) -> None: + """Apply ``write_batch`` to ``items`` a batch per transaction, appending the count + each batch reports to ``committed`` once it commits. A failure leaves the batches + before it committed, and ``committed`` says how much they wrote.""" + for start in range(0, len(items), WRITE_BATCH_ROWS): + if should_stop(): + break + opened = time.perf_counter() + with create_write_session() as session: + count = write_batch(session, items[start : start + WRITE_BATCH_ROWS]) + session.commit() + committed.append(count) + held = time.perf_counter() - opened + if start + WRITE_BATCH_ROWS < len(items): + time.sleep(min(max(held, WRITE_YIELD_MIN_SECONDS), WRITE_YIELD_MAX_SECONDS)) + + def _log_scan_error(phase: str, error: OSError) -> None: error_type = ( "permission_denied" if isinstance(error, PermissionError) else "os_error" @@ -241,7 +282,9 @@ def apply_reference_observations( session: Session, observations: list[_ReferenceObservation] ) -> int: """Apply the observations; returns how many rows were marked missing.""" - marked = 0 + # One query loads the rows, so the session.get calls below never go to the database. + list(session.scalars(sa.select(AssetContent).where(AssetContent.id.in_([o.content_id for o in observations])))) + gone: list[str] = [] for observation in observations: content = session.get(AssetContent, observation.content_id) # Skip a row another writer changed since it was observed; the next scan sees it afresh. @@ -253,8 +296,7 @@ def apply_reference_observations( ): continue if observation.stat_result is None: - mark_content_missing(session, content.id) - marked += 1 + gone.append(content.id) continue detect_content_change( session, @@ -262,35 +304,38 @@ def apply_reference_observations( observation.stat_result, hashing_is_enabled=mode.hashing_enabled(), ) - return marked + return len(mark_contents_missing(session, gone)) -def _sync_prefixes_in_write_txn( - prefixes: list[str], progress: _ScanProgress | None -) -> tuple[set[str], int]: - """Returns the surviving paths and how many rows were marked missing.""" +def _sync_prefixes( + prefixes: list[str], + progress: _ScanProgress | None, + should_stop: ShouldStop, + marked: list[int], +) -> set[str]: + """Returns the surviving paths; ``marked`` gets each committed batch's count of rows + marked missing, which a caller still has if a later batch raises.""" with create_session() as session: observations, survivors = observe_references_on_filesystem( session, prefixes, progress ) - marked = 0 - if observations: - with create_write_session() as session: - marked = apply_reference_observations(session, observations) - session.commit() - return survivors, marked + _write_in_batches(observations, apply_reference_observations, should_stop, marked) + return survivors def sync_root_safely( - root: RootType, progress: _ScanProgress | None = None + root: RootType, + progress: _ScanProgress | None = None, + should_stop: ShouldStop = _never_stop, ) -> set[str]: """Sync a single root's references with the filesystem. Returns survivors (existing paths) or empty set on failure. """ + marked: list[int] = [] try: - survivors, marked = _sync_prefixes_in_write_txn( - get_scan_prefixes_for_root(root), progress + survivors = _sync_prefixes( + get_scan_prefixes_for_root(root), progress, should_stop, marked ) except Exception as exc: logging.exception("fast DB scan failed for %s: %s", root, exc) @@ -300,18 +345,19 @@ def sync_root_safely( error_type=error_type(exc), error_kind=error_kind(exc), ) - return set() + survivors = set() if progress is not None: - progress.missing_marked += marked + progress.missing_marked += sum(marked) return survivors def sync_temp_references_safely( progress: _ScanProgress | None = None, + should_stop: ShouldStop = _never_stop, ) -> None: """Retire temp references whose file is gone; temp is never scanned, so nothing else stats them.""" try: - _sync_prefixes_in_write_txn(get_temp_prefixes(), progress) + _sync_prefixes(get_temp_prefixes(), progress, should_stop, []) except Exception as exc: logging.exception("temp reference sync failed: %s", exc) emit( @@ -322,19 +368,27 @@ def sync_temp_references_safely( ) -def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int | None: +def mark_missing_outside_prefixes_safely( + prefixes: list[str], should_stop: ShouldStop = _never_stop +) -> int | None: """Mark references as missing when outside the given prefixes. This is a non-destructive soft-delete. Returns the count marked, or None when - the operation fails. + the operation fails; batches committed before a failure stay committed. """ + marked: list[int] = [] try: with create_session() as sess: - count = mark_contents_missing_outside_prefixes(sess, prefixes) - sess.commit() - return count + content_ids = content_ids_outside_prefixes(sess, prefixes) + _write_in_batches( + content_ids, + lambda session, batch: len(mark_contents_missing(session, batch)), + should_stop, + marked, + ) + return sum(marked) except Exception as exc: - logging.exception("marking missing assets failed: %s", exc) + logging.exception("marking missing assets failed after marking %d: %s", sum(marked), exc) emit( "scanner.mark_missing_failed", error_type=error_type(exc), @@ -343,19 +397,16 @@ def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int | None: return None -def mark_contents_missing_outside_prefixes( - session: Session, prefixes: list[str] -) -> int: - contents = session.scalars( - sa.select(AssetContent) - .where(AssetContent.is_missing.is_(False)) +def content_ids_outside_prefixes(session: Session, prefixes: list[str]) -> list[str]: + """The live rows outside every prefix. A read: the marking re-checks each is still + live inside its own write transaction.""" + is_owned = path_prefix_matcher(prefixes) + rows = session.execute( + sa.select(AssetContent.id, AssetContent.path) + .where(AssetContent.is_missing == sa.false()) .execution_options(yield_per=500) ) - is_owned = path_prefix_matcher(prefixes) - missing = [content for content in contents if not is_owned(content.path)] - for content in missing: - mark_content_missing(session, content.id) - return len(missing) + return [content_id for content_id, path in rows if not is_owned(path)] def collect_paths_for_roots( @@ -509,15 +560,13 @@ def mark_unlisted_references_missing_safely( root: RootType, observations: list[_ReferenceObservation], progress: _ScanProgress | None = None, + should_stop: ShouldStop = _never_stop, ) -> None: """Retire rows whose file the listing lacks, through the same guarded write sync_root applies to a row whose file has vanished.""" - if not observations: - return + marked: list[int] = [] try: - with create_write_session() as session: - marked = apply_reference_observations(session, observations) - session.commit() + _write_in_batches(observations, apply_reference_observations, should_stop, marked) except Exception as exc: logging.exception("fast DB scan failed for %s: %s", root, exc) emit( @@ -526,9 +575,8 @@ def mark_unlisted_references_missing_safely( error_type=error_type(exc), error_kind=error_kind(exc), ) - return if progress is not None: - progress.missing_marked += marked + progress.missing_marked += sum(marked) def list_output_for_rescan() -> ListingWalk: diff --git a/app/assets/seeder.py b/app/assets/seeder.py index 3476b90d2..29b444f7d 100644 --- a/app/assets/seeder.py +++ b/app/assets/seeder.py @@ -43,6 +43,15 @@ class ScanInProgressError(Exception): """Raised when an operation cannot proceed because a scan is running.""" +class PruneCancelledError(Exception): + """A standalone prune stopped by a cancel. The batches before it stay committed; + ``marked`` counts them.""" + + def __init__(self, marked: int) -> None: + super().__init__(f"prune cancelled after marking {marked}") + self.marked = marked + + class State(Enum): """Seeder state machine states.""" @@ -172,6 +181,11 @@ class _AssetSeeder: self._cancel_event = threading.Event() self._run_gate = threading.Event() self._run_gate.set() # Start unpaused (set = running, clear = paused) + # Clear while a standalone prune holds the seeder (set = no prune running). + self._prune_idle = threading.Event() + self._prune_idle.set() + # Set by shutdown(): a standalone prune that has not started by then does not. + self._shutting_down = False self._roots: tuple[RootType, ...] = () self._phase: ScanPhase = ScanPhase.FULL self._compute_hashes: bool = False @@ -447,8 +461,14 @@ class _AssetSeeder: Returns: True if the scan thread joined cleanly; False on timeout. """ + with self._lock: + self._shutting_down = True self.cancel() + # A standalone prune stops at its next batch once cancelled. One deadline + # covers both waits. + deadline = time.monotonic() + timeout joined = self.wait(timeout=timeout) + joined = self.wait_for_standalone_prune(max(0.0, deadline - time.monotonic())) and joined if not joined: logging.warning( "Asset seeder thread did not exit within %ss", @@ -480,13 +500,18 @@ class _AssetSeeder: Raises: ScanInProgressError: If a scan is currently running + PruneCancelledError: If a cancel stopped it part way """ with self._lock: if self._state != State.IDLE: raise ScanInProgressError( "Cannot mark missing assets while scan is running" ) + if self._shutting_down: + raise PruneCancelledError(0) self._state = State.RUNNING + self._cancel_event.clear() + self._prune_idle.clear() try: if not dependencies_available(): @@ -496,7 +521,16 @@ class _AssetSeeder: return 0 all_prefixes = get_owned_prefixes() - marked = mark_missing_outside_prefixes_safely(all_prefixes) + # Not pausable (the API waits on it), but a cancel or shutdown stops it + # between batches: it runs on a worker thread that exit would wait for. + stopped = False + + def should_stop() -> bool: + nonlocal stopped + stopped = self._cancel_event.is_set() + return stopped + + marked = mark_missing_outside_prefixes_safely(all_prefixes, should_stop) if marked is None: return None emit( @@ -504,12 +538,32 @@ class _AssetSeeder: count=marked, stage=_ScanStage.MARK_MISSING.value, ) + if stopped: + logging.info("Marking missing assets cancelled after marking %d", marked) + raise PruneCancelledError(marked) if marked > 0: logging.info("Marked %d references as missing", marked) return marked finally: + # The API runs this off the event loop, so a prompt can finish meanwhile + # and queue its output rescan; start it now. Not after a cancel: shutdown + # cancels, and a scan started here would run on into teardown. It stays + # queued for the next scan to start. with self._lock: - self._reset_to_idle() + try: + if self._cancel_event.is_set(): + self._reset_to_idle() + else: + self._finish_and_start_pending() + finally: + self._prune_idle.set() + + def standalone_prune_running(self) -> bool: + return not self._prune_idle.is_set() + + def wait_for_standalone_prune(self, timeout: float | None = None) -> bool: + """Block until no standalone prune holds the seeder. True unless it timed out.""" + return self._prune_idle.wait(timeout) def _reset_to_idle(self) -> None: """Reset state to IDLE, preserving last progress. Caller must hold _lock.""" @@ -655,11 +709,13 @@ class _AssetSeeder: if self._prune_first: all_prefixes = get_owned_prefixes() - marked = mark_missing_outside_prefixes_safely(all_prefixes) + marked = mark_missing_outside_prefixes_safely( + all_prefixes, lambda: self._check_pause_and_cancel(_ScanStage.PRUNING) + ) marked_count = 0 if marked is None else marked if marked is None: self._add_error( - "Marking missing assets failed; scan continued without pruning" + "Marking missing assets failed; scan continued with the prune incomplete" ) else: emit( @@ -671,7 +727,9 @@ class _AssetSeeder: logging.info( "Marked %d refs as missing before scan", marked_count ) - sync_temp_references_safely(scan_state) + sync_temp_references_safely( + scan_state, lambda: self._check_pause_and_cancel(_ScanStage.PRUNING) + ) if self._check_pause_and_cancel(_ScanStage.PRUNING): logging.info("Asset scan cancelled after pruning phase") @@ -801,23 +859,28 @@ class _AssetSeeder: ) finally: with self._lock: - start_paused = self._state is State.PAUSED - self._reset_to_idle() - pending = self._pending_scan - if pending is not None: - self._pending_scan = None - if not self.start( - roots=pending["roots"], - phase=pending["phase"], - prune_first=False, - compute_hashes=pending["compute_hashes"], - _start_paused=start_paused, - ): - logging.warning( - "Pending scan could not start (roots=%s, phase=%s)", - pending["roots"], - pending["phase"].value, - ) + self._finish_and_start_pending() + + def _finish_and_start_pending(self) -> None: + """Reset to IDLE, then start the scan queued while this run held the seeder, + paused if this run was. Caller must hold _lock.""" + start_paused = self._state is State.PAUSED + self._reset_to_idle() + pending = self._pending_scan + if pending is not None: + self._pending_scan = None + if not self.start( + roots=pending["roots"], + phase=pending["phase"], + prune_first=False, + compute_hashes=pending["compute_hashes"], + _start_paused=start_paused, + ): + logging.warning( + "Pending scan could not start (roots=%s, phase=%s)", + pending["roots"], + pending["phase"].value, + ) @staticmethod def _emit_marked_missing(root: RootType, marked: int) -> None: @@ -854,7 +917,11 @@ class _AssetSeeder: existing_paths.update(live_references) else: marked_before = scan_state.missing_marked - existing_paths.update(sync_root_safely(r, scan_state)) + existing_paths.update( + sync_root_safely( + r, scan_state, lambda: self._check_pause_and_cancel(_ScanStage.FAST_SCAN) + ) + ) self._emit_marked_missing(r, scan_state.missing_marked - marked_before) logging.debug( "Fast scan: sync_root phase took %.3fs (%d existing paths)", @@ -877,7 +944,12 @@ class _AssetSeeder: scan_state.dirs_listed += walk.dirs_listed vanished, unlisted = unlisted_references(live_references, walk.listings, scan_state) marked_before = scan_state.missing_marked - mark_unlisted_references_missing_safely("output", vanished, scan_state) + mark_unlisted_references_missing_safely( + "output", + vanished, + scan_state, + lambda: self._check_pause_and_cancel(_ScanStage.FAST_SCAN), + ) self._emit_marked_missing("output", scan_state.missing_marked - marked_before) logging.debug( "Fast scan: output listing: %d dirs listed, %d rows retired, " diff --git a/openapi.yaml b/openapi.yaml index d2288354e..d1146ffe7 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -2402,6 +2402,7 @@ paths: description: Number of assets marked as missing type: integer status: + description: '"completed", or "cancelled" when a cancel stopped the prune part way; the assets it marked before that stay marked' type: string type: object description: Prune result diff --git a/tests-unit/assets_test/services/test_batched_marking.py b/tests-unit/assets_test/services/test_batched_marking.py new file mode 100644 index 000000000..7f77c4168 --- /dev/null +++ b/tests-unit/assets_test/services/test_batched_marking.py @@ -0,0 +1,722 @@ +"""The prune and the offline marking write in short batches rather than one long +transaction, so a foreground write gets the lock between batches. These tests pin the +batch boundaries, what a failure or a stop leaves behind, the pause between batches, +what a batch does with rows another writer changed, and the set-based mark.""" + +import asyncio +import json +import os +import sqlite3 +import threading +import time +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import patch + +import pytest +import sqlalchemy as sa +from aiohttp.test_utils import make_mocked_request +from sqlalchemy import event +from sqlalchemy.orm import Session as SASession, sessionmaker +from sqlalchemy.pool import StaticPool + +from app.assets import scanner, seeder as seeder_module +from app.assets.api import routes +from app.assets.database.models import Asset, AssetContent, AssetTag, Base, Tag +from app.assets.database.queries.records import ( + create_content, + create_record, + ensure_tag, + ensure_tag_link, + mark_content_missing, + mark_contents_missing, +) +from app.assets.scanner_admission import _WATCH_LIST +from app.assets.seeder import State + + +@pytest.fixture +def db_engine(): + """One in-memory database every thread shares, for the tests that prune on a worker.""" + engine = sa.create_engine("sqlite://", poolclass=StaticPool, connect_args={"check_same_thread": False}) + Base.metadata.create_all(engine) + return engine + + +@pytest.fixture(autouse=True) +def no_yield(monkeypatch): + monkeypatch.setattr(scanner, "WRITE_YIELD_MIN_SECONDS", 0.0) + monkeypatch.setattr(scanner, "WRITE_YIELD_MAX_SECONDS", 0.0) + + +@pytest.fixture +def catalog(db_engine): + """Routes the scanner's sessions to the test engine and counts write transactions.""" + opened: list[int] = [] + write_factory = sessionmaker(bind=db_engine) + + @contextmanager + def _create_session(): + with SASession(db_engine) as sess: + yield sess + + def _create_write_session(): + opened.append(1) + return write_factory() + + _WATCH_LIST.clear() + with patch("app.assets.scanner.create_session", _create_session), \ + patch("app.assets.seeder.create_session", _create_session), \ + patch("app.assets.scanner.create_write_session", _create_write_session): + yield opened + _WATCH_LIST.clear() + + +def _rows(session, directory: Path, count: int, *, files: bool = False) -> list[str]: + ids = [] + for i in range(count): + path = directory / f"f{i:05d}.png" + if files: + path.write_bytes(f"bytes-{i}".encode()) + stat = path.stat() if files else None + content = create_content( + session, + path=str(path), + size_bytes=stat.st_size if stat else 1, + mtime_ns=stat.st_mtime_ns if stat else 1, + ) + create_record(session, content_id=content.id, name=path.name, tags=["output"]) + ids.append(content.id) + session.commit() + return ids + + +def _live_ids(session) -> set[str]: + session.expire_all() + return set(session.scalars(sa.select(AssetContent.id).where(AssetContent.is_missing == sa.false()))) + + +def _prune(owned: list[str]) -> int | None: + return scanner.mark_missing_outside_prefixes_safely(owned) + + +@pytest.mark.parametrize("count, batches", [(0, 0), (1, 1), (256, 1), (257, 2), (600, 3)]) +def test_prune_commits_one_batch_per_256_rows(session, catalog, temp_dir, count, batches): + _rows(session, temp_dir, count) + + assert _prune([]) == count + assert len(catalog) == batches + assert _live_ids(session) == set() + + +def test_rows_under_an_owned_prefix_are_not_pruned(session, catalog, temp_dir): + owned, gone = temp_dir / "owned", temp_dir / "gone" + owned.mkdir() + gone.mkdir() + kept = set(_rows(session, owned, 5)) + _rows(session, gone, 5) + + assert _prune([str(owned)]) == 5 + assert _live_ids(session) == kept + + +def test_a_row_re_registered_between_batches_is_left_to_its_new_owner(session, catalog, temp_dir): + """register_executed_output retires the live row at a path and inserts a new one. + Doing that between two batches must leave exactly the new row live.""" + ids = _rows(session, temp_dir, 300) + last = session.get(AssetContent, ids[-1]) + replacement: list[str] = [] + + def between_batches() -> bool: + if len(catalog) == 1: + mark_content_missing(session, last.id) + new = create_content(session, path=last.path, size_bytes=2, mtime_ns=2) + session.commit() + replacement.append(new.id) + return False + + marked = scanner.mark_missing_outside_prefixes_safely([], between_batches) + + # The replacement was never a candidate; the retired row is not counted twice. + assert marked == len(ids) - 1 + assert _live_ids(session) == set(replacement) + + +def test_a_foreground_write_gets_the_lock_between_batches(tmp_path, monkeypatch): + """On a real file database with the production write-session setup, another + connection can take the write lock at every point between two batches.""" + db = tmp_path / "catalog.db" + read_engine = sa.create_engine(f"sqlite:///{db}") + write_engine = sa.create_engine(f"sqlite:///{db}") + + @event.listens_for(write_engine, "connect") + def _connect(dbapi_connection, _record): + dbapi_connection.isolation_level = None + + @event.listens_for(write_engine, "begin") + def _begin(connection): + connection.exec_driver_sql("BEGIN IMMEDIATE") + + with read_engine.connect() as conn: + conn.exec_driver_sql("PRAGMA journal_mode=WAL") + Base.metadata.create_all(read_engine) + with SASession(read_engine) as session: + _rows(session, tmp_path, 600) + + foreground: list[bool] = [] + + def between_batches() -> bool: + other = sqlite3.connect(db, timeout=0, isolation_level=None) + try: + other.execute("BEGIN IMMEDIATE") + other.execute("INSERT INTO tags (name) VALUES (?)", (f"fg-{len(foreground)}",)) + other.execute("COMMIT") + foreground.append(True) + except sqlite3.OperationalError: + foreground.append(False) + finally: + other.close() + return False + + monkeypatch.setattr(scanner, "create_session", lambda: SASession(read_engine)) + monkeypatch.setattr(scanner, "create_write_session", sessionmaker(bind=write_engine)) + assert scanner.mark_missing_outside_prefixes_safely([], between_batches) == 600 + + assert foreground == [True, True, True] + + +def test_a_failed_batch_keeps_the_batches_before_it(session, catalog, temp_dir): + _rows(session, temp_dir, 600) + real = scanner.mark_contents_missing + + def fail_in_the_second_batch(sess, ids): + if len(catalog) == 2: + raise RuntimeError("disk I/O error") + return real(sess, ids) + + with patch("app.assets.scanner.mark_contents_missing", fail_in_the_second_batch): + assert _prune([]) is None + + assert len(_live_ids(session)) == 600 - scanner.WRITE_BATCH_ROWS + + +def test_sync_root_counts_the_batches_committed_before_a_failure(session, catalog, temp_dir, monkeypatch): + output = temp_dir / "output" + output.mkdir() + monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(output)) + _rows(session, output, 300) + real = scanner.mark_contents_missing + + def fail_in_the_second_batch(sess, ids): + if len(catalog) == 2: + raise RuntimeError("disk I/O error") + return real(sess, ids) + + progress = seeder_module._ScanState() + with patch("app.assets.scanner.mark_contents_missing", fail_in_the_second_batch): + assert scanner.sync_root_safely("output", progress) == set() + + assert progress.missing_marked == scanner.WRITE_BATCH_ROWS + + +def _gone_observations(session, temp_dir: Path, count: int) -> list[scanner._ReferenceObservation]: + ids = _rows(session, temp_dir, count, files=True) + observations = [] + for content_id in ids: + content = session.get(AssetContent, content_id) + os.remove(content.path) + observations.append( + scanner._ReferenceObservation(content.id, content.size_bytes, content.mtime_ns, None) + ) + return observations + + +def test_stop_between_batches_leaves_the_rest_live(session, catalog, temp_dir): + observations = _gone_observations(session, temp_dir, 600) + + committed: list[int] = [] + scanner._write_in_batches(observations, scanner.apply_reference_observations, lambda: bool(committed), committed) + + assert committed == [scanner.WRITE_BATCH_ROWS] + assert len(_live_ids(session)) == 600 - scanner.WRITE_BATCH_ROWS + + +def test_a_scan_waits_between_batches_while_a_prompt_runs(session, catalog, temp_dir, monkeypatch): + output = temp_dir / "output" + output.mkdir() + monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(output)) + _rows(session, output, 600) + instance = seeder_module._AssetSeeder() + instance._state = State.RUNNING + instance._scan_state = seeder_module._ScanState() + instance._run_gate.set() + real = scanner.mark_contents_missing + first_batch = threading.Event() + + def mark(sess, ids): + if not first_batch.is_set(): + assert instance.pause() # a prompt starts during the first batch + first_batch.set() + return real(sess, ids) + + def sync(): + scanner.sync_root_safely( + "output", + instance._scan_state, + lambda: instance._check_pause_and_cancel(seeder_module._ScanStage.FAST_SCAN), + ) + + with patch("app.assets.scanner.mark_contents_missing", mark): + worker = threading.Thread(target=sync) + worker.start() + assert first_batch.wait(5) + time.sleep(0.2) + assert len(catalog) == 1 # no batch while paused + assert instance.resume() + worker.join(5) + + assert len(catalog) == 3 + assert instance._scan_state.missing_marked == 600 + + +def test_offline_rows_retired_across_batches_recover_when_the_drive_returns( + session, catalog, temp_dir, monkeypatch +): + """#16646's hashing-off recovery with the marking split over several batches: files + that come back while the marking is part way through are retired by the batches + that follow, then recovered by the walk that follows in the same scan, every record + keeping its id.""" + output = temp_dir / "output" + (temp_dir / "input").mkdir() + output.mkdir() + monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(output)) + monkeypatch.setattr("folder_paths.get_input_directory", lambda: str(temp_dir / "input")) + monkeypatch.setattr(scanner, "get_comfy_models_folders", lambda: []) + ids = _rows(session, output, 700, files=True) + records = {record.id: record.content_id for record in session.scalars(sa.select(Asset))} + parked = temp_dir / "parked" + output.rename(parked) + output.mkdir() + + instance = seeder_module._AssetSeeder() + instance._scan_state = seeder_module._ScanState() + instance._phase = seeder_module.ScanPhase.FAST + instance._run_gate.set() + real = scanner.mark_contents_missing + + def mark(sess, content_ids): + marked = real(sess, content_ids) + if len(catalog) == 1: + # The drive comes back after the first batch. + for name in os.listdir(parked): + os.rename(parked / name, output / name) + return marked + + with patch("app.assets.scanner.mark_contents_missing", mark): + instance._run_fast_phase(("input", "output")) + + session.expire_all() + assert instance._scan_state.missing_marked == 700 + assert instance._scan_state.recovered == 700 + assert _live_ids(session) == set(ids) + assert {record.id: record.content_id for record in session.scalars(sa.select(Asset))} == records + live_paths = session.scalars(sa.select(AssetContent.path).where(AssetContent.is_missing == sa.false())).all() + assert len(live_paths) == len(set(live_paths)) == 700 + + +def _link_state(session) -> tuple[dict[str, bool], set[tuple[str, str, str]]]: + session.expire_all() + contents = {c.path: c.is_missing for c in session.scalars(sa.select(AssetContent))} + links = { + (record.name, link.tag_name, link.origin) + for record in session.scalars(sa.select(Asset)) + for link in session.scalars(sa.select(AssetTag).where(AssetTag.asset_id == record.id)) + } + return contents, links + + +def _equivalence_fixture(session) -> list[str]: + """Rows covering each case the mark handles: no record, one, two, a record already + tagged missing by hand, an already-missing row, and an id that does not exist.""" + none = create_content(session, path="/c/none.png") + one = create_content(session, path="/c/one.png") + create_record(session, content_id=one.id, name="one", tags=["output"]) + two = create_content(session, path="/c/two.png") + create_record(session, content_id=two.id, name="two-a") + create_record(session, content_id=two.id, name="two-b", tags=["input"]) + tagged = create_content(session, path="/c/tagged.png") + record = create_record(session, content_id=tagged.id, name="tagged") + ensure_tag(session, "missing") + ensure_tag_link(session, asset_id=record.id, tag_name="missing", origin="manual") + already = create_content(session, path="/c/already.png") + create_record(session, content_id=already.id, name="already") + mark_content_missing(session, already.id) + session.commit() + return [none.id, one.id, two.id, tagged.id, already.id, "no-such-id"] + + +def test_the_set_mark_matches_marking_row_by_row(session): + engine = sa.create_engine("sqlite:///:memory:") + Base.metadata.create_all(engine) + with SASession(engine) as per_row: + # Same fixture in both catalogs, then the ids made to match by path. + ids = _equivalence_fixture(session) + _equivalence_fixture(per_row) + by_path = {c.path: c.id for c in per_row.scalars(sa.select(AssetContent))} + paths = {c.id: c.path for c in session.scalars(sa.select(AssetContent))} + + marked = mark_contents_missing(session, ids) + session.commit() + for content_id in ids: + path = paths.get(content_id) + content = per_row.get(AssetContent, by_path[path]) if path else None + if content is not None and not content.is_missing: + mark_content_missing(per_row, content.id) + per_row.commit() + + # create_content stores os.path.abspath(path), so compare in that form (a drive letter on Windows). + expected = sorted(os.path.abspath(f"/c/{name}.png") for name in ("none", "one", "tagged", "two")) + assert sorted(paths[i] for i in marked) == expected + assert _link_state(session) == _link_state(per_row) + assert session.get(Tag, "missing") is not None + + +def test_the_set_mark_settles_a_link_race_row_by_row(session, monkeypatch): + content = create_content(session, path="/c/raced.png") + record = create_record(session, content_id=content.id, name="raced") + session.commit() + real_execute = session.execute + + def insert_loses_the_race(statement, *args, **kwargs): + if isinstance(statement, sa.Insert) and statement.table is AssetTag.__table__: + raise sa.exc.IntegrityError("INSERT", {}, Exception("UNIQUE constraint failed")) + return real_execute(statement, *args, **kwargs) + + monkeypatch.setattr(session, "execute", insert_loses_the_race) + assert mark_contents_missing(session, [content.id]) == [content.id] + session.commit() + + link = session.get(AssetTag, (record.id, "missing")) + assert link is not None and link.origin == "automatic" + + +def test_the_batched_writes_run_no_table_scan_inside_their_transactions(session, db_engine, temp_dir, monkeypatch): + """Each statement a batch runs while it holds the write lock is looked up by key or + index, so the lock window grows with the batch, not the catalog. (The prune's + candidate read is a full read by design, and runs before any write transaction.)""" + in_write: list[bool] = [] + + class _Tracked(SASession): + def __enter__(self): + in_write.append(True) + return super().__enter__() + + def __exit__(self, *exc_info): + in_write.clear() + return super().__exit__(*exc_info) + + monkeypatch.setattr(scanner, "create_write_session", sessionmaker(bind=db_engine, class_=_Tracked)) + monkeypatch.setattr(scanner, "create_session", lambda: SASession(db_engine)) + # A size change splits the row, and the new record's tags come from its root. + monkeypatch.setattr("folder_paths.get_output_directory", lambda: str(temp_dir)) + gone_dir, changed_dir, pruned_dir = (temp_dir / name for name in ("gone", "changed", "pruned")) + for directory in (gone_dir, changed_dir, pruned_dir): + directory.mkdir() + observations = _gone_observations(session, gone_dir, 50) + for content_id in _rows(session, changed_dir, 2, files=True): + content = session.get(AssetContent, content_id) + if content.path.endswith("0.png"): + Path(content.path).write_bytes(b"a different size") + os.utime(content.path, ns=(10**18, 10**18)) + observations.append( + scanner._ReferenceObservation(content.id, content.size_bytes, content.mtime_ns, os.stat(content.path)) + ) + statements: list[tuple[str, object]] = [] + + def capture(conn, cursor, statement, parameters, context, executemany): + if in_write and statement.lstrip().upper().startswith(("SELECT", "UPDATE", "INSERT", "DELETE")): + statements.append((statement, parameters)) + + event.listen(db_engine, "before_cursor_execute", capture) + try: + scanner._write_in_batches(observations, scanner.apply_reference_observations, lambda: False, []) + _rows(session, pruned_dir, 50) + assert scanner.mark_missing_outside_prefixes_safely([str(gone_dir), str(changed_dir)]) == 50 + finally: + event.remove(db_engine, "before_cursor_execute", capture) + + kinds = {statement.split()[0].upper() for statement, _ in statements} + assert {"SELECT", "UPDATE", "INSERT"} <= kinds + with db_engine.connect() as conn: + for statement, parameters in statements: + plan = conn.exec_driver_sql(f"EXPLAIN QUERY PLAN {statement}", parameters).all() + scans = [row[-1] for row in plan if row[-1].startswith("SCAN")] + assert not scans, (statement, plan) + + +@pytest.mark.asyncio +async def test_the_prune_endpoint_keeps_the_event_loop_serving(monkeypatch): + def slow_prune() -> int: + time.sleep(0.5) + return 3 + + monkeypatch.setattr(routes.asset_seeder, "mark_missing_outside_prefixes", slow_prune) + ticks = 0 + + async def ticker(): + nonlocal ticks + while True: + await asyncio.sleep(0.01) + ticks += 1 + + ticking = asyncio.create_task(ticker()) + response = await routes.mark_missing_assets.__wrapped__(make_mocked_request("POST", "/api/assets/prune")) + ticking.cancel() + + assert json.loads(response.body) == {"status": "completed", "marked": 3} + assert ticks >= 20 + + + +def test_the_standalone_prune_starts_the_scan_queued_while_it_ran(session, catalog, temp_dir, monkeypatch): + """The API runs the prune off the event loop, so a prompt can finish meanwhile and + queue its output rescan, which cannot start while the prune holds the seeder.""" + _rows(session, temp_dir, 10) + instance = seeder_module._AssetSeeder() + monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True) + monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) + started: list[dict] = [] + + def start(**kwargs) -> bool: + if instance._state is not State.IDLE: + return False + started.append(kwargs) + return True + + monkeypatch.setattr(instance, "start", start) + real = scanner.mark_contents_missing + + def mark(sess, ids): + assert instance.enqueue_scan(roots=("output",), phase=seeder_module.ScanPhase.FULL) is False + return real(sess, ids) + + with patch("app.assets.scanner.mark_contents_missing", mark): + assert instance.mark_missing_outside_prefixes() == 10 + + assert [kwargs["roots"] for kwargs in started] == [("output",)] + assert instance._pending_scan is None + + +def _seeder_with_recorded_starts(monkeypatch) -> tuple[seeder_module._AssetSeeder, list[tuple]]: + instance = seeder_module._AssetSeeder() + monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True) + monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) + started: list[tuple] = [] + + def start(roots=("models", "input", "output"), **kwargs) -> bool: + if instance._state is not State.IDLE: + return False + started.append(tuple(roots)) + return True + + monkeypatch.setattr(instance, "start", start) + monkeypatch.setattr(routes, "asset_seeder", instance) + monkeypatch.setattr(routes, "_ASSETS_ENABLED", True) + return instance, started + + +@pytest.mark.asyncio +async def test_a_seed_request_during_an_api_prune_waits_for_it_then_starts(monkeypatch): + """The prune runs off the event loop now, so a seed request can arrive while it + holds the seeder. A 409 there would tell the client a scan is coming when none is.""" + instance, started = _seeder_with_recorded_starts(monkeypatch) + release = threading.Event() + pruning = threading.Event() + + def blocking_prune(prefixes, should_stop=None): + pruning.set() + assert release.wait(5) + return 0 + + monkeypatch.setattr(seeder_module, "mark_missing_outside_prefixes_safely", blocking_prune) + monkeypatch.setattr(routes, "_PRUNE_POLL_SECONDS", 0.01) + prune = asyncio.create_task(asyncio.to_thread(instance.mark_missing_outside_prefixes)) + assert await asyncio.to_thread(pruning.wait, 5) + + def must_not_block_a_thread(timeout=None): + raise AssertionError("the seed route held an executor thread for the prune") + + # The route waits on the loop; a blocking wait would hold an executor thread per request. + monkeypatch.setattr(instance, "wait_for_standalone_prune", must_not_block_a_thread) + seed = asyncio.create_task(routes.seed_assets.__wrapped__(make_mocked_request("POST", "/api/assets/seed"))) + await asyncio.sleep(0.2) + assert not seed.done() # waiting out the prune, not answering 409 + release.set() + response = await seed + await prune + + assert response.status == 202 + assert started == [("models", "input", "output")] + + +@pytest.mark.asyncio +async def test_a_seed_request_during_a_scan_still_gets_409(monkeypatch): + instance, started = _seeder_with_recorded_starts(monkeypatch) + instance._state = State.RUNNING + + response = await routes.seed_assets.__wrapped__(make_mocked_request("POST", "/api/assets/seed")) + + assert response.status == 409 + assert started == [] + + +def test_a_cancel_stops_a_standalone_prune_and_shutdown_waits_for_it(session, catalog, temp_dir, monkeypatch): + """The API prune runs on a worker thread that interpreter exit joins, so shutdown's + cancel must stop it between batches rather than let it run to the end.""" + _rows(session, temp_dir, 600) + instance = seeder_module._AssetSeeder() + monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True) + monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) + real = scanner.mark_contents_missing + first_batch = threading.Event() + shut_down = threading.Event() + + def mark(sess, ids): + marked = real(sess, ids) + first_batch.set() + assert shut_down.wait(5) # hold the first batch until shutdown has cancelled + return marked + + result: list[object] = [] + + def prune() -> None: + try: + result.append(instance.mark_missing_outside_prefixes()) + except seeder_module.PruneCancelledError as cancelled: + result.append(cancelled) + + with patch("app.assets.scanner.mark_contents_missing", mark): + worker = threading.Thread(target=prune) + worker.start() + assert first_batch.wait(5) + threading.Timer(0.1, shut_down.set).start() + assert instance.shutdown(timeout=5) + worker.join(5) + + assert len(catalog) == 1 + assert len(result) == 1 and isinstance(result[0], seeder_module.PruneCancelledError) + assert result[0].marked == scanner.WRITE_BATCH_ROWS + assert not instance.standalone_prune_running() + + +@pytest.mark.asyncio +async def test_a_cancelled_api_prune_is_not_reported_as_completed(monkeypatch): + def cancelled_prune() -> int: + raise seeder_module.PruneCancelledError(256) + + monkeypatch.setattr(routes.asset_seeder, "mark_missing_outside_prefixes", cancelled_prune) + response = await routes.mark_missing_assets.__wrapped__(make_mocked_request("POST", "/api/assets/prune")) + + assert response.status == 200 + assert json.loads(response.body) == {"status": "cancelled", "marked": 256} + + +def test_a_prune_that_finishes_before_a_late_cancel_reports_completed(session, catalog, temp_dir, monkeypatch): + """The cancel only counts if it stopped a batch from running.""" + _rows(session, temp_dir, 10) + instance = seeder_module._AssetSeeder() + monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True) + monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) + real = scanner.mark_contents_missing + + def mark(sess, ids): + marked = real(sess, ids) + instance.cancel() # arrives during the last (only) batch + return marked + + with patch("app.assets.scanner.mark_contents_missing", mark): + assert instance.mark_missing_outside_prefixes() == 10 + + +def test_shutdown_during_a_prune_does_not_start_the_scan_a_prompt_queued(session, catalog, temp_dir, monkeypatch): + """A scan started after shutdown cancelled the prune would run on into teardown.""" + _rows(session, temp_dir, 600) + instance = seeder_module._AssetSeeder() + monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True) + monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) + started: list[tuple] = [] + + def start(roots=("models", "input", "output"), **kwargs) -> bool: + if instance._state is not State.IDLE: + return False + started.append(tuple(roots)) + return True + + monkeypatch.setattr(instance, "start", start) + real = scanner.mark_contents_missing + first_batch = threading.Event() + shut_down = threading.Event() + + def mark(sess, ids): + marked = real(sess, ids) + # A prompt finishes and queues its output rescan; the prune holds the seeder. + assert instance.enqueue_scan(roots=("output",), phase=seeder_module.ScanPhase.FULL) is False + first_batch.set() + assert shut_down.wait(5) + return marked + + outcome: list[object] = [] + + def prune() -> None: + try: + outcome.append(instance.mark_missing_outside_prefixes()) + except seeder_module.PruneCancelledError as cancelled: + outcome.append(cancelled) + + with patch("app.assets.scanner.mark_contents_missing", mark): + worker = threading.Thread(target=prune) + worker.start() + assert first_batch.wait(5) + threading.Timer(0.1, shut_down.set).start() + assert instance.shutdown(timeout=5) + worker.join(5) + + # Asserted here, not in the worker: a failure there would only surface as a warning. + assert len(outcome) == 1 and isinstance(outcome[0], seeder_module.PruneCancelledError) + assert started == [] + assert instance._state is State.IDLE + + +def test_shutdown_before_a_prune_starts_keeps_it_from_starting(session, catalog, temp_dir, monkeypatch): + """The API hands the prune to a worker thread; a shutdown that lands before it takes + the seeder must still keep it from running into teardown.""" + _rows(session, temp_dir, 10) + instance = seeder_module._AssetSeeder() + monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True) + monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) + + assert instance.shutdown(timeout=1) + with pytest.raises(seeder_module.PruneCancelledError) as cancelled: + instance.mark_missing_outside_prefixes() + + assert cancelled.value.marked == 0 + assert len(catalog) == 0 + assert len(_live_ids(session)) == 10 + assert not instance.standalone_prune_running() + + +def test_the_prune_flag_clears_even_if_its_cleanup_raises(session, catalog, temp_dir, monkeypatch): + _rows(session, temp_dir, 10) + instance = seeder_module._AssetSeeder() + monkeypatch.setattr(seeder_module, "dependencies_available", lambda: True) + monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) + + def start_fails(): + raise RuntimeError("can't start new thread") + + monkeypatch.setattr(instance, "_finish_and_start_pending", start_fails) + with pytest.raises(RuntimeError): + instance.mark_missing_outside_prefixes() + + assert not instance.standalone_prune_running() + assert instance.wait_for_standalone_prune(0) diff --git a/tests-unit/assets_test/services/test_scanner_b.py b/tests-unit/assets_test/services/test_scanner_b.py index ad0828a9d..729d6d41a 100644 --- a/tests-unit/assets_test/services/test_scanner_b.py +++ b/tests-unit/assets_test/services/test_scanner_b.py @@ -5,14 +5,13 @@ from pathlib import Path from unittest.mock import patch from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, sessionmaker from app.assets.database.models import Asset, AssetContent, AssetTag from app.assets.helpers import to_stored_hash from app.assets.scanner import ( build_asset_specs, enrich_asset, - mark_contents_missing_outside_prefixes, mark_missing_outside_prefixes_safely, seed_asset_specs, apply_reference_observations, @@ -151,7 +150,7 @@ def test_seed_creates_content_and_record(session, temp_dir: Path): ] -def test_prune_marks_missing_not_deletes(session, temp_dir: Path): +def test_prune_marks_missing_not_deletes(session, db_engine, temp_dir: Path): input_root = temp_dir / "input" input_root.mkdir() file_path = input_root / "removed-from-registry.png" @@ -161,8 +160,11 @@ def test_prune_marks_missing_not_deletes(session, temp_dir: Path): seed_asset_specs(session, _build_seed_specs(input_root)) session.commit() - marked = mark_contents_missing_outside_prefixes(session, prefixes=[]) - session.commit() + with patch("app.assets.scanner.create_session", sessionmaker(bind=db_engine)), \ + patch("app.database.db.WriteSession", sessionmaker(bind=db_engine)), \ + patch("app.assets.scanner.get_owned_prefixes", return_value=[]): + marked = mark_missing_outside_prefixes_safely([]) + session.expire_all() content = session.scalar(select(AssetContent)) record = session.scalar(select(Asset)) diff --git a/tests-unit/seeder_test/test_seeder.py b/tests-unit/seeder_test/test_seeder.py index a403170d7..6b1c12cc3 100644 --- a/tests-unit/seeder_test/test_seeder.py +++ b/tests-unit/seeder_test/test_seeder.py @@ -93,7 +93,7 @@ def _configure_fast_phase( specs: list[SeedAssetSpec], ) -> None: monkeypatch.setattr( - seeder_module, "sync_root_safely", lambda _root, _progress: set() + seeder_module, "sync_root_safely", lambda _root, _progress, _should_stop=None: set() ) monkeypatch.setattr( seeder_module, @@ -435,7 +435,7 @@ def test_idle_reset_survives_a_raising_cancellation_emit( monkeypatch.setattr(scan_seeder, "_check_pause_and_cancel", cancel_at_pruning) monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ()) monkeypatch.setattr( - seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 0 + seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 0 ) original_emit = seeder_module.emit @@ -462,7 +462,7 @@ def test_scan_paused_after_its_last_phase_still_completes( scan_seeder._phase = ScanPhase.ENRICH monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ()) monkeypatch.setattr( - seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 0 + seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 0 ) def pause_while_finishing(roots) -> tuple[bool, int]: @@ -507,10 +507,10 @@ def test_prune_before_scan_emits_marked_missing_with_pruning_stage( scan_seeder._phase = ScanPhase.FAST monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ()) monkeypatch.setattr( - seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 5 + seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 5 ) monkeypatch.setattr( - seeder_module, "sync_temp_references_safely", lambda _progress: None + seeder_module, "sync_temp_references_safely", lambda _progress, _should_stop=None: None ) monkeypatch.setattr(scan_seeder, "_run_fast_phase", lambda roots: (0, 0, 0)) @@ -530,7 +530,7 @@ def test_standalone_mark_missing_emits_count_with_mark_missing_stage( scan_seeder._state = State.IDLE monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: ()) monkeypatch.setattr( - seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes: 7 + seeder_module, "mark_missing_outside_prefixes_safely", lambda prefixes, _should_stop=None: 7 ) with caplog.at_level(logging.INFO): @@ -575,10 +575,10 @@ def test_scan_prune_failure_is_reported_and_the_scan_still_runs( fast_phase_roots: list[tuple[str, ...]] = [] monkeypatch.setattr(seeder_module, "get_owned_prefixes", lambda: []) monkeypatch.setattr( - seeder_module, "mark_missing_outside_prefixes_safely", lambda _prefixes: None + seeder_module, "mark_missing_outside_prefixes_safely", lambda _prefixes, _should_stop=None: None ) monkeypatch.setattr( - seeder_module, "sync_temp_references_safely", lambda _progress: None + seeder_module, "sync_temp_references_safely", lambda _progress, _should_stop=None: None ) def run_fast_phase(roots: tuple[str, ...]) -> tuple[int, int, int]: @@ -591,7 +591,7 @@ def test_scan_prune_failure_is_reported_and_the_scan_still_runs( scan_seeder._run_scan() assert scan_seeder._errors == [ - "Marking missing assets failed; scan continued without pruning" + "Marking missing assets failed; scan continued with the prune incomplete" ] assert fast_phase_roots == [("models", "input")] assert events_named(caplog, "seeder.marked_missing") == [] @@ -604,7 +604,7 @@ def test_batch_insert_failure_emits_only_the_exception_type( ) -> None: session = Mock() monkeypatch.setattr( - seeder_module, "sync_root_safely", lambda _root, _progress: set() + seeder_module, "sync_root_safely", lambda _root, _progress, _should_stop=None: set() ) monkeypatch.setattr( seeder_module,