From ebebf0d69c88663d8a53fab9202768a8ef2a0579 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 22 Oct 2026 19:17:11 -0700 Subject: [PATCH] fix(jobs): refuse legacy record-changing operations while a job runs The job registry gains admission classes. A records job is refused while any records or guarded job runs, a guarded job while a records job runs, and read jobs share only the one-job-per-kind rule. record_changing=True stays as the spelling of records. Job summaries add cancellable and admission, and a non-cancellable job ignores cancel. JobRegistry.hold() registers a running, non-cancellable job for a legacy-page operation and ends it when the request does. The legacy restore, maintenance, prompt cleanup, duplicate image delete, LoRA import and batch auto-tag routes hold as records; thumbnail generation (both routes) and thumbnail clear hold as guarded. A maintenance run of only read checks holds as read. So a restore during a scan, or a scan during a legacy restore, is refused with the running job named instead of interleaving; this closes the gap noted on #269. Refusals keep each route's shape: the JSON routes return the legacy 409 " already running", and the thumbnail and LoRA streams send their existing error event. --- py/api/__init__.py | 21 +- py/api/admin.py | 74 +++++- py/api/autotag_routes.py | 22 ++ py/api/images.py | 52 +++- py/api/jobs.py | 172 ++++++++++++-- py/api/lora_integration.py | 32 ++- tests/test_api_backup_restore.py | 83 +++++++ tests/test_api_images.py | 89 +++++++ tests/test_image_tag_jobs.py | 17 ++ tests/test_job_admission.py | 395 +++++++++++++++++++++++++++++++ tests/test_jobs.py | 5 + 11 files changed, 918 insertions(+), 44 deletions(-) create mode 100644 tests/test_job_admission.py diff --git a/py/api/__init__.py b/py/api/__init__.py index fbebf5c..a7ef474 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -34,7 +34,7 @@ from .images import ImageRoutesMixin from .admin import AdminRoutesMixin from .job_routes import JobRoutesMixin from .image_tag_jobs import ImageTagJobsMixin -from .jobs import JobRegistry +from .jobs import JobBusy, JobRegistry from .logging_routes import LoggingRoutesMixin from .autotag_routes import AutotagRoutesMixin from .lora_integration import LoraIntegrationMixin @@ -361,6 +361,25 @@ class PromptManagerAPI( """409 envelope for a long job that is already running.""" return _json_error(f"{label} already running", 409) + async def _run_held(self, kind, title, admission, operation): + """Run a legacy-page operation under a registry hold (D80). + + While ``operation()`` runs, the registry lists it as a running, + non-cancellable job, so workspace jobs its admission class conflicts + with are refused, naming it. When the registry refuses the hold, the + legacy 409 envelope names the running job and nothing is done. A + response of 400 or above ends the hold ``failed``. + """ + try: + held = self._jobs.hold(kind, title, admission=admission) + except JobBusy as busy: + return self._long_job_busy_response(busy.running.title) + async with held as job: + response = await operation() + if getattr(response, "status", 200) >= 400: + job.finish("failed") + return response + async def _run_in_executor(self, func, *args, **kwargs): """Run a blocking function in the default thread pool executor. diff --git a/py/api/admin.py b/py/api/admin.py index af16fbc..bb93b42 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -22,6 +22,9 @@ from .jobs import ( SCAN_GROUP_NAMES, SCAN_GROUPS, SCAN_SCOPES, + DUPLICATE_CLEANUP_JOB_KIND, + MAINTENANCE_JOB_KIND, + RESTORE_JOB_KIND, Job, JobBusy, ) @@ -696,7 +699,15 @@ class AdminRoutesMixin: job.finish() async def cleanup_duplicates_endpoint(self, request): - """Cleanup duplicate prompts endpoint.""" + """Cleanup duplicate prompts endpoint (holds the registry: records).""" + return await self._run_held( + MAINTENANCE_JOB_KIND, + "Prompt cleanup (admin page)", + "records", + self._cleanup_duplicates, + ) + + async def _cleanup_duplicates(self): try: removed_count = await self._run_in_executor(self.db.cleanup_duplicates) @@ -811,12 +822,26 @@ class AdminRoutesMixin: return hash_sha256.hexdigest() async def delete_duplicate_images_endpoint(self, request): - """Delete duplicate image files from disk (single flight with the scan).""" - job = self._acquire_long_job(self.LONG_JOB_DUPLICATES) - if job is None: + """Delete duplicate image files from disk. + + Single flight with the duplicate scan, and a records hold on the + registry, so no scan, restore or tagging job runs meanwhile. + """ + guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES) + if guard is None: return self._long_job_busy_response("Duplicate scan or delete") - async with job: - return await self._delete_duplicate_images(request) + + async def delete(): + # No await since the check above: the lock is still free. + async with guard: + return await self._delete_duplicate_images(request) + + return await self._run_held( + DUPLICATE_CLEANUP_JOB_KIND, + "Duplicate image delete (gallery page)", + "records", + delete, + ) async def _delete_duplicate_images(self, request): try: @@ -1268,6 +1293,11 @@ class AdminRoutesMixin: "vacuum", "cleanup_orphaned_images", ) + # Operations that only read; a run of only these holds the registry as + # "read", every other run as "records". + READ_ONLY_MAINTENANCE_OPERATIONS = frozenset( + {"check_hash_duplicates", "statistics", "check_consistency"} + ) async def run_maintenance(self, request): """Perform database maintenance operations and report each outcome.""" @@ -1280,7 +1310,27 @@ class AdminRoutesMixin: operations = data.get( "operations", list(self.DEFAULT_MAINTENANCE_OPERATIONS) ) + except Exception as e: + self.logger.error(f"Maintenance error: {e}", exc_info=True) + return web.json_response( + {"success": False, "error": "Maintenance failed"}, status=500 + ) + writes = not ( + isinstance(operations, list) + and all( + isinstance(name, str) and name in self.READ_ONLY_MAINTENANCE_OPERATIONS + for name in operations + ) + ) + return await self._run_held( + MAINTENANCE_JOB_KIND, + "Maintenance (admin page)", + "records" if writes else "read", + lambda: self._run_maintenance(operations), + ) + async def _run_maintenance(self, operations): + try: results = await self._run_in_executor( self._run_maintenance_operations, operations ) @@ -1434,8 +1484,18 @@ class AdminRoutesMixin: The upload is streamed to a temp file (capped at restore_max_bytes), verified with PromptModel.verify_database_file and only then copied - over the live database, which is backed up first. + over the live database, which is backed up first. The whole request + holds the registry (records): it is refused while a scan or another + record-changing job runs, and those are refused while it runs. """ + return await self._run_held( + RESTORE_JOB_KIND, + "Database restore (admin page)", + "records", + lambda: self._restore_database(request), + ) + + async def _restore_database(self, request): max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024) temp_path = None try: diff --git a/py/api/autotag_routes.py b/py/api/autotag_routes.py index 999b941..df922e2 100644 --- a/py/api/autotag_routes.py +++ b/py/api/autotag_routes.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import Optional from urllib.parse import quote +from .jobs import AUTOTAG_BATCH_JOB_KIND, JobBusy from .prompts import thumbnail_url_for from .. import model_gate @@ -440,6 +441,27 @@ class AutotagRoutesMixin: "message": _INTERNAL_ERROR, } + # The batch writes prompt tags: a records hold on the registry, so it + # never runs beside a scan, restore or the workspace tagging job. + try: + held = self._jobs.hold( + AUTOTAG_BATCH_JOB_KIND, + "Auto-tag images (legacy page)", + admission="records", + ) + except JobBusy as busy: + if busy.running.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload + return web.json_response( + {"success": False, "error": _BUSY_ERROR}, status=409 + ) + return self._long_job_busy_response(busy.running.title) + async with held as job: + response = await self._start_autotag_batch(request, opts, events) + if response.status >= 400: + job.finish("failed") + return response + + async def _start_autotag_batch(self, request, opts, events): if not _autotag_batch_lock.acquire(blocking=False): return web.json_response( {"success": False, "error": _BUSY_ERROR}, status=409 diff --git a/py/api/images.py b/py/api/images.py index 6edfdc1..bf271a5 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -19,6 +19,7 @@ from typing import Callable, NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError +from .jobs import THUMBNAILS_JOB_KIND, JobBusy from .prompts import bad_request, parse_page_params, publish_image_paths from .prompts import thumbnail_url_for from .prompts import safe_error_message as _safe_error @@ -104,6 +105,10 @@ def _card_header(path: Path) -> tuple: return inspection.width, inspection.height, records +# The legacy gallery page's thumbnail generation, as Activity lists it (D80). +THUMBNAIL_GENERATION_TITLE = "Thumbnail generation (gallery page)" + + class _CardFacts(NamedTuple): """One image file's card facts, derived from its header alone. @@ -1736,7 +1741,19 @@ class ImageRoutesMixin: return cls.THUMBNAIL_SIZES.get(quality, cls.THUMBNAIL_SIZES["medium"]) async def generate_thumbnails(self, request): - """Generate every thumbnail in one executor job (POST, blocking).""" + """Generate every thumbnail in one executor job (POST, blocking). + + Holds the registry as guarded work: refused while a record-changing + job runs, and such jobs are refused while it runs. + """ + return await self._run_held( + THUMBNAILS_JOB_KIND, + THUMBNAIL_GENERATION_TITLE, + "guarded", + lambda: self._generate_thumbnails(request), + ) + + async def _generate_thumbnails(self, request): try: data, error = await _json_body_or_error(request) if error is not None: @@ -1935,14 +1952,31 @@ class ImageRoutesMixin: "Connection": "keep-alive", }, ) - await response.prepare(request) + try: + held = self._jobs.hold( + THUMBNAILS_JOB_KIND, THUMBNAIL_GENERATION_TITLE, admission="guarded" + ) + except JobBusy as busy: + # The gallery reads this stream: refuse with its error event. + await response.prepare(request) + message = f"{busy.running.title} already running" + await self._emit_sse( + response, "error", {"error": message, "message": message} + ) + return response + async with held: + await response.prepare(request) + await self._thumbnail_progress_stream(response, thumbnail_size) + return response + + async def _thumbnail_progress_stream(self, response, thumbnail_size): try: output_dirs = self._get_all_output_dirs() if not output_dirs: await self._emit_sse( response, "error", {"error": "ComfyUI output directory not found"} ) - return response + return await self._emit_sse( response, @@ -1970,7 +2004,6 @@ class ImageRoutesMixin: "message": "Thumbnail generation failed. Check server logs.", }, ) - return response async def _stream_thumbnails( self, response, targets, thumbnail_size, skipped=0, remaining=0 @@ -2179,8 +2212,17 @@ class ImageRoutesMixin: """Clear generated thumbnails across gallery roots, preserving originals. Invalid thumbnail directories return 400; unavailable gallery roots - return 404, and unexpected filesystem failures return 500. + return 404, and unexpected filesystem failures return 500. Holds the + registry as guarded work, like generation. """ + return await self._run_held( + THUMBNAILS_JOB_KIND, + "Thumbnail clear (gallery page)", + "guarded", + self._clear_thumbnails, + ) + + async def _clear_thumbnails(self) -> web.Response: try: output_dirs = self._get_all_output_dirs() if not output_dirs: diff --git a/py/api/jobs.py b/py/api/jobs.py index 27e4ecb..112c464 100644 --- a/py/api/jobs.py +++ b/py/api/jobs.py @@ -9,9 +9,21 @@ as a snapshot and then every event published after it, until the job ends. Jobs live only in server memory and are gone after a restart. The registry keeps the running jobs and this session's last ``MAX_FINISHED_JOBS`` finished -ones. Record-changing jobs (the metadata scan, bulk image tagging, and later -maintenance, cleanup and restore) never run together, and one job of a kind -runs at a time. +ones. One job of a kind runs at a time, and each job has an admission class +(``ADMISSIONS``) that decides what may run beside it: + +- ``records`` changes library records or files (the metadata scan, bulk image + tagging, and the legacy restore, maintenance, cleanup, duplicate delete, + LoRA import and batch auto-tag routes). It is refused while any + ``records`` or ``guarded`` job runs. +- ``guarded`` must not overlap record changes but may overlap other guarded + work (thumbnail generation and clearing, later backup). It is refused while + a ``records`` job runs. +- ``read`` only reads; it is refused only by a running job of its own kind. + +A short legacy-page operation takes a ``hold``: a running, non-cancellable +job that lasts as long as the request, so the Activity page lists it and a +workspace job is refused while it runs (and the other way round). """ import asyncio @@ -60,6 +72,41 @@ SCAN_GROUPS = { IMAGE_TAG_JOB_KIND = "image-tags" IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors") +ADMISSIONS = ("read", "guarded", "records") +# What each admission class is refused beside (besides its own kind). +_REFUSED_BESIDE = { + "read": frozenset(), + "guarded": frozenset({"records"}), + "records": frozenset({"records", "guarded"}), +} + +# Kinds of the legacy-page operations that hold the registry (D80; titles in +# their handlers end in "(admin page)" or "(gallery page)"). +RESTORE_JOB_KIND = "restore" +MAINTENANCE_JOB_KIND = "maintenance" +DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup" +THUMBNAILS_JOB_KIND = "thumbnails" +LORA_IMPORT_JOB_KIND = "lora_import" +AUTOTAG_BATCH_JOB_KIND = "autotag" + + +def resolve_admission(admission=None, record_changing=None) -> str: + """The admission class for a start call; ``record_changing`` is an alias. + + ``record_changing=True`` means ``records`` and ``False`` means ``read``; + with neither given the class is ``read``. Raises ValueError for an + unknown class or one that contradicts ``record_changing``. + """ + if admission is None: + return "records" if record_changing else "read" + if admission not in ADMISSIONS: + raise ValueError(f"Unknown job admission {admission!r}") + if record_changing is not None and record_changing != (admission == "records"): + raise ValueError( + f"record_changing={record_changing!r} contradicts admission {admission!r}" + ) + return admission + def _relative_path(path) -> str: """``path`` as a root-relative POSIX path; an absolute one keeps its name.""" @@ -85,13 +132,23 @@ class Job: """One tracked operation: state, counters, outcomes, log and subscribers.""" def __init__( - self, kind, title, *, record_changing=False, scope=None, groups=(), clock=None + self, + kind, + title, + *, + admission=None, + record_changing=None, + cancellable=True, + scope=None, + groups=(), + clock=None, ): self._clock = clock or time.time self.id = uuid.uuid4().hex self.kind = kind self.title = title - self.record_changing = record_changing + self.admission = resolve_admission(admission, record_changing) + self.cancellable = bool(cancellable) self.scope = dict(scope or {}) self.state = "running" self.started_at = self._clock() @@ -110,6 +167,10 @@ class Job: self.task = None self._subscribers = set() + @property + def record_changing(self): + return self.admission == "records" + @property def running(self): return self.state == "running" @@ -240,6 +301,9 @@ class Job: # What detail() lists, so a view re-reads it only when it changed. "listed": {group: len(rows) for group, rows in self.outcomes.items()}, "log_length": len(self.log), + # A view offers Cancel only when the server can stop the job. + "cancellable": self.cancellable, + "admission": self.admission, } def detail(self): @@ -291,33 +355,76 @@ class JobRegistry: self._logger = logger_ or logger self._jobs = {} - def start(self, kind, title, *, record_changing, scope, runner, groups=()) -> Job: + def start( + self, + kind, + title, + *, + admission=None, + record_changing=None, + cancellable=True, + scope=None, + runner, + groups=(), + ) -> Job: """Create a job and run ``runner(job)`` as a task. - Raises JobBusy, naming the job in the way, when a job of ``kind`` is - running or when ``record_changing`` and another record-changing job - runs. Must be called from the event loop. + ``admission`` is one of ``ADMISSIONS`` (default ``read``); + ``record_changing=True`` is the older spelling of ``records``. Raises + JobBusy, naming the job in the way, when the admission class refuses + it (see the module docstring). Must be called from the event loop. """ - blocker = self.blocker(kind, record_changing=record_changing) - if blocker is not None: - raise JobBusy(blocker) - self._prune() - job = Job( - kind, title, record_changing=record_changing, scope=scope, groups=groups + job = self._admit( + kind, + title, + admission=admission, + record_changing=record_changing, + cancellable=cancellable, + scope=scope, + groups=groups, ) - self._jobs[job.id] = job job.task = asyncio.ensure_future(self._run(job, runner)) return job - def blocker(self, kind, *, record_changing): - """The running job a new ``kind`` job would be refused for, or None.""" + def hold(self, kind, title, *, admission, scope=None): + """Register a running, non-cancellable job for a legacy-page operation. + + Returns an async context manager yielding the job; the job ends when + the block exits: ``finished``, or ``failed`` when it raised (unless + the body already ended it). Raises JobBusy at once, without waiting, + when the admission class refuses it. Use it directly in ``async with`` + so the job cannot outlive the operation. + """ + job = self._admit( + kind, title, admission=admission, cancellable=False, scope=scope + ) + return _Hold(job) + + def _admit(self, kind, title, *, admission, record_changing=None, **job_args): + resolved = resolve_admission(admission, record_changing) + blocker = self.blocker(kind, admission=resolved) + if blocker is not None: + raise JobBusy(blocker) + self._prune() + job = Job(kind, title, admission=resolved, **job_args) + self._jobs[job.id] = job + return job + + def blocker(self, kind, *, admission=None, record_changing=None): + """The running job a new ``kind`` job would be refused for, or None. + + The same kind refuses first; then the oldest running job whose class + the new job's admission class may not run beside. + """ same_kind = self.active(kind) if same_kind is not None: return same_kind - if record_changing: - for job in self._jobs.values(): - if job.running and job.record_changing: - return job + refused = _REFUSED_BESIDE[resolve_admission(admission, record_changing)] + running = [job for job in self._jobs.values() if job.running] + running.sort(key=lambda job: job._started_seq) + for job in running: + if job.admission in refused: + return job return None async def _run(self, job, runner): @@ -362,11 +469,26 @@ class JobRegistry: def cancel(self, job_id) -> bool: """Ask a running job to stop at its next safe boundary. - Returns whether a cancellation was requested: False for a finished or - unknown job, so a cancel after the end is a no-op. + Returns whether a cancellation was requested: False for a finished, + unknown or non-cancellable job, so a cancel after the end is a no-op. """ job = self._jobs.get(job_id) - if job is None or job.finished: + if job is None or job.finished or not job.cancellable: return False job.cancelled = True return True + + +class _Hold: + """The async context manager ``JobRegistry.hold`` returns.""" + + def __init__(self, job): + self._job = job + + async def __aenter__(self): + return self._job + + async def __aexit__(self, exc_type, exc, tb): + if self._job.running: + self._job.finish("finished" if exc_type is None else "failed") + return False diff --git a/py/api/lora_integration.py b/py/api/lora_integration.py index 3a39982..c24b852 100644 --- a/py/api/lora_integration.py +++ b/py/api/lora_integration.py @@ -6,6 +6,7 @@ from pathlib import Path from aiohttp import web from ..config_store import ConfigUnreadable, RevisionConflict, get_config_store +from .jobs import LORA_IMPORT_JOB_KIND, JobBusy # Flags of POST /lora/enable; each must be a JSON boolean when present. _LORA_FLAGS = ("enabled", "trigger_words_enabled", "clear_civitai_api_key") @@ -273,19 +274,38 @@ class LoraIntegrationMixin: reason="OK", headers={"Content-Type": "text/event-stream"}, ) - await response.prepare(request) async def send_progress(data): line = f"data: {json.dumps(data)}\n\n" await response.write(line.encode("utf-8")) + # The import replaces prompts and images: a records hold, so it + # never runs beside a scan, restore or tagging job. try: - await self._lora_scan_stream(allowed, send_progress) - except Exception as e: - self.logger.exception("LoRA scan failed") - await send_progress({"type": "error", "message": self._public_error(e)}) + held = self._jobs.hold( + LORA_IMPORT_JOB_KIND, + "LoRA import (admin page)", + admission="records", + ) + except JobBusy as busy: + await response.prepare(request) + message = f"{busy.running.title} already running" + await send_progress({"type": "error", "message": message}) + await response.write_eof() + return response - await response.write_eof() + async with held as job: + await response.prepare(request) + try: + await self._lora_scan_stream(allowed, send_progress) + except Exception as e: + self.logger.exception("LoRA scan failed") + job.finish("failed") + await send_progress( + {"type": "error", "message": self._public_error(e)} + ) + + await response.write_eof() return response except Exception as e: diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index d91b302..aa211ff 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -5,12 +5,16 @@ stream the upload to disk, reject anything that is not a healthy prompts database and leave the live database untouched when it does. """ +import asyncio import glob +import hashlib import os import sqlite3 import sys import tempfile +import threading import unittest +from pathlib import Path from unittest.mock import MagicMock, patch _mock_server = MagicMock() @@ -142,6 +146,85 @@ class TestBackupRoute(BackupRestoreTestCase): self.assertIn("backup", data["error"].lower()) +class TestRestoreHoldsTheRegistry(BackupRestoreTestCase): + """A legacy restore and a workspace job refuse each other (#269 gap).""" + + def _digest(self): + return { + suffix: hashlib.sha256(Path(self.db_path + suffix).read_bytes()).hexdigest() + for suffix in ("", "-wal") + if os.path.exists(self.db_path + suffix) + } + + async def test_legacy_restore_refused_while_scan_job_runs(self): + self._save("keep me") + upload = self._make_db_file(["restored"]) + gate = asyncio.Event() + + async def runner(job): + await gate.wait() + + scan = self.api._jobs.start( + "scan", "Scan for metadata", record_changing=True, scope={}, runner=runner + ) + before = self._digest() + try: + resp = await self._restore(upload) + self.assertEqual(resp.status, 409) + self.assertEqual( + await resp.json(), + {"success": False, "error": "Scan for metadata already running"}, + ) + self.assertEqual(self._digest(), before) + self.assertEqual( + [p["text"] for p in self.api.db.search_prompts()], ["keep me"] + ) + self.assertEqual(glob.glob(self.db_path + ".backup_*"), []) + finally: + gate.set() + await asyncio.wait_for(scan.task, 5) + + async def test_scan_refused_while_legacy_restore_holds(self): + upload = self._make_db_file(["restored"]) + entered = threading.Event() + release = threading.Event() + model = self.api.db.model + real_verify = model.verify_database_file + + def gated(path, **kwargs): + entered.set() + release.wait(10) + return real_verify(path, **kwargs) + + loop = asyncio.get_running_loop() + with patch.object(model, "verify_database_file", gated): + restore = asyncio.ensure_future(self._restore(upload)) + self.assertTrue(await loop.run_in_executor(None, entered.wait, 10)) + resp = await self.client.post("/prompt_manager/jobs/scan") + body = await resp.json() + listed = (await (await self.client.get("/prompt_manager/jobs")).json())[ + "jobs" + ] + release.set() + restored = await restore + + self.assertEqual((resp.status, body["code"]), (409, "busy")) + self.assertEqual(body["running"]["title"], "Database restore (admin page)") + self.assertEqual( + [(j["kind"], j["state"], j["cancellable"]) for j in listed], + [("restore", "running", False)], + ) + self.assertEqual(restored.status, 200) + self.assertEqual(self.api._jobs.list()[0].state, "finished") + self.assertIsNone(self.api._jobs.active("restore")) + + async def test_rejected_upload_ends_its_hold_failed(self): + resp = await self._restore(self._write("notes.txt", b"not a database")) + self.assertEqual(resp.status, 400) + job = self.api._jobs.list()[0] + self.assertEqual((job.kind, job.state), ("restore", "failed")) + + class TestRestoreRoute(BackupRestoreTestCase): async def test_restore_replaces_database_and_reports_backup(self): self._save("old prompt") diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 18062bb..8c0361c 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -5,11 +5,13 @@ Uses aiohttp's test client against a real PromptManagerAPI wired to a temporary SQLite database and a temporary output directory. """ +import asyncio import json import os import pathlib import sys import tempfile +import threading import unittest from contextlib import ExitStack from pathlib import Path @@ -638,6 +640,93 @@ class TestGenerateThumbnailsPost(ThumbnailTestCase): self.assertEqual(data["total_images"], 0) +class TestThumbnailRoutesHoldTheRegistry(ThumbnailTestCase): + """Thumbnail work is ``guarded``: never beside a record-changing job.""" + + def _run_scan(self): + gate = asyncio.Event() + + async def runner(job): + await gate.wait() + + job = self.api._jobs.start( + "scan", "Scan for metadata", admission="records", scope={}, runner=runner + ) + + async def release(): + gate.set() + await asyncio.wait_for(job.task, 5) + + self.addAsyncCleanup(release) + return job + + async def test_thumbnail_generate_refused_while_scan_runs(self): + make_png(self.output_dir / "a.png") + self._run_scan() + + resp = await self.client.request( + "POST", "/prompt_manager/images/generate-thumbnails", json={} + ) + self.assertEqual(resp.status, 409) + self.assertEqual( + await resp.json(), + {"success": False, "error": "Scan for metadata already running"}, + ) + + resp = await self.client.request("POST", self.PROGRESS_URL) + text = await resp.text() + self.assertEqual(resp.status, 200) # the gallery reads an SSE stream + self.assertEqual( + text, + "event: error\ndata: " + + json.dumps( + { + "error": "Scan for metadata already running", + "message": "Scan for metadata already running", + } + ) + + "\n\n", + ) + + thumbs = self.output_dir / "thumbnails" + make_png(thumbs / "b_thumb.png") + resp = await self.client.request( + "POST", "/prompt_manager/images/clear-thumbnails" + ) + self.assertEqual(resp.status, 409) + self.assertFalse(self._thumb_path("a").exists()) + self.assertTrue((thumbs / "b_thumb.png").exists()) + + async def test_scan_refused_while_thumbnails_generate(self): + make_png(self.output_dir / "a.png") + entered = threading.Event() + release = threading.Event() + real = self.api._generate_gallery_thumbnails_sync + + def gated(*args): + entered.set() + release.wait(10) + return real(*args) + + self.api._generate_gallery_thumbnails_sync = gated + loop = asyncio.get_running_loop() + generate = asyncio.ensure_future( + self.client.request( + "POST", "/prompt_manager/images/generate-thumbnails", json={} + ) + ) + self.assertTrue(await loop.run_in_executor(None, entered.wait, 10)) + resp = await self.client.request("POST", "/prompt_manager/jobs/scan") + body = await resp.json() + self.assertEqual((resp.status, body["code"]), (409, "busy")) + self.assertEqual( + body["running"]["title"], "Thumbnail generation (gallery page)" + ) + release.set() + self.assertEqual((await generate).status, 200) + self.assertTrue(self._thumb_path("a").exists()) + + class TestRecentImagesCaps(ImageAPITestCase): """Every path through /images/recent hands the DB a bounded window.""" diff --git a/tests/test_image_tag_jobs.py b/tests/test_image_tag_jobs.py index 3332422..8964b07 100644 --- a/tests/test_image_tag_jobs.py +++ b/tests/test_image_tag_jobs.py @@ -420,6 +420,23 @@ class TestStartAndExclusion(ImageTagJobTestCase): # A refused start never leaves the gate taken. self.assertIsNone(model_gate.holder()) + async def test_image_tag_job_refused_while_legacy_restore_holds(self): + self.png("a.png") + async with self.api._jobs.hold( + "restore", "Database restore (admin page)", admission="records" + ): + status, body = await self.start(self.add_body(["a.png"], ["x"])) + self.assertEqual((status, body["code"]), (409, "busy")) + self.assertEqual(body["running"]["title"], "Database restore (admin page)") + status, body = await self.start(self.auto_body(["a.png"])) + self.assertEqual((status, body["code"]), (409, "busy")) + self.assertIsNone(model_gate.holder()) + async with self.api._jobs.hold( + "thumbnails", "Thumbnail generation (gallery page)", admission="guarded" + ): + status, body = await self.start(self.add_body(["a.png"], ["x"])) + self.assertEqual((status, body["code"]), (409, "busy")) + async def test_record_changing_excludes_running_scan(self): self.png("a.png") gate = asyncio.Event() diff --git a/tests/test_job_admission.py b/tests/test_job_admission.py new file mode 100644 index 0000000..266d71e --- /dev/null +++ b/tests/test_job_admission.py @@ -0,0 +1,395 @@ +"""Registry admission classes and the legacy routes that hold the registry. + +Record-changing work never runs together (Activity design): a ``records`` +job is refused while any ``records`` or ``guarded`` job runs, a ``guarded`` +job (thumbnails) while a ``records`` job runs, and ``read`` jobs only share +the one-per-kind rule. The legacy admin and gallery routes that change +records or files take a registry hold for their duration, so a workspace job +and a legacy-page operation refuse each other, naming the one that runs. + +Fake runners and gated executor calls only; ordering is driven by events. +""" + +import asyncio +import hashlib +import json +import os +import sys +import threading +import unittest +from pathlib import Path +from unittest.mock import MagicMock, patch + +_mock_server = MagicMock() +_mock_server.PromptServer.instance.routes = MagicMock() +sys.modules.setdefault("server", _mock_server) + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from py.api.jobs import ( # noqa: E402 + ADMISSIONS, + Job, + JobBusy, + JobRegistry, +) + +try: + from tests.test_api_admin import AdminAPITestCase # noqa: E402 +except ImportError: # discovered with tests/ as the top-level directory + from test_api_admin import AdminAPITestCase # noqa: E402 + + +def _gated(gate): + async def runner(job): + await gate.wait() + + return runner + + +async def _noop(job): + return None + + +class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase): + + def _start(self, registry, kind, admission, gate=None, **kwargs): + return registry.start( + kind, + kind.title(), + admission=admission, + scope={}, + runner=_gated(gate) if gate is not None else _noop, + **kwargs, + ) + + async def test_admission_names_and_default(self): + self.assertEqual(ADMISSIONS, ("read", "guarded", "records")) + registry = JobRegistry() + job = registry.start("report", "Report", scope={}, runner=_noop) + await asyncio.wait_for(job.task, 2) + self.assertEqual(job.admission, "read") + self.assertFalse(job.record_changing) + with self.assertRaises(ValueError): + registry.start("x", "X", admission="exclusive", scope={}, runner=_noop) + with self.assertRaises(ValueError): + registry.start( + "x", + "X", + admission="read", + record_changing=True, + scope={}, + runner=_noop, + ) + + async def test_record_changing_is_an_alias_for_records(self): + registry = JobRegistry() + gate = asyncio.Event() + job = registry.start( + "scan", "Scan", record_changing=True, scope={}, runner=_gated(gate) + ) + self.assertEqual(job.admission, "records") + self.assertTrue(job.record_changing) + self.assertEqual(job.summary()["admission"], "records") + self.assertIs( + registry.blocker("image-tags", record_changing=True), job + ) # the existing call shape keeps working + gate.set() + await asyncio.wait_for(job.task, 2) + + async def test_records_job_refused_while_guarded_runs(self): + registry = JobRegistry() + gate = asyncio.Event() + thumbs = self._start(registry, "thumbnails", "guarded", gate) + with self.assertRaises(JobBusy) as raised: + self._start(registry, "scan", "records") + self.assertIs(raised.exception.running, thumbs) + self.assertIs(registry.blocker("restore", admission="records"), thumbs) + gate.set() + await asyncio.wait_for(thumbs.task, 2) + scan = self._start(registry, "scan", "records") + await asyncio.wait_for(scan.task, 2) + self.assertEqual(scan.state, "finished") + + async def test_guarded_refused_while_records_runs(self): + registry = JobRegistry() + gate = asyncio.Event() + scan = self._start(registry, "scan", "records", gate) + with self.assertRaises(JobBusy) as raised: + self._start(registry, "thumbnails", "guarded") + self.assertIs(raised.exception.running, scan) + # A second guarded kind may run beside a guarded one, not beside records. + gate.set() + await asyncio.wait_for(scan.task, 2) + hold = asyncio.Event() + thumbs = self._start(registry, "thumbnails", "guarded", hold) + backup = self._start(registry, "backup", "guarded") + await asyncio.wait_for(backup.task, 2) + self.assertEqual(backup.state, "finished") + hold.set() + await asyncio.wait_for(thumbs.task, 2) + + async def test_read_jobs_run_in_parallel(self): + registry = JobRegistry() + gate = asyncio.Event() + scan = self._start(registry, "scan", "records", gate) + health = self._start(registry, "health", "read", gate) + duplicates = self._start(registry, "duplicates", "read", gate) + self.assertTrue(all(j.running for j in (scan, health, duplicates))) + with self.assertRaises(JobBusy) as raised: # one job of a kind + self._start(registry, "health", "read") + self.assertIs(raised.exception.running, health) + gate.set() + await asyncio.wait_for( + asyncio.gather(scan.task, health.task, duplicates.task), 2 + ) + + async def test_non_cancellable_job_ignores_cancel(self): + registry = JobRegistry() + gate = asyncio.Event() + job = self._start(registry, "restore", "records", gate, cancellable=False) + self.assertIs(job.summary()["cancellable"], False) + self.assertFalse(registry.cancel(job.id)) + self.assertFalse(job.cancelled) + gate.set() + await asyncio.wait_for(job.task, 2) + self.assertEqual(job.state, "finished") + + async def test_hold_registers_a_running_non_cancellable_job(self): + registry = JobRegistry() + held = registry.hold( + "restore", "Database restore (admin page)", admission="records" + ) + async with held as job: + self.assertIs(registry.active("restore"), job) + self.assertEqual(job.state, "running") + self.assertEqual( + (job.summary()["cancellable"], job.summary()["admission"]), + (False, "records"), + ) + self.assertFalse(registry.cancel(job.id)) + with self.assertRaises(JobBusy) as raised: + self._start(registry, "scan", "records") + self.assertIn("Database restore (admin page)", str(raised.exception)) + self.assertEqual(job.state, "finished") + self.assertIsNone(registry.active("restore")) + self.assertIs(registry.list()[0], job) + + async def test_hold_refused_without_waiting(self): + registry = JobRegistry() + gate = asyncio.Event() + scan = self._start(registry, "scan", "records", gate) + with self.assertRaises(JobBusy) as raised: + registry.hold("maintenance", "Maintenance", admission="records") + self.assertIs(raised.exception.running, scan) + self.assertEqual([j.kind for j in registry.list()], ["scan"]) + gate.set() + await asyncio.wait_for(scan.task, 2) + + async def test_hold_released_on_exception(self): + registry = JobRegistry() + with self.assertRaises(RuntimeError): + async with registry.hold( + "maintenance", "Maintenance", admission="records" + ) as job: + raise RuntimeError("boom") + self.assertEqual(job.state, "failed") + again = registry.hold("maintenance", "Maintenance", admission="records") + async with again as second: + second.finish("failed") # the body may end it itself + self.assertEqual(second.state, "failed") + + def test_job_defaults_stay_read_and_cancellable(self): + job = Job("duplicates", "Duplicate scan") + self.assertEqual((job.admission, job.cancellable), ("read", True)) + + +class _LegacyHoldCase(AdminAPITestCase): + """The admin test app plus a running workspace job on demand.""" + + def _run_job(self, kind="scan", title="Scan for metadata", admission="records"): + gate = asyncio.Event() + job = self.api._jobs.start( + kind, title, admission=admission, scope={}, runner=_gated(gate) + ) + self.addAsyncCleanup(self._release, gate, job) + return gate, job + + @staticmethod + async def _release(gate, job): + gate.set() + await asyncio.wait_for(job.task, 5) + + def _db_digest(self): + path = Path(self.api.db.model.db_path) + return { + suffix: hashlib.sha256(Path(f"{path}{suffix}").read_bytes()).hexdigest() + for suffix in ("", "-wal") + if Path(f"{path}{suffix}").exists() + } + + async def _busy(self, resp, title="Scan for metadata"): + body = await resp.json() + self.assertEqual(resp.status, 409) + self.assertEqual(body, {"success": False, "error": f"{title} already running"}) + + +class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase): + + async def test_hold_is_listed_as_running_job(self): + held = self.api._jobs.hold( + "restore", "Database restore (admin page)", admission="records" + ) + async with held: + resp = await self.client.request("GET", "/prompt_manager/jobs") + jobs = (await resp.json())["jobs"] + self.assertEqual(len(jobs), 1) + job = jobs[0] + self.assertEqual( + (job["kind"], job["title"], job["state"]), + ("restore", "Database restore (admin page)", "running"), + ) + self.assertIs(job["cancellable"], False) + self.assertEqual(job["admission"], "records") + + async def test_legacy_maintenance_and_cleanup_refused_while_scan_runs(self): + self.api.db.save_prompt(text="keep me", prompt_hash="h1") + before = self._db_digest() + self._run_job() + resp = await self.client.request( + "POST", + "/prompt_manager/maintenance", + json={"operations": ["cleanup_duplicates", "vacuum"]}, + ) + await self._busy(resp) + resp = await self.client.request("POST", "/prompt_manager/cleanup") + await self._busy(resp) + self.assertEqual(self._db_digest(), before) + + async def test_legacy_maintenance_refused_while_image_tags_runs(self): + self._run_job("image-tags", "Add tags to 3 images") + resp = await self.client.request("POST", "/prompt_manager/maintenance") + await self._busy(resp, "Add tags to 3 images") + + async def test_scan_refused_while_legacy_maintenance_holds(self): + entered = threading.Event() + release = threading.Event() + real = self.api._run_maintenance_operations + + def gated(operations): + entered.set() + release.wait(10) + return real(operations) + + self.api._run_maintenance_operations = gated + loop = asyncio.get_running_loop() + maintenance = asyncio.ensure_future( + self.client.request( + "POST", "/prompt_manager/maintenance", json={"operations": ["vacuum"]} + ) + ) + self.assertTrue(await loop.run_in_executor(None, entered.wait, 10)) + resp = await self.client.request("POST", "/prompt_manager/jobs/scan") + body = await resp.json() + self.assertEqual((resp.status, body["code"]), (409, "busy")) + self.assertEqual(body["running"]["title"], "Maintenance (admin page)") + release.set() + self.assertEqual((await maintenance).status, 200) + listed = ( + await (await self.client.request("GET", "/prompt_manager/jobs")).json() + )["jobs"] + self.assertEqual( + [(j["kind"], j["state"]) for j in listed], [("maintenance", "finished")] + ) + + async def test_legacy_duplicate_delete_refused_while_scan_runs(self): + victim = self.output_dir / "dup.png" + victim.write_bytes(b"x") + self._run_job() + resp = await self.client.request( + "POST", + "/prompt_manager/delete_duplicate_images", + json={"image_paths": [str(victim)]}, + ) + await self._busy(resp) + self.assertTrue(victim.exists()) + # The duplicate family lock was not taken by the refused request. + self.assertIsNotNone(self.api._acquire_long_job(self.api.LONG_JOB_DUPLICATES)) + + async def test_lora_scan_refused_while_scan_runs(self): + from py.config import IntegrationConfig + + lm = Path(self.tmpdir) / "custom_nodes" / "ComfyUI-Lora-Manager" + (lm / "py").mkdir(parents=True) + self._run_job() + imported = [] + with ( + patch.multiple( + IntegrationConfig, LORA_MANAGER_ENABLED=True, LORA_MANAGER_PATH=str(lm) + ), + patch("py.lora_utils.resolve_lora_manager_path", return_value=str(lm)), + patch.object( + self.api, "_lora_scan_stream", side_effect=lambda *a: imported.append(a) + ), + ): + resp = await self.client.request("POST", "/prompt_manager/lora/scan") + text = await resp.text() + self.assertEqual(resp.status, 200) # the legacy page reads a stream + events = [json.loads(line[6:]) for line in text.split("\n\n") if line] + self.assertEqual( + events, + [{"type": "error", "message": "Scan for metadata already running"}], + ) + self.assertEqual(imported, []) + + async def test_legacy_autotag_refused_while_scan_runs(self): + self._run_job() + resp = await self.client.request( + "POST", "/prompt_manager/autotag/start", json={} + ) + await self._busy(resp) + + async def test_legacy_payloads_unchanged_when_idle(self): + resp = await self.client.request( + "POST", "/prompt_manager/maintenance", json={"operations": ["vacuum"]} + ) + body = await resp.json() + self.assertEqual(resp.status, 200) + self.assertEqual( + set(body), + {"success", "operations_completed", "all_successful", "results", "message"}, + ) + resp = await self.client.request("POST", "/prompt_manager/cleanup") + body = await resp.json() + self.assertEqual(resp.status, 200) + self.assertEqual(set(body), {"success", "message", "duplicates_removed"}) + resp = await self.client.request( + "POST", "/prompt_manager/delete_duplicate_images", json={"image_paths": []} + ) + self.assertEqual( + (resp.status, await resp.json()), + (400, {"success": False, "error": "No image paths provided"}), + ) + # Each hold ended with the request; a refused input shows as failed. + listed = ( + await (await self.client.request("GET", "/prompt_manager/jobs")).json() + )["jobs"] + self.assertEqual( + [(j["kind"], j["state"]) for j in listed], + [ + ("duplicate_cleanup", "failed"), + ("maintenance", "finished"), + ("maintenance", "finished"), + ], + ) + self.assertEqual( + [j["title"] for j in listed], + [ + "Duplicate image delete (gallery page)", + "Prompt cleanup (admin page)", + "Maintenance (admin page)", + ], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_jobs.py b/tests/test_jobs.py index f561be0..6848fa8 100644 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -182,8 +182,13 @@ class TestJobRegistry(unittest.IsolatedAsyncioTestCase): "scope", "listed", "log_length", + "cancellable", + "admission", }, ) + self.assertEqual( + (summary["cancellable"], summary["admission"]), (True, "records") + ) self.assertEqual(summary["scope"], {"files": "new", "output_folders": 2}) # What the detail would list, so a view reads it only when that changed. self.assertEqual((summary["listed"], summary["log_length"]), ({}, 0))