From 9ac471eb72f843eba7cdb2337601d2e3460f7424 Mon Sep 17 00:00:00 2001 From: Vito Date: Mon, 26 Oct 2026 11:46:54 -0700 Subject: [PATCH] perf(images): share and warm the Images listing (#348, PR 1) (#349) * perf(images): share and warm one listing per output root (#348) The Images page, its sidebar counts, the subfolder list and the legacy gallery each rebuilt the output listing when the 30 s cache expired, all at once, and the subfolder list computed relative paths on the event loop. Thumbnail writes also threw the listing away. Keep one RootListing per output root (py/api/gallery_listing.py): - one walk per root at a time; concurrent readers wait for it (shielded) - a listing older than 30 s is served at once and refreshed in the background - the walk, the comparison with the previous listing and both folder sets run in the executor; an unchanged walk keeps the same list object - os.scandir probe: a missing root fails the walk and keeps the previous listing (one warning per failure streak) - Rescan waits for a walk started after the click and is a 500 when it fails, even with a previous listing - deletes keep the blocking invalidation; a generation counter that survives drops discards walks that raced it - thumbnail jobs, conversion and the legacy generate/clear routes no longer drop the listing - warm-up on the host's on_startup, cancelled on on_cleanup Adds tools/vendor/perf/images_load_probe.mjs (maintainer tool, not CI) and the TESTING.md section "Images listing cache (#348)". * fix(images): drop listings on the loop when invalidated from a worker thread (#348) A settings save (gallery roots change) and the browser-test seeds call invalidate_gallery_cache() from the executor. Listing records are only touched on the event loop, so an off-loop call now hands the drop to the loop and waits for it, keeping the call synchronous for its callers. The shared-Rescan test waits for the later Rescans to take their walk number instead of sleeping. * test(images): bound the shared-listing test waits (#348) A macOS py3.12 run stopped inside the shared-listing tests and sat until the 6 h job limit with no traceback. It does not reproduce locally (about 300 runs on 3.12 and 3.14, pinned to one core). Requests and background settles in TestSharedListing now give up after 30 s, so a recurrence fails with the stuck await instead of hanging the run. * test(images): let the settle loops yield so a finished walk leaves the set (#348) On CI (ubuntu and macOS py3.12) the runner's watchdog dumped test_a_root_that_recovers_warns_again_on_its_next_failure spinning in _settle: asyncio.gather over tasks that are already done returns without yielding, so the set's discard callback, queued when the walk finished, never ran and the loop never ended (wait_for could not fire either). Delaying discard by one loop turn reproduces it locally; both settle loops now yield one turn after each gather. * test(browser): poll the import review's rows instead of reading them once gridOf reads the dialog's rows once; right after the click the review may not have rendered them yet (firefox CI: received undefined). * perf(images): walk the output with scandir entries and build folders from strings (#348) The live test showed a read during a background refresh taking 2.2 s instead of 0.07 s: the walk held the GIL for about 2 s in the server. py-spy split that walk into two halves: Path.relative_to per file in folder_sets, and Path.stat + Path.is_symlink + Path joins per file in the scan. - _walk_media visits folders in os.walk's order with os.scandir entries: the file type comes from the listing (no lstat per file) and the mtime from entry.stat (one call, none on Windows). _iter_media_files is now a thin wrapper, so its other caller is unchanged. - folder_sets slices the path strings (in the path flavour's separator) instead of calling relative_to, falling back to it for a path that does not start with the root string. On the maintainer's 19,935-file output folder: scan 0.386 s -> 0.112 s, folder sets 0.452 s -> 0.007 s, both results identical to the previous code. A reference copy of the os.walk version checks order, filters, the depth cap, symlinks and the file cap. * perf(images): build the workspace caches behind the warm-up and changed walks (#348) The live test after the faster walk still took 2.9-5.2 s for the first Images open after a restart. py-spy showed the walk was now small; the time went to the workspace's own caches, cold after a restart and built by three parallel first requests at once: the favorites listing (Path.relative_to per file), the per-file values and newest order, and the prompt-link read. - Once the host's startup warm-up has run, a stored walk that changed a listing schedules one executor job that builds what a first page reads, through the request path itself (_select_workspace_page with the default filters, then _workspace_favorite_paths). A walk stored during a build asks for one more build; an unchanged walk (same list kept) builds nothing; a failed build only warns; on_cleanup stops it. Servers that never ran the warm-up (the test servers) start no build. - _favorite_listing slices the path strings as _listed does, falling back to relative_to for a path outside the root string. * perf(images): start a stale listing's refresh walk a second after the read (#348) After the faster walk, a read that found the listing stale still took 0.5-1.0 s against 0.07 s: it was served from the old listing but ran beside the refresh walk it had just started, sharing the GIL. The refresh now starts REFRESH_DELAY (1 s) after that read through a per-root rewalk_handle, so the page's own requests finish first. One pending refresh per root and none while a walk runs; any walk start (a cold read, Rescan, the warm-up) cancels a pending one; on_cleanup cancels them before its first await. * test(browser): accept KB or MB for the picked safety backup's size The restore preview showed '2.5 MB' for the test server's database copy in a CI coverage shard: its size depends on what earlier tests on that server wrote, and the formatter switches to MB at 1 MiB. * fix(images): keep the 3.10-3.13 suffix rule in the scandir walk and harden the warm-up (#348) Review of the walk speedup: - os.path.splitext drops names like ..png that Path.suffix listed on Python 3.10-3.13 (3.14 changed Path.suffix, which hid it from the parity test on a 3.14 venv). _name_suffix spells out the 3.10-3.13 rule; the test now asserts ..png and ...webp are listed and g. is not. - _warm_workspace reads the roots inside its try, so a failing _get_all_output_dirs warns instead of leaving an unretrieved task error. - The folder_sets fallback test now takes the fallback (a Windows case mismatch); the coalescing test pins the walk count. --- TESTING.md | 47 ++ __init__.py | 5 + py/api/__init__.py | 112 ++- py/api/gallery_listing.py | 146 ++++ py/api/images.py | 305 +++++-- py/api/thumbnail_convert.py | 1 - py/api/thumbnail_jobs.py | 1 - .../browser/workspace-system-backups.spec.mjs | 3 +- .../workspace-system-integrations.spec.mjs | 3 +- tests/test_api_core.py | 18 +- tests/test_api_images.py | 766 +++++++++++++++++- tests/test_api_workspace_images.py | 38 + tests/test_gallery_listing.py | 185 +++++ tests/test_thumbnail_convert.py | 8 + tests/test_thumbnail_jobs.py | 14 + tools/vendor/perf/images_load_probe.mjs | 81 ++ 16 files changed, 1661 insertions(+), 72 deletions(-) create mode 100644 py/api/gallery_listing.py create mode 100644 tests/test_gallery_listing.py create mode 100644 tools/vendor/perf/images_load_probe.mjs diff --git a/TESTING.md b/TESTING.md index 6a0fd43..9501625 100644 --- a/TESTING.md +++ b/TESTING.md @@ -294,6 +294,53 @@ change the budget in the spec with a reason) in a follow-up PR. Never rerun it t green. It is not a required check, and it cannot be one while it skips pull requests (`.github/branch-protection.json`). +## Images listing cache (#348) + +`/prompt_manager/images/output`, `/prompt_manager/gallery/subfolders` and +`POST /prompt_manager/gallery/rescan` read one shared listing per output root +(`RootListing` in `py/api/gallery_listing.py`, `_listing` in +`py/api/images.py`). The tests are `tests/test_gallery_listing.py` (pure +helpers) and `TestSharedListing` in `tests/test_api_images.py` (routes); they +hold a walk open with a gated scan instead of sleeping. + +- **Single flight.** At most one walk per root runs at a time; every reader + that needs that root while it runs waits for it (shielded, so a client that + disconnects never cancels it). The walk, the comparison with the previous + listing and both folder sets run in the executor (`walk_root`). A walk that + finds the same `(path, mtime)` list keeps the previous list object, so the + favorites lookup and the sort orders are not rebuilt. +- **30 s background refresh.** A listing older than `REFRESH_AFTER` (30 s) is + served at once and one background walk replaces it. Files added or removed + by hand therefore show up one open later than before; Rescan shows them at + once. +- **Blocking invalidation after deletes.** Duplicate cleanup and the legacy + duplicate delete still call `invalidate_gallery_cache()`: the listing is + dropped and the next read waits for a walk, so a deleted file is never + listed. A walk running during the drop is discarded (the record's + `generation` outlives every drop); a cold reader gives up with a 500 after 3 + discarded walks in a row. Thumbnail writes (jobs, conversion, the legacy + generate and clear routes) no longer drop the listing: thumbnail URLs are + computed per request. +- **Startup warm-up.** The root `__init__.py` registers + `api.warm_gallery_listings` on the host's `on_startup`, which starts one + background walk per root, and `api.cleanup_gallery_listings` on + `on_cleanup`, which cancels running walks. A warm-up failure never stops + startup; the first read retries. +- **Failures and Rescan.** `walk_root` opens the root with `os.scandir` first, + so a deleted or unreadable root fails the walk instead of listing nothing. + A failed walk keeps the previous listing, logs one warning per streak of + failures and is retried at most once per `REFRESH_AFTER`; a reader with no + listing gets a 500. Rescan waits for a walk that started after the click + (by walk number, not by clock), Rescans clicked during one walk share the + next, and a failed walk is a 500 even when a previous listing exists. +- **Measuring it.** `tools/vendor/perf/images_load_probe.mjs` is a maintainer + tool, never run in CI. After `npm ci --prefix tools/vendor` and + `npx --prefix tools/vendor playwright install chromium`, run + `node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188` against + a running ComfyUI. It prints the time to the first decoded thumbnail, the + time to network idle and the 10 slowest requests, from a fresh browser + context each run. + ## Windows portability preflight Before changing paths, SQLite, streams, restore or cleanup: diff --git a/__init__.py b/__init__.py index e22e745..0e7d294 100644 --- a/__init__.py +++ b/__init__.py @@ -100,6 +100,11 @@ try: if host_app is not None and hasattr(host_app, "on_cleanup"): host_app.on_cleanup.append(api.cleanup_history_reads) host_app.on_cleanup.append(api.cleanup_annotation_captures) + host_app.on_cleanup.append(api.cleanup_gallery_listings) + if host_app is not None and hasattr(host_app, "on_startup"): + # Build each output root's Images listing once the server's loop runs, + # so the first open does not wait for the walk (#348). + host_app.on_startup.append(api.warm_gallery_listings) if host_app is not None and hasattr(host_app, "on_shutdown"): # Open Needs review change streams end before the server waits on them. host_app.on_shutdown.append(api.close_review_events) diff --git a/py/api/__init__.py b/py/api/__init__.py index 2084878..35efee4 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -31,6 +31,7 @@ import math import ntpath import os import re +import threading from pathlib import Path from aiohttp import web @@ -47,6 +48,7 @@ from .review_events import ReviewEventsMixin from .saved_views import SavedViewRoutesMixin from .history import HistoryRoutesMixin from .images import ImageRoutesMixin +from .gallery_listing import WALK_STOPPED, RootListing, WalkOutcome from .admin import AdminRoutesMixin from .settings_routes import SettingsRoutesMixin from .job_routes import JobRoutesMixin @@ -87,6 +89,18 @@ def _get_project_root(): return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +# How long an executor thread waits for the loop to drop the listings. +_OFF_LOOP_DROP_WAIT = 30 + + +def _running_loop(): + """The event loop running in this thread, or None.""" + try: + return asyncio.get_running_loop() + except RuntimeError: + return None + + def _public_path(path) -> str: """Render a server path for an API response without revealing the layout. @@ -344,8 +358,17 @@ class PromptManagerAPI( self._captures_closed = False self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir() self._html_cache = {} # Cached HTML file contents keyed by path - self._gallery_cache = {} # dict: path_str -> (files, timestamp) - self._gallery_cache_ttl = 30 # Cache TTL in seconds + # One shared listing per output root, keyed by str(root), and the + # background walks refreshing them (py/api/gallery_listing.py, #348). + self._root_listings: dict[str, RootListing] = {} + self._listing_tasks: set = set() + # The loop the walks run on, set by the first walk. + self._listing_loop = None + # The workspace caches rebuilt behind a stored walk that changed the + # listing, once the host's startup warm-up has run (#348). + self._workspace_warm_enabled = False + self._workspace_warm_task = None + self._workspace_warm_again = False # Workspace listing values and sort orders derived from a cached scan # (py/api/images.py _workspace_listed/_workspace_orders). self._workspace_listed_cache = {} @@ -465,20 +488,87 @@ class PromptManagerAPI( return await loop.run_in_executor(None, call) return await loop.run_in_executor(None, func, *args) - def invalidate_gallery_cache(self): - """Invalidate the cached gallery file listing. + def _root_listing(self, root) -> RootListing: + """The shared listing record of one output root (created empty).""" + return self._root_listings.setdefault(str(root), RootListing()) - Called by the image monitor when new files are detected. The workspace - keeps its per-file values and card facts, so the next request rebuilds - only new or changed files (see ImageRoutesMixin._workspace_listed and - _card_facts); its sort orders - are dropped, as the rescan replaces them anyway, and so is the - favorites lookup, which holds the old scan list. + def invalidate_gallery_cache(self): + """Drop every root's listing, so the next read waits for a new walk. + + Called after files are deleted through the app (duplicate cleanup and + the legacy duplicate delete), never for thumbnail writes. A walk + running now is not stored (its generation moved on). The workspace + keeps its per-file values and card facts, so the next request + rebuilds only new or changed files (see ImageRoutesMixin. + _workspace_listed and _card_facts); its sort orders are dropped, as + the walk replaces them anyway, and so is the favorites lookup, which + holds the old list. """ - self._gallery_cache = {} + loop = self._listing_loop + if loop is not None and loop.is_running() and _running_loop() is not loop: + # Called from an executor thread (a settings save, a test seed): + # listing records are touched only on the loop, and the caller + # expects the drop done when this returns. + dropped = threading.Event() + + def drop(): + try: + self._drop_listings() + finally: + dropped.set() + + loop.call_soon_threadsafe(drop) + if not dropped.wait(_OFF_LOOP_DROP_WAIT): + self.logger.warning("The Images listing drop is still queued") + return + self._drop_listings() + + def _drop_listings(self): + for record in list(self._root_listings.values()): + record.drop() self._workspace_order_cache = {} self._favorite_listing_cache = {} + async def warm_gallery_listings(self, app=None): + """Start one background walk per output root (aiohttp on_startup). + + Registered by the root ``__init__.py``, so the first Images open after + ComfyUI starts finds the listing built. It never fails startup: a + root that cannot be walked is retried by the first read. + """ + # Each stored walk that changes a listing then rebuilds the workspace's + # per-file values, orders and favorites behind it (_warm_workspace). + self._workspace_warm_enabled = True + try: + for root in self._get_all_output_dirs(): + self._schedule_walk(root) + except Exception as e: + self.logger.warning(f"Could not warm the Images listing: {e}") + + async def cleanup_gallery_listings(self, app=None): + """Stop the background listing walks (aiohttp on_cleanup).""" + self._workspace_warm_enabled = False + # Before any await, so no delayed refresh can start a walk meanwhile. + for record in self._root_listings.values(): + handle, record.rewalk_handle = record.rewalk_handle, None + if handle is not None: + handle.cancel() + warm, self._workspace_warm_task = self._workspace_warm_task, None + if warm is not None: + warm.cancel() + await asyncio.gather(warm, return_exceptions=True) + tasks = list(self._listing_tasks) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._listing_tasks.clear() + # A walk cancelled before its first step never ran its own cleanup. + for record in self._root_listings.values(): + flight, record.flight = record.flight, None + if flight is not None and not flight.done(): + flight.set_result(WalkOutcome.failed(RuntimeError(WALK_STOPPED))) + def reset_gallery_roots_cache(self): """Drop every cached listing and derived workspace value (roots changed).""" self.invalidate_gallery_cache() diff --git a/py/api/gallery_listing.py b/py/api/gallery_listing.py new file mode 100644 index 0000000..b873099 --- /dev/null +++ b/py/api/gallery_listing.py @@ -0,0 +1,146 @@ +"""One shared listing per output root, and the pure helpers behind it (#348). + +The API keeps one ``RootListing`` per output root (``_root_listings`` in +``py/api/__init__.py``). The record is read and written only on the event +loop; the walk itself, the comparison with the previous listing and both +folder sets run in the executor through ``walk_root``. Routes live in +``py/api/images.py`` (``_listing``); this module has no aiohttp import so +tests can exercise it without a server. +""" + +from __future__ import annotations + +import asyncio +import os +from dataclasses import dataclass +from pathlib import PurePath +from typing import Callable, NamedTuple + +# A listing older than this (seconds) is served as it is, and one background +# walk replaces it. Read as ``gallery_listing.REFRESH_AFTER`` so tests can +# patch it. +REFRESH_AFTER = 30 + +# The stale read is served at once; its refresh walk starts this many seconds +# later, so the walk does not share the GIL with the page's own requests +# (#348 live test: reads during the walk took 2.2 s instead of 0.07 s). +REFRESH_DELAY = 1.0 + +# A reader with no listing gives up (500) after this many walks in a row were +# discarded by invalidations made while they ran. +MAX_DISCARDED_WALKS = 3 + +WALK_STOPPED = "The listing walk was stopped" + +STORED = "stored" +DISCARDED = "discarded" +FAILED = "failed" + + +class WalkOutcome(NamedTuple): + """What one walk ended with; a walk's flight future resolves to this.""" + + kind: str + error: BaseException | None = None + + @classmethod + def stored(cls) -> "WalkOutcome": + return cls(STORED) + + @classmethod + def discarded(cls) -> "WalkOutcome": + return cls(DISCARDED) + + @classmethod + def failed(cls, error: BaseException) -> "WalkOutcome": + return cls(FAILED, error) + + +@dataclass +class RootListing: + """The shared listing of one output root. + + A record is never deleted: ``drop()`` clears the listing and bumps + ``generation``, so a walk that started before a drop can never be stored + after it. + """ + + # (path, mtime) newest first, as _scan_gallery_files_sync returns it. + files: list | None = None + folders: frozenset | None = None + folders_with_ancestors: frozenset | None = None + # Monotonic time the last walk finished, stored or failed. + walked_at: float | None = None + # Number of the last walk started; a running walk is always that one. + walk_seq: int = 0 + # Number of the walk whose files are stored (Rescan waits on it). + stored_seq: int = 0 + generation: int = 0 + # The running walk's future (resolves to a WalkOutcome), or None. + flight: asyncio.Future | None = None + # True while walks keep failing: one warning per streak. + failing: bool = False + # The delayed refresh not started yet (loop.call_later), or None. + rewalk_handle: asyncio.TimerHandle | None = None + + def drop(self) -> None: + self.files = None + self.folders = None + self.folders_with_ancestors = None + self.walked_at = None + self.generation += 1 + + +def folder_sets(root, files) -> tuple[frozenset, frozenset]: + """``(folders, folders_with_ancestors)`` of the files listed under *root*. + + Folders keep the spelling ``/gallery/subfolders`` has always answered: + ``str(relative.parent)`` (native separators) for the folders holding + files, plus every ancestor joined by ``/``. + """ + # From the path strings: relative_to per file was half of a large walk. + # The separator is the path flavour's own (tests use PureWindowsPath). + sep = str(type(root)("a", "b"))[1] + prefix = str(root) + if not prefix.endswith(sep): + prefix += sep + folders = set() + for path, _mtime in files: + text = str(path) + if text.startswith(prefix): + rel_dir = text[len(prefix) :].rpartition(sep)[0] + else: + rel_dir = str(path.relative_to(root).parent) + if rel_dir and rel_dir != ".": + folders.add(rel_dir) + with_ancestors = set(folders) + for folder in folders: + parts = folder.replace("\\", "/").split("/") + for i in range(1, len(parts)): + with_ancestors.add("/".join(parts[:i])) + return frozenset(folders), frozenset(with_ancestors) + + +def same_listing(old, new) -> bool: + """True when two listings hold the same ``(path, mtime)`` pairs in order.""" + if old is None or new is None: + return False + return old is new or old == new + + +def walk_root( + scan_fn: Callable[[PurePath], list], root, previous_files +) -> tuple[list, frozenset, frozenset]: + """The single executor job of a walk: ``(files, folders, with_ancestors)``. + + Opens the root first, so a missing or unreadable root raises ``OSError`` + instead of listing nothing. Returns *previous_files* itself when nothing + changed, which keeps the identity-keyed caches built on it warm. + """ + with os.scandir(root): + pass + files = scan_fn(root) + if same_listing(previous_files, files): + files = previous_files + folders, with_ancestors = folder_sets(root, files) + return files, folders, with_ancestors diff --git a/py/api/images.py b/py/api/images.py index 0d2691a..763de13 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -22,6 +22,15 @@ from aiohttp import web from PIL import Image, UnidentifiedImageError from .jobs import MAX_OUTCOME_ROWS, THUMBNAILS_JOB_KIND +from . import gallery_listing +from .gallery_listing import ( + DISCARDED, + FAILED, + MAX_DISCARDED_WALKS, + WALK_STOPPED, + WalkOutcome, + walk_root, +) from . import thumbnail_settings from .thumbnail_encode import ( PREVIEW_HELD, @@ -763,32 +772,65 @@ def _favorite_root_key(root) -> str: return _favorite_relative_key(os.path.realpath(str(root))) -def _iter_media_files(output_path, max_files): - """Yield media files under *output_path* with bounded effort. +def _name_suffix(name): + """``Path(name).suffix`` as Python 3.10-3.13 compute it. + + The last ``.``-part unless the dot is the first or last character, so + ``..png`` is ``.png`` and ``.png`` has none. Spelled out because 3.14 + changed ``Path.suffix`` and ``os.path.splitext`` differs for leading dots. + """ + dot = name.rfind(".") + return name[dot:] if 0 < dot < len(name) - 1 else "" + + +def _walk_media(output_path, max_files): + """Yield the ``os.DirEntry`` of each media file under *output_path*. Symlinks are never followed, ``thumbnails/`` is skipped, descent stops - at MAX_SCAN_DEPTH and at most *max_files* paths are produced. + at MAX_SCAN_DEPTH and at most *max_files* entries are produced. The order + is ``os.walk``'s (top-down, subfolders in ``os.scandir`` order, files by + name), which the gallery's ties in mtime keep. Each entry's type comes + from the directory listing itself, so a walk builds no ``Path`` and makes + no extra ``lstat`` per file (#348: a 20,000-file walk spent half its time + there). """ - output_path = Path(output_path) - root_depth = len(output_path.parts) + stack = [(os.fspath(output_path), ())] produced = 0 - for root, dirs, files in os.walk(output_path): - root_path = Path(root) - rel_parts = root_path.parts[root_depth:] + while stack: + top, rel_parts = stack.pop() if "thumbnails" in rel_parts: - dirs[:] = [] continue - if len(rel_parts) >= MAX_SCAN_DEPTH: - dirs[:] = [] - for name in sorted(files): - if Path(name).suffix.lower() not in MEDIA_EXTENSIONS: + dirs, files = [], [] + try: + with os.scandir(top) as entries: + for entry in entries: + try: + is_dir = entry.is_dir() + except OSError: + is_dir = False + (dirs if is_dir else files).append(entry) + except OSError: + continue # as os.walk: an unreadable folder is skipped + for entry in sorted(files, key=lambda item: item.name): + if _name_suffix(entry.name).lower() not in MEDIA_EXTENSIONS: continue - if (root_path / name).is_symlink(): + if entry.is_symlink(): continue # never read through a link out of the output tree if produced >= max_files: return produced += 1 - yield root_path / name + yield entry + if len(rel_parts) >= MAX_SCAN_DEPTH: + continue + for entry in reversed(dirs): + if not entry.is_symlink(): + stack.append((entry.path, rel_parts + (entry.name,))) + + +def _iter_media_files(output_path, max_files): + """Yield the ``Path`` of each media file under *output_path* (see _walk_media).""" + for entry in _walk_media(output_path, max_files): + yield Path(entry.path) def media_thumbnail_rel(rel_path, settings=None): @@ -1278,6 +1320,60 @@ class ImageRoutesMixin: {"success": False, "error": _safe_error(e)}, status=500 ) + def _schedule_workspace_warm(self): + """Rebuild the workspace caches behind a changed listing (on the loop). + + Only after the host's startup warm-up (warm_gallery_listings), so the + first Images open after a restart, and the first read after a refresh + found new files, pay no per-file rebuild. A walk stored while a build + runs asks for one more build instead of a second one at once. + """ + if not self._workspace_warm_enabled: + return + task = self._workspace_warm_task + if task is not None and not task.done(): + self._workspace_warm_again = True + return + self._workspace_warm_task = asyncio.get_running_loop().create_task( + self._warm_workspace() + ) + + async def _warm_workspace(self): + while True: + self._workspace_warm_again = False + try: + listings = self._stored_listings() + if listings is None: + return # a root has no listing yet; its walk asks again + await self._run_in_executor(self._build_workspace_caches, listings) + except Exception as e: + self.logger.warning(f"Could not prepare the Images page: {e}") + return + if not self._workspace_warm_again: + return + + def _stored_listings(self): + """``[(index, root, files)]`` of every output root, or None if one has none.""" + listings = [] + for index, root in enumerate(self._get_all_output_dirs()): + record = self._root_listings.get(str(root)) + if record is None or record.files is None: + return None + listings.append((index, root, record.files)) + return listings + + def _build_workspace_caches(self, listings): + """Build what a first workspace page reads, as that page builds it (blocking). + + The default page's selection builds the per-file values, the newest + order, the wanted paths and reads the prompt links (warming SQLite's + cache); the favorites filter builds each root's favorite listing. + """ + self._select_workspace_page( + listings, _workspace_filters({}), None, 1, "workspace-warm-up" + ) + self._workspace_favorite_paths(listings) + def _scan_gallery_files_sync(self, output_path): """Scan output directory for media files (blocking I/O, run in executor). @@ -1285,27 +1381,141 @@ class ImageRoutesMixin: walk is bounded by MAX_SCAN_DEPTH and MAX_GALLERY_FILES. """ found = [] - for media_path in _iter_media_files(output_path, MAX_GALLERY_FILES): + for entry in _walk_media(output_path, MAX_GALLERY_FILES): try: - found.append((media_path, media_path.stat().st_mtime)) + # Never a symlink (_walk_media skips them), so this is the + # file's own mtime, read from the listing where it can be. + mtime = entry.stat(follow_symlinks=False).st_mtime except OSError: continue + found.append((Path(entry.path), mtime)) if len(found) >= MAX_GALLERY_FILES: self.logger.warning(f"Gallery scan capped at {MAX_GALLERY_FILES} files") found.sort(key=lambda x: x[1], reverse=True) return found async def _get_gallery_files(self, output_path): - """Get gallery files with per-directory TTL cache.""" - now = _time.monotonic() - key = str(output_path) - cached = self._gallery_cache.get(key) - if cached and (now - cached[1]) < self._gallery_cache_ttl: - return cached[0] + """The shared listing of one output root: ``[(path, mtime)]``, newest first.""" + return (await self._listing(output_path)).files - files = await self._run_in_executor(self._scan_gallery_files_sync, output_path) - self._gallery_cache[key] = (files, now) - return files + async def _listing(self, output_path, *, fresh=False): + """The shared ``RootListing`` of *output_path*, with a listing (#348). + + A listing that exists is returned at once; when it is older than + ``REFRESH_AFTER`` one background walk replaces it. With no listing, + the caller waits for the root's single walk. ``fresh`` (Rescan) waits + for a walk that started after this call, and raises when it fails + even if a previous listing exists. + """ + record = self._root_listing(output_path) + needed = record.walk_seq + 1 if fresh else 0 + discarded = 0 + while True: + if record.files is not None and record.stored_seq >= needed: + if not fresh: + self._refresh_if_stale(output_path, record) + return record + flight = record.flight + if flight is None: + flight = self._start_walk(output_path, record) + joined = record.walk_seq + # Shielded: a client that goes away never cancels a shared walk. + outcome = await asyncio.shield(flight) + if joined < needed: + continue # an older walk ended; Rescan needs a newer one + if outcome.kind == FAILED: + raise outcome.error + if outcome.kind == DISCARDED: + discarded += 1 + if discarded >= MAX_DISCARDED_WALKS: + raise RuntimeError( + "The output listing kept changing while it was read" + ) + + def _refresh_if_stale(self, output_path, record): + """Walk a stale listing again, REFRESH_DELAY after this read (on the loop). + + At most one pending refresh per root, none while a walk runs. + """ + age = _time.monotonic() - (record.walked_at or 0.0) + if age < gallery_listing.REFRESH_AFTER: + return + if record.flight is not None or record.rewalk_handle is not None: + return + record.rewalk_handle = asyncio.get_running_loop().call_later( + gallery_listing.REFRESH_DELAY, self._delayed_walk, output_path, record + ) + + def _delayed_walk(self, output_path, record): + record.rewalk_handle = None + self._schedule_walk(output_path) + + def _schedule_walk(self, output_path): + """Start a background walk of *output_path* unless one is running.""" + record = self._root_listing(output_path) + if record.flight is None: + self._start_walk(output_path, record) + + def _start_walk(self, output_path, record): + """Start the root's one walk; returns its flight future (on the loop).""" + loop = asyncio.get_running_loop() + self._listing_loop = loop + handle, record.rewalk_handle = record.rewalk_handle, None + if handle is not None: + handle.cancel() # this walk is the refresh it was waiting to start + record.walk_seq += 1 + flight = loop.create_future() + record.flight = flight + task = loop.create_task( + self._run_walk( + output_path, record, flight, record.generation, record.walk_seq + ) + ) + self._listing_tasks.add(task) + task.add_done_callback(self._listing_tasks.discard) + return flight + + async def _run_walk(self, output_path, record, flight, generation, seq): + """Walk in the executor, then store, discard or fail (on the loop). + + The flight is always resolved with a WalkOutcome, never failed, so an + error never reaches a future nobody awaits. + """ + outcome = None + try: + files, folders, with_ancestors = await self._run_in_executor( + walk_root, self._scan_gallery_files_sync, output_path, record.files + ) + except asyncio.CancelledError: + outcome = WalkOutcome.failed(RuntimeError(WALK_STOPPED)) + raise + except Exception as exc: + record.walked_at = _time.monotonic() + if not record.failing: + self.logger.warning( + f"Could not list output folder {output_path}: {exc}" + ) + record.failing = True + outcome = WalkOutcome.failed(exc) + else: + record.failing = False + if record.generation == generation: + changed = files is not record.files + record.files = files + record.folders = folders + record.folders_with_ancestors = with_ancestors + record.walked_at = _time.monotonic() + record.stored_seq = seq + outcome = WalkOutcome.stored() + if changed: + self._schedule_workspace_warm() + else: + outcome = WalkOutcome.discarded() + finally: + if record.flight is flight: + record.flight = None + if not flight.done(): + flight.set_result(outcome) async def get_output_images(self, request): """Get all images from ComfyUI output folder(s).""" @@ -1832,17 +2042,18 @@ class ImageRoutesMixin: ) async def rescan_gallery(self, request): - """Forget the cached listing and rewalk every output root now. + """Rewalk every output root now and answer with the new listing. - The listing is cached for ``_gallery_cache_ttl`` seconds; the gallery's - Rescan button needs files added or removed on disk to show up at once. + Readers are served the shared listing and refresh it in the + background (``gallery_listing.REFRESH_AFTER``); the gallery's Rescan + button needs files added or removed on disk to show up at once, so it + waits for a walk that started after the click and fails when it does. """ try: - self.invalidate_gallery_cache() output_dirs = list(self._get_all_output_dirs()) total = 0 for output_path in output_dirs: - total += len(await self._get_gallery_files(output_path)) + total += len((await self._listing(output_path, fresh=True)).files) return web.json_response( {"success": True, "total": total, "roots": len(output_dirs)} ) @@ -1853,20 +2064,14 @@ class ImageRoutesMixin: ) async def _root_subfolders(self, output_path, include_ancestors): - """The distinct folders holding listed files under one output root.""" - subfolders = set() - for f, _mtime in await self._get_gallery_files(output_path): - rel_dir = str(f.relative_to(output_path).parent) - if rel_dir and rel_dir != ".": - subfolders.add(rel_dir) + """The distinct folders holding listed files under one output root. + + Built with the walk, in the executor (gallery_listing.folder_sets). + """ + listing = await self._listing(output_path) if include_ancestors: - ancestors = set() - for folder in subfolders: - parts = folder.replace("\\", "/").split("/") - for i in range(1, len(parts)): - ancestors.add("/".join(parts[:i])) - subfolders.update(ancestors) - return subfolders + return listing.folders_with_ancestors + return listing.folders async def get_gallery_subfolders(self, request): """Get distinct subfolders from gallery output directories. @@ -1940,7 +2145,6 @@ class ImageRoutesMixin: result = await self._run_in_executor( self._generate_gallery_thumbnails_sync, output_dirs ) - self.invalidate_gallery_cache() return web.json_response(result) except Exception as e: @@ -2416,7 +2620,6 @@ class ImageRoutesMixin: cleared_count, cleared_size = await self._run_in_executor( self._clear_gallery_thumbnails_sync, output_dirs ) - self.invalidate_gallery_cache() def format_size(bytes_size: int) -> str: size = float(bytes_size) @@ -2764,9 +2967,17 @@ class ImageRoutesMixin: cached = self._favorite_listing_cache.get(cache_key) if cached is not None and cached[0] is files: return cached[1] + # From the path strings, as _listed does: relative_to per file was the + # largest cost of a cold first open (#348). + prefix = os.path.join(str(root), "") listing = {} for path, _mtime in files: - posix = path.relative_to(root).as_posix() + text = str(path) + if text.startswith(prefix): + relative = text[len(prefix) :] + posix = relative if os.sep == "/" else relative.replace(os.sep, "/") + else: + posix = path.relative_to(root).as_posix() listing.setdefault(_favorite_relative_key(posix), posix) self._favorite_listing_cache[cache_key] = (files, listing) return listing diff --git a/py/api/thumbnail_convert.py b/py/api/thumbnail_convert.py index 2ef961d..3a9db4e 100644 --- a/py/api/thumbnail_convert.py +++ b/py/api/thumbnail_convert.py @@ -175,7 +175,6 @@ class ThumbnailConvertMixin: job.log_line(f"Conversion finished · {NOTHING_OUTDATED.lower()}") return "finished" totals, left = await self._convert_in_batches(job, options, plan, output_dirs) - self.invalidate_gallery_cache() if totals["stopped"]: job.update(status_line=FORMAT_CHANGED.format(converted=totals["converted"])) job.log_line( diff --git a/py/api/thumbnail_jobs.py b/py/api/thumbnail_jobs.py index fadf154..6477f7d 100644 --- a/py/api/thumbnail_jobs.py +++ b/py/api/thumbnail_jobs.py @@ -261,7 +261,6 @@ class ThumbnailJobsMixin: stats, start, left, attempted = await self._generate_in_batches( job, options, plan, roots, output_dirs ) - self.invalidate_gallery_cache() generated = _plural(stats["generated"], "preview") if job.cancelled: return self._end_cancelled( diff --git a/tests/browser/workspace-system-backups.spec.mjs b/tests/browser/workspace-system-backups.spec.mjs index fb4bdae..6856490 100644 --- a/tests/browser/workspace-system-backups.spec.mjs +++ b/tests/browser/workspace-system-backups.spec.mjs @@ -485,5 +485,6 @@ test('a safety backup file (no .db extension) can be picked and reviewed (M3)', const answer = await chooseBackup(page, bytes, 'prompts.db.backup_20261025_093000_000000'); expect((await answer).status()).toBe(200); await expect(dialog(page).locator('#modal-title')).toHaveText('Restore preview'); - expect((await gridOf(page))[0][1]).toMatch(/^prompts\.db\.backup_20261025_093000_000000 · \d+ KB · .+$/); + // The copy's size depends on what earlier tests on this server wrote: KB or MB. + expect((await gridOf(page))[0][1]).toMatch(/^prompts\.db\.backup_20261025_093000_000000 · [\d,.]+ (KB|MB) · .+$/); }); diff --git a/tests/browser/workspace-system-integrations.spec.mjs b/tests/browser/workspace-system-integrations.spec.mjs index 47358ef..f8ee02c 100644 --- a/tests/browser/workspace-system-integrations.spec.mjs +++ b/tests/browser/workspace-system-integrations.spec.mjs @@ -381,7 +381,8 @@ test('saved pairs block the import: the review explains it and offers only Close await openIntegrations(page); await enable(page); await layoutOf(page).getByRole('button', { name: 'Review import', exact: true }).click(); - expect((await gridOf(page))[2]).toEqual(['Protected', '1 imported prompt used by saved pairs']); + // gridOf reads once; poll until the review has rendered its rows. + await expect.poll(async () => (await gridOf(page))[2]).toEqual(['Protected', '1 imported prompt used by saved pairs']); await expect(dialog(page).locator('p.small')).toHaveText( 'Delete or change the saved pairs that use these prompts, then review the import again. Nothing was changed.'); expect(await dialog(page).locator('.dialog-actions button').allTextContents()).toEqual(['Close']); diff --git a/tests/test_api_core.py b/tests/test_api_core.py index b48e73c..e6e06be 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -687,9 +687,23 @@ class TestInstanceHelpers(FolderPathsFixture): self.assertEqual(self._run(self.api._run_in_executor(add, 1, 2)), 3) def test_invalidate_gallery_cache(self): - self.api._gallery_cache = {"x": 1} + record = self.api._root_listing(self.output_dir) + record.files = [(self.output_dir / "a.png", 1.0)] + record.folders = record.folders_with_ancestors = frozenset() + record.walked_at = 1.0 + self.api._favorite_listing_cache = {"x": 1} + self.api._workspace_order_cache = {"x": 1} self.api.invalidate_gallery_cache() - self.assertEqual(self.api._gallery_cache, {}) + # The record stays: its generation outlives every drop (#348). + self.assertIs(self.api._root_listings[str(self.output_dir)], record) + self.assertIsNone(record.files) + self.assertIsNone(record.folders) + self.assertIsNone(record.walked_at) + self.assertEqual(record.generation, 1) + self.assertEqual(self.api._favorite_listing_cache, {}) + self.assertEqual(self.api._workspace_order_cache, {}) + self.api.reset_gallery_roots_cache() + self.assertEqual(record.generation, 2) def test_enrich_images_adds_urls_and_thumbnails(self): GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)] diff --git a/tests/test_api_images.py b/tests/test_api_images.py index e3f3ab6..4f8cfc6 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -14,7 +14,7 @@ import tempfile import threading import unittest from contextlib import ExitStack -from pathlib import Path +from pathlib import Path, PurePath from unittest.mock import AsyncMock, MagicMock, patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -34,6 +34,7 @@ from PIL import Image # noqa: E402 from database.operations import PromptDatabase # noqa: E402 import py.api.images as images_module +from py.api import gallery_listing # noqa: E402 from py.api.thumbnail_settings import ThumbnailSettings # noqa: E402 from py.api import PromptManagerAPI # noqa: E402 from py.api.prompts import MAX_PAGE_LIMIT # noqa: E402 @@ -1454,12 +1455,15 @@ class TestGalleryListing(ImageRouteCoverageCase): self.assertEqual(data["images"], []) def test_gallery_scan_skips_unreadable_entries(self): - ghost = self.output_dir / "ghost.png" + make_png(self.output_dir / "real.png") + listed = list(images_module._walk_media(self.output_dir, 10)) + ghost = MagicMock(path=str(self.output_dir / "ghost.png")) + ghost.stat.side_effect = FileNotFoundError("gone") with patch.object( - images_module, "_iter_media_files", return_value=iter([ghost]) + images_module, "_walk_media", return_value=iter([ghost, *listed]) ): found = self.api._scan_gallery_files_sync(self.output_dir) - self.assertEqual(found, []) + self.assertEqual([path.name for path, _mtime in found], ["real.png"]) class TestThumbnailUnits(ImageRouteCoverageCase): @@ -1805,10 +1809,13 @@ class TestClearThumbnails(ImageRouteCoverageCase): self.assertFalse(data["success"]) self.assertTrue(original.is_file()) - async def test_cleanup_runs_off_loop_and_invalidates_gallery_cache(self): + async def test_cleanup_runs_off_loop_and_keeps_the_listing(self): import threading make_png(self.output_dir / "thumbnails" / "a_thumb.png") + make_png(self.output_dir / "a.png") + await self._json("GET", "/prompt_manager/images/output") + listing = self.api._root_listings[str(self.output_dir)].files loop_thread = threading.get_ident() deletion_threads = [] original_unlink = Path.unlink @@ -1817,14 +1824,14 @@ class TestClearThumbnails(ImageRouteCoverageCase): deletion_threads.append(threading.get_ident()) return original_unlink(path, *args, **kwargs) - self.api._gallery_cache = {"previous": object()} with patch.object(Path, "unlink", unlink): status, data = await self._json("POST", self.URL) self.assertEqual(status, 200) self.assertEqual(data["cleared_files"], 1) self.assertTrue(deletion_threads) self.assertNotIn(loop_thread, deletion_threads) - self.assertEqual(self.api._gallery_cache, {}) + # Previews are not part of the listing (#348): it stays as it was. + self.assertIs(self.api._root_listings[str(self.output_dir)].files, listing) class TestImagePromptLookup(ImageRouteCoverageCase): @@ -1979,7 +1986,10 @@ class TestGalleryRescan(ImageAPITestCase): return resp.status, await resp.json() async def test_rescan_makes_new_files_visible_before_the_cache_expires(self): - self.api._gallery_cache_ttl = 3600 # the cache alone would hide the new file + # The listing alone would hide the new file for an hour. + refresh = patch.object(gallery_listing, "REFRESH_AFTER", 3600) + refresh.start() + self.addCleanup(refresh.stop) make_png(self.output_dir / "a.png") _, first = await self._json("GET", self.URL) self.assertEqual(first["total"], 1) @@ -2026,6 +2036,78 @@ class TestOutputScanHelpers(ImageAPITestCase): expected = hashlib.sha1(str(media).encode("utf-8")).hexdigest()[:16] self.assertEqual(entry["id"], expected) + def test_media_walk_matches_the_os_walk_reference(self): + # The scandir walk (#348) must list exactly what the os.walk version + # listed, in the same order: ties in mtime keep this order. + output = Path(self.output_dir) + for rel in ( + "z.png", + "a.PNG", + "notes.txt", + "b/2.png", + "b/1.webp", + "b/c/deep.mp4", + "b/c/d/too_deep.png", + "thumbnails/t.png", + "b/thumbnails/u.png", + "e/.hidden.png", + "e/f.jpeg", + "e/..png", + "e/...webp", + "e/g.", + ): + make_png(output / rel) + try: + os.symlink(output / "z.png", output / "b" / "link.png") + os.symlink(output / "b", output / "dirlink", target_is_directory=True) + except (OSError, NotImplementedError): + pass # the rest of the tree still checks order and filters + with patch.object(images_module, "MAX_SCAN_DEPTH", 2): + listed = list(images_module._iter_media_files(output, 100)) + expected = list(_reference_iter_media_files(output, 100)) + capped = list(images_module._iter_media_files(output, 3)) + self.assertEqual(listed, expected) + self.assertEqual(capped, expected[:3]) + # Path.suffix of Python 3.10-3.13, which CI and hosts run (3.14 changed + # it): "..png" has the suffix ".png", ".png" and "g." have none. + names = {path.name for path in listed} + self.assertTrue({"..png", "...webp", ".hidden.png"} <= names, names) + self.assertFalse({"g.", "notes.txt"} & names, names) + self.assertTrue(all(isinstance(path, Path) for path in listed)) + + def test_gallery_scan_reads_each_file_type_and_mtime_from_scandir(self): + # Path.stat and Path.is_symlink per file were the other half of the + # live walk's cost (#348); the scandir entries already carry both. + output = Path(self.output_dir) + make_png(output / "a.png") + make_png(output / "sub" / "b.png") + with ( + patch.object(Path, "stat", side_effect=AssertionError), + patch.object(Path, "is_symlink", side_effect=AssertionError), + ): + found = self.api._scan_gallery_files_sync(output) + self.assertEqual( + sorted(path.name for path, _mtime in found), ["a.png", "b.png"] + ) + for path, mtime in found: + self.assertEqual(mtime, os.stat(path).st_mtime) + + def test_favorite_listing_slices_the_path_strings(self): + # relative_to per file was the largest cold cost of the first open. + output = Path(self.output_dir) + files = [ + (output / "a.png", 1.0), + (output / "Sub" / "B.png", 2.0), + (output / "sub" / "deep" / "c.webp", 3.0), + ] + expected = {} + for path, _mtime in files: + posix = path.relative_to(output).as_posix() + expected.setdefault(images_module._favorite_relative_key(posix), posix) + with patch.object(PurePath, "relative_to", side_effect=AssertionError): + listing = self.api._favorite_listing(output, files) + self.assertEqual(listing, expected) + def test_iter_media_files_skips_symlinked_files(self): output = Path(self.output_dir) make_png(output / "real.png") @@ -2342,3 +2424,671 @@ class TestRootScopeServing(ImageAPITestCase): self._make_symlink(roots[1], link) for path in ("image.png", "thumbnails/image_thumb.png"): await self.assert_no_access(path, scope=token) + + +# Seconds any single wait in the shared-listing tests may take. +_HANG_LIMIT = 30 + + +def _reference_iter_media_files(output_path, max_files): + """``_iter_media_files`` as it was before the scandir walk (#348).""" + output_path = Path(output_path) + root_depth = len(output_path.parts) + produced = 0 + for root, dirs, files in os.walk(output_path): + root_path = Path(root) + rel_parts = root_path.parts[root_depth:] + if "thumbnails" in rel_parts: + dirs[:] = [] + continue + if len(rel_parts) >= images_module.MAX_SCAN_DEPTH: + dirs[:] = [] + for name in sorted(files): + # Path.suffix as Python 3.10-3.13 computed it (3.14 changed it). + dot = name.rfind(".") + suffix = name[dot:] if 0 < dot < len(name) - 1 else "" + if suffix.lower() not in images_module.MEDIA_EXTENSIONS: + continue + if (root_path / name).is_symlink(): + continue + if produced >= max_files: + return + produced += 1 + yield root_path / name + + +class _GatedScan: + """Wraps the real scan, counts calls and can hold one call after it listed. + + ``hold_next()`` makes the next call list the files, then wait on the + returned gate, so a test can change the disk or the cache while a walk + is in flight and decide when it ends. + """ + + def __init__(self, real): + self.real = real + self.calls = 0 + self._gate = None + self.entered = threading.Event() + self._gates = [] + + def hold_next(self): + self._gate = threading.Event() + self._gates.append(self._gate) + self.entered.clear() + return self._gate + + def release_all(self): + for gate in self._gates: + gate.set() + + def __call__(self, root): + self.calls += 1 + files = self.real(root) + gate, self._gate = self._gate, None + if gate is not None: + self.entered.set() + gate.wait(30) + return files + + async def wait_entered(self): + loop = asyncio.get_running_loop() + entered = await loop.run_in_executor(None, self.entered.wait, 30) + if not entered: + raise AssertionError("the walk never reached the gate") + + +class TestSharedListing(ImageAPITestCase): + """One listing per output root, shared and refreshed in the background (#348).""" + + URL = "/prompt_manager/images/output" + + async def asyncSetUp(self): + await super().asyncSetUp() + self.scan = _GatedScan(self.api._scan_gallery_files_sync) + self.api._scan_gallery_files_sync = self.scan + self.addCleanup(self.scan.release_all) + # The stale refresh waits REFRESH_DELAY after the read (#348); these + # tests start it on the next loop turn and _settle waits for it. + delay = patch.object(gallery_listing, "REFRESH_DELAY", 0) + delay.start() + self.addCleanup(delay.stop) + + async def _json(self, method, path, **kwargs): + # Bounded, like every wait in this class: a stuck walk fails the test + # with a traceback instead of hanging the whole run. + async def call(): + resp = await self.client.request(method, path, **kwargs) + return resp.status, await resp.json() + + return await asyncio.wait_for(call(), _HANG_LIMIT) + + def _record(self): + return self.api._root_listings[str(self.output_dir)] + + def _pending(self): + return any( + record.rewalk_handle is not None + for record in self.api._root_listings.values() + ) + + def _warming(self): + task = self.api._workspace_warm_task + return task is not None and not task.done() + + def _age(self, seconds=None): + from py.api import gallery_listing + + record = self._record() + record.walked_at -= (seconds or gallery_listing.REFRESH_AFTER) + 1 + + async def _settle(self): + """Wait for every background walk this API started.""" + + async def settle(): + while self.api._listing_tasks or self._warming() or self._pending(): + if self._pending(): + await asyncio.sleep(0.01) # a delayed refresh not started yet + await asyncio.gather(*list(self.api._listing_tasks)) + # A finished walk leaves the set in a done callback, which needs + # a loop turn; gather over finished tasks does not yield one. + await asyncio.sleep(0) + if self._warming(): + await asyncio.wait([self.api._workspace_warm_task]) + + await asyncio.wait_for(settle(), _HANG_LIMIT) + + async def test_concurrent_cold_readers_share_one_walk(self): + make_png(self.output_dir / "sub" / "a.png") + results = await asyncio.gather( + self._json("GET", self.URL), + self._json("GET", self.URL + "?limit=1"), + self._json("GET", self.URL + "?mode=workspace&limit=1"), + self._json("GET", "/prompt_manager/gallery/subfolders"), + ) + self.assertEqual([status for status, _ in results], [200] * 4) + self.assertEqual(self.scan.calls, 1) + self.assertEqual(results[0][1]["total"], 1) + self.assertEqual(results[3][1]["subfolders"], ["sub"]) + + async def test_a_stale_listing_is_served_at_once_and_refreshed_behind(self): + make_png(self.output_dir / "a.png") + _, first = await self._json("GET", self.URL) + make_png(self.output_dir / "b.png") + self._age() + gate = self.scan.hold_next() + _, stale = await self._json("GET", self.URL) + # Served from the old listing while the walk is still held. + self.assertEqual((first["total"], stale["total"]), (1, 1)) + await self.scan.wait_entered() + # A second stale read joins the running walk instead of starting one. + _, again = await self._json("GET", self.URL) + self.assertEqual(again["total"], 1) + gate.set() + await self._settle() + _, fresh = await self._json("GET", self.URL) + self.assertEqual(fresh["total"], 2) + self.assertEqual(self.scan.calls, 2) + + async def test_a_fresh_listing_is_not_walked_again(self): + make_png(self.output_dir / "a.png") + for _ in range(3): + await self._json("GET", self.URL) + await self._json("GET", "/prompt_manager/gallery/subfolders") + self.assertEqual(self.scan.calls, 1) + self.assertEqual(self.api._listing_tasks, set()) + + async def test_rescan_waits_for_a_walk_started_after_the_click(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + self._age() + gate = self.scan.hold_next() + await self._json("GET", self.URL) # starts a background walk + await self.scan.wait_entered() # it has listed: it cannot see c.png + make_png(self.output_dir / "c.png") + rescan = asyncio.ensure_future( + self._json("POST", "/prompt_manager/gallery/rescan") + ) + await asyncio.sleep(0) + self.assertFalse(rescan.done()) + gate.set() + status, data = await rescan + self.assertEqual(status, 200) + self.assertEqual((data["total"], data["roots"]), (2, 1)) + self.assertEqual(self.scan.calls, 3) + + async def test_rescans_clicked_during_a_walk_share_the_next_one(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + original = self.api._listing + rescans = [] + clicked = asyncio.Event() + + async def listing(root, *, fresh=False): + if fresh: + rescans.append(root) + if len(rescans) == 3: + clicked.set() + # Runs on in this step up to its first wait, so the Rescan's + # walk number is taken before the test resumes. + return await original(root, fresh=fresh) + + self.api._listing = listing + gate = self.scan.hold_next() + first = asyncio.ensure_future( + self._json("POST", "/prompt_manager/gallery/rescan") + ) + await self.scan.wait_entered() + # Clicked while the first Rescan's walk runs: neither may use it, + # and both wait for the same next walk. + later = [ + asyncio.ensure_future(self._json("POST", "/prompt_manager/gallery/rescan")) + for _ in range(2) + ] + await asyncio.wait_for(clicked.wait(), 30) + gate.set() + results = await asyncio.gather(first, *later) + self.assertEqual([status for status, _ in results], [200] * 3) + self.assertEqual(self.scan.calls, 3) + + async def test_rescan_failure_with_a_previous_listing_is_500(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + self.api._scan_gallery_files_sync = _raise + with self.assertLogs(self.api.logger, "WARNING"): + status, data = await self._json("POST", "/prompt_manager/gallery/rescan") + self.assertEqual(status, 500) + self.assertFalse(data["success"]) + self.assertNotIn("boom", json.dumps(data)) + # Plain readers keep the previous listing, and do not retry at once. + status, listed = await self._json("GET", self.URL) + self.assertEqual((status, listed["total"]), (200, 1)) + self.assertEqual(self.api._listing_tasks, set()) + + async def test_a_cold_read_whose_walk_fails_is_500(self): + self.api._scan_gallery_files_sync = _raise + with self.assertLogs(self.api.logger, "WARNING"): + status, data = await self._json("GET", self.URL) + self.assertEqual(status, 500) + self.assertFalse(data["success"]) + + async def test_a_deleted_root_keeps_the_previous_listing_and_warns_once(self): + import shutil + + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + shutil.rmtree(self.output_dir) + with self.assertLogs(self.api.logger, "WARNING") as logs: + for _ in range(2): + self._age() + status, data = await self._json("GET", self.URL) + self.assertEqual((status, data["total"]), (200, 1)) + await self._settle() + warnings = [r for r in logs.records if r.levelname == "WARNING"] + self.assertEqual(len(warnings), 1, [r.getMessage() for r in warnings]) + self.assertEqual(self.scan.calls, 1) # the probe failed before the scan + + async def test_a_root_that_recovers_warns_again_on_its_next_failure(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + real = self.api._scan_gallery_files_sync + with self.assertLogs(self.api.logger, "WARNING") as logs: + for scan in (_raise, real, _raise): + self.api._scan_gallery_files_sync = scan + self._age() + await self._json("GET", self.URL) + await self._settle() + warnings = [r for r in logs.records if r.levelname == "WARNING"] + self.assertEqual(len(warnings), 2) + + async def test_an_invalidation_during_a_walk_discards_it(self): + make_png(self.output_dir / "a.png") + gate = self.scan.hold_next() + reader = asyncio.ensure_future(self._json("GET", self.URL)) + await self.scan.wait_entered() + make_png(self.output_dir / "b.png") + self.api.invalidate_gallery_cache() + gate.set() + status, data = await reader + self.assertEqual((status, data["total"]), (200, 2)) + self.assertEqual(self.scan.calls, 2) + + async def test_an_invalidation_from_a_worker_thread_drops_on_the_loop(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + record = self._record() + loop_thread = threading.get_ident() + threads = [] + original = record.drop + + def drop(): + threads.append(threading.get_ident()) + original() + + record.drop = drop + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self.api.invalidate_gallery_cache) + # Dropped on the loop, before invalidate_gallery_cache returned. + self.assertEqual(threads, [loop_thread]) + self.assertIsNone(record.files) + + async def test_three_discarded_walks_are_a_500(self): + loop = asyncio.get_running_loop() + real = self.scan.real + + async def invalidate(): + self.api.invalidate_gallery_cache() + + def scan(root): + asyncio.run_coroutine_threadsafe(invalidate(), loop).result(30) + return real(root) + + self.scan.real = scan + status, data = await self._json("GET", self.URL) + self.assertEqual(status, 500) + self.assertFalse(data["success"]) + self.assertEqual(self.scan.calls, 3) + + async def test_a_disconnected_waiter_does_not_cancel_the_walk(self): + make_png(self.output_dir / "a.png") + gate = self.scan.hold_next() + waiter = asyncio.ensure_future(self.api._listing(self.output_dir)) + await self.scan.wait_entered() + waiter.cancel() + with self.assertRaises(asyncio.CancelledError): + await waiter + flight = self._record().flight + self.assertIsNotNone(flight) + gate.set() + await asyncio.shield(flight) + self.assertEqual(len(self._record().files), 1) + self.assertIsNone(self._record().flight) + + async def test_a_duplicate_delete_is_gone_on_the_very_next_read(self): + make_png(self.output_dir / "keep.png") + make_png(self.output_dir / "dupe.png") + _, before = await self._json("GET", self.URL) + self.assertEqual(before["total"], 2) + status, data = await self._json( + "POST", + "/prompt_manager/delete_duplicate_images", + json={"image_paths": ["dupe.png"]}, + ) + self.assertEqual((status, data["deleted_count"]), (200, 1)) + _, after = await self._json("GET", self.URL) + self.assertEqual([image["filename"] for image in after["images"]], ["keep.png"]) + + async def test_response_keys_are_unchanged(self): + make_png(self.output_dir / "sub" / "a.png") + _, legacy = await self._json("GET", self.URL) + self.assertEqual( + set(legacy), + {"success", "images", "total", "offset", "limit", "has_more"}, + ) + _, workspace = await self._json("GET", self.URL + "?mode=workspace") + self.assertEqual( + set(workspace), + { + "success", + "images", + "total", + "positive_count", + "limit", + "has_more", + "next_cursor", + }, + ) + _, folders = await self._json("GET", "/prompt_manager/gallery/subfolders") + self.assertEqual(set(folders), {"success", "subfolders"}) + _, scoped = await self._json( + "GET", "/prompt_manager/gallery/subfolders?mode=workspace" + ) + self.assertEqual(set(scoped), {"success", "subfolders", "roots"}) + self.assertEqual(set(scoped["roots"][0]), {"index", "label", "subfolders"}) + _, rescan = await self._json("POST", "/prompt_manager/gallery/rescan") + self.assertEqual(set(rescan), {"success", "total", "roots"}) + + async def test_subfolders_with_ancestors_come_from_the_listing(self): + make_png(self.output_dir / "a" / "b" / "c.png") + _, plain = await self._json("GET", "/prompt_manager/gallery/subfolders") + _, nested = await self._json( + "GET", "/prompt_manager/gallery/subfolders?include_ancestors=true" + ) + self.assertEqual(plain["subfolders"], [str(Path("a") / "b")]) + self.assertEqual( + sorted(nested["subfolders"]), sorted(["a", str(Path("a") / "b")]) + ) + self.assertEqual(self.scan.calls, 1) + + async def test_warm_up_walks_every_root_before_the_first_read(self): + second = self.comfy_dir / "second" + make_png(second / "x.png") + make_png(self.output_dir / "a.png") + self.api._get_all_output_dirs = lambda: [self.output_dir, second] + await self.api.warm_gallery_listings(web.Application()) + self.assertEqual(len(self.api._listing_tasks), 2) + await self._settle() + self.assertEqual(self.scan.calls, 2) + _, data = await self._json("GET", self.URL) + self.assertEqual(data["total"], 2) + self.assertEqual(self.scan.calls, 2) + + async def test_warm_up_never_fails_startup(self): + def broken(): + raise RuntimeError("no roots") + + self.api._get_all_output_dirs = broken + with self.assertLogs(self.api.logger, "WARNING"): + await self.api.warm_gallery_listings(web.Application()) + self.assertEqual(self.api._listing_tasks, set()) + + async def test_warm_up_failure_is_retried_by_the_first_read(self): + make_png(self.output_dir / "a.png") + self.api._scan_gallery_files_sync = _raise + with self.assertLogs(self.api.logger, "WARNING"): + await self.api.warm_gallery_listings(web.Application()) + await self._settle() + self.api._scan_gallery_files_sync = self.scan + status, data = await self._json("GET", self.URL) + self.assertEqual((status, data["total"]), (200, 1)) + + async def test_cleanup_cancels_running_walks(self): + make_png(self.output_dir / "a.png") + gate = self.scan.hold_next() + await self.api.warm_gallery_listings(web.Application()) + await self.scan.wait_entered() + tasks = set(self.api._listing_tasks) + flight = self._record().flight + await self.api.cleanup_gallery_listings(web.Application()) + gate.set() + self.assertTrue(all(task.cancelled() for task in tasks)) + self.assertEqual(self.api._listing_tasks, set()) + self.assertTrue(flight.done()) + self.assertIsNone(self._record().flight) + self.assertIsNone(self._record().files) + + async def test_cleanup_resolves_a_walk_cancelled_before_it_ran(self): + await self.api.warm_gallery_listings(web.Application()) + flight = self._record().flight + waiter = asyncio.ensure_future(asyncio.shield(flight)) + # No await between the warm-up and the cleanup: the task never ran. + await self.api.cleanup_gallery_listings(web.Application()) + outcome = await waiter + self.assertEqual(outcome.kind, gallery_listing.FAILED) + self.assertIsNone(self._record().flight) + self.assertEqual(self.scan.calls, 0) + + async def test_a_stale_read_starts_its_refresh_after_a_delay(self): + # The read is served from the old listing; the walk starts + # REFRESH_DELAY later so it does not share the GIL with the page's + # own requests (#348 live test: 2.2 s reads during the walk). + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + self._age() + with patch.object(gallery_listing, "REFRESH_DELAY", 30): + status, data = await self._json("GET", self.URL) + await self._json("GET", self.URL) # still one pending refresh + record = self._record() + handle = record.rewalk_handle + self.assertEqual((status, data["total"]), (200, 1)) + self.assertIsNotNone(handle) + self.assertIsNone(record.flight) + self.assertEqual(self.scan.calls, 1) + self.assertAlmostEqual( + handle.when() - asyncio.get_running_loop().time(), 30, delta=2 + ) + handle.cancel() + record.rewalk_handle = None + + async def test_the_delayed_refresh_walks_once_when_it_fires(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + make_png(self.output_dir / "b.png") + self._age() + await self._json("GET", self.URL) + await self._json("GET", self.URL) + await self._settle() + self.assertEqual(self.scan.calls, 2) + self.assertIsNone(self._record().rewalk_handle) + _, data = await self._json("GET", self.URL) + self.assertEqual(data["total"], 2) + + async def test_a_walk_started_meanwhile_cancels_the_delayed_refresh(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + self._age() + with patch.object(gallery_listing, "REFRESH_DELAY", 30): + await self._json("GET", self.URL) + handle = self._record().rewalk_handle + status, _ = await self._json("POST", "/prompt_manager/gallery/rescan") + self.assertEqual(status, 200) + self.assertTrue(handle.cancelled()) + self.assertIsNone(self._record().rewalk_handle) + self.assertEqual(self.scan.calls, 2) + + async def test_cleanup_cancels_a_delayed_refresh(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + self._age() + with patch.object(gallery_listing, "REFRESH_DELAY", 30): + await self._json("GET", self.URL) + handle = self._record().rewalk_handle + await self.api.cleanup_gallery_listings(web.Application()) + self.assertTrue(handle.cancelled()) + self.assertIsNone(self._record().rewalk_handle) + self.assertEqual(self.scan.calls, 1) + + WORKSPACE_URL = "/prompt_manager/images/output?mode=workspace&limit=1" + + async def _warmed(self): + await self.api.warm_gallery_listings(web.Application()) + await self._settle() + + def _workspace_ready(self): + """True when a first workspace page needs no per-file rebuild.""" + files = self._record().files + favorites = self.api._favorite_listing_cache.get(str(self.output_dir)) + listed = self.api._workspace_listed_cache.get((0, str(self.output_dir))) + return ( + favorites is not None + and favorites[0] is files + and listed is not None + and listed[0] is files + and bool(self.api._workspace_order_cache) + ) + + async def test_warm_up_also_builds_the_workspace_caches(self): + # The first open after a restart paid ~3 s building these (#348 live). + make_png(self.output_dir / "a.png") + make_png(self.output_dir / "sub" / "b.png") + await self._warmed() + self.assertTrue(self._workspace_ready()) + with patch.object(images_module, "_listed", side_effect=AssertionError): + status, data = await self._json("GET", self.WORKSPACE_URL) + self.assertEqual((status, data["total"]), (200, 2)) + + async def test_a_changed_refresh_rebuilds_the_workspace_caches_behind(self): + make_png(self.output_dir / "a.png") + await self._warmed() + make_png(self.output_dir / "b.png") + self._age() + await self._json("GET", self.URL) # stale: one background walk + await self._settle() + self.assertEqual(len(self._record().files), 2) + self.assertTrue(self._workspace_ready()) + + async def test_an_unchanged_refresh_does_not_rebuild_them(self): + make_png(self.output_dir / "a.png") + await self._warmed() + with patch.object( + self.api, "_build_workspace_caches", side_effect=AssertionError + ): + self._age() + await self._json("GET", self.URL) + await self._settle() + self.assertTrue(self._workspace_ready()) + + async def test_without_the_warm_up_reads_start_no_workspace_build(self): + make_png(self.output_dir / "a.png") + await self._json("GET", self.URL) + await self._settle() + self.assertIsNone(self.api._workspace_warm_task) + self.assertEqual(self.api._favorite_listing_cache, {}) + + async def test_a_failed_workspace_build_warns_and_reads_still_work(self): + make_png(self.output_dir / "a.png") + with ( + patch.object( + self.api, "_select_workspace_page", side_effect=RuntimeError("boom") + ), + self.assertLogs(self.api.logger, "WARNING") as logs, + ): + await self._warmed() + self.assertTrue(any("Images page" in line for line in logs.output)) + status, data = await self._json("GET", self.URL) + self.assertEqual((status, data["total"]), (200, 1)) + + async def test_walks_during_a_workspace_build_coalesce_into_one_more(self): + make_png(self.output_dir / "a.png") + release, entered = threading.Event(), threading.Event() + built = [] + real = self.api._build_workspace_caches + + def build(listings): + built.append(len(listings[0][2])) + if len(built) == 1: + entered.set() + release.wait(_HANG_LIMIT) + real(listings) + + self.addCleanup(release.set) + with patch.object(self.api, "_build_workspace_caches", side_effect=build): + await self.api.warm_gallery_listings(web.Application()) + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, entered.wait, _HANG_LIMIT) + for name in ("b.png", "c.png"): # two changed walks while it builds + make_png(self.output_dir / name) + self._age() + await self._json("GET", self.URL) + while self.api._listing_tasks: + await asyncio.gather(*list(self.api._listing_tasks)) + await asyncio.sleep(0) + release.set() + await self._settle() + self.assertEqual(built, [1, 3]) + self.assertEqual(self.scan.calls, 3) # two stored walks, one more build + self.assertTrue(self._workspace_ready()) + + async def test_a_build_that_cannot_list_the_roots_only_warns(self): + make_png(self.output_dir / "a.png") + await self._warmed() + + def broken(): + raise RuntimeError("no roots") + + self.api._get_all_output_dirs = broken + with self.assertLogs(self.api.logger, "WARNING") as logs: + self.api._schedule_workspace_warm() + task = self.api._workspace_warm_task + await asyncio.wait([task]) + self.assertIsNone(task.exception()) + self.assertTrue(any("Images page" in line for line in logs.output)) + + async def test_cleanup_stops_a_running_workspace_build(self): + make_png(self.output_dir / "a.png") + release, entered = threading.Event(), threading.Event() + + def build(_listings): + entered.set() + release.wait(_HANG_LIMIT) + + self.addCleanup(release.set) + with patch.object(self.api, "_build_workspace_caches", side_effect=build): + await self.api.warm_gallery_listings(web.Application()) + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, entered.wait, _HANG_LIMIT) + task = self.api._workspace_warm_task + await self.api.cleanup_gallery_listings(web.Application()) + release.set() + self.assertTrue(task.done()) + self.assertIsNone(self.api._workspace_warm_task) + + +class TestGalleryListingHostHooks(unittest.IsolatedAsyncioTestCase): + async def test_root_registers_the_warm_up_and_its_cleanup(self): + try: + from test_init_package import fake_server_module, load_package + except ImportError: + from tests.test_init_package import fake_server_module, load_package + + with tempfile.TemporaryDirectory() as directory: + server = fake_server_module() + server.PromptServer.instance.app = web.Application() + with load_package( + "gallery_listing_hooks_pkg", + str(Path(directory) / "db.sqlite"), + directory, + server, + ) as (package, host): + self.assertIn(package.api.warm_gallery_listings, host.app.on_startup) + self.assertIn(package.api.cleanup_gallery_listings, host.app.on_cleanup) diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 1b1113b..d5f1aaa 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -171,6 +171,44 @@ class TestWorkspaceImages(ImageAPITestCase): self.assertIsNone(second["next_cursor"]) self.assertEqual(second["total"], 4) + async def test_an_unchanged_refresh_keeps_the_listing_and_derived_caches(self): + from py.api import gallery_listing + + self.media("a.png", 300) + self.media("b.png", 200) + _, first = await self.get() + response = await self.client.put( + "/prompt_manager/images/favorite", + json={"root_index": 0, "path": "a.png"}, + ) + self.assertEqual(response.status, 200) + record = self.api._root_listings[str(self.output_dir)] + files = record.files + key = ((0, str(self.output_dir)),) + newest = self.api._workspace_order_cache[key][1]["newest"] + favorites = self.api._favorite_listing_cache[str(self.output_dir)] + record.walked_at -= gallery_listing.REFRESH_AFTER + 1 + with patch.object( + self.api, + "_scan_gallery_files_sync", + wraps=self.api._scan_gallery_files_sync, + ) as scan: + with patch.object(gallery_listing, "REFRESH_DELAY", 0): + await self.get() # served at once; one refresh walk follows + while record.rewalk_handle is not None: + await asyncio.sleep(0.01) + while self.api._listing_tasks: + await asyncio.gather(*list(self.api._listing_tasks)) + # A finished walk leaves the set in a done callback, which needs + # a loop turn; gather over finished tasks does not yield one. + await asyncio.sleep(0) + _, again = await self.get() + self.assertEqual(scan.call_count, 1) + self.assertIs(record.files, files) + self.assertIs(self.api._workspace_order_cache[key][1]["newest"], newest) + self.assertIs(self.api._favorite_listing_cache[str(self.output_dir)], favorites) + self.assertEqual(again["total"], first["total"]) + async def test_filters_apply_before_limit_and_sorts_are_complete(self): self.media("folder/Zebra.png", 200) self.media("folder/apple.png", 300) diff --git a/tests/test_gallery_listing.py b/tests/test_gallery_listing.py new file mode 100644 index 0000000..3a57922 --- /dev/null +++ b/tests/test_gallery_listing.py @@ -0,0 +1,185 @@ +"""Pure helpers of the shared Images listing (py/api/gallery_listing.py, #348).""" + +import os +import sys +import tempfile +import unittest +from pathlib import Path, PurePath, PureWindowsPath +from unittest.mock import patch + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from py.api import gallery_listing # noqa: E402 +from py.api.gallery_listing import ( # noqa: E402 + DISCARDED, + FAILED, + STORED, + RootListing, + WalkOutcome, + folder_sets, + same_listing, + walk_root, +) + + +def _legacy_subfolders(root, files, include_ancestors): + """``_root_subfolders`` as it was before #348, the reference for folder_sets.""" + subfolders = set() + for f, _mtime in files: + rel_dir = str(f.relative_to(root).parent) + if rel_dir and rel_dir != ".": + subfolders.add(rel_dir) + if include_ancestors: + ancestors = set() + for folder in subfolders: + parts = folder.replace("\\", "/").split("/") + for i in range(1, len(parts)): + ancestors.add("/".join(parts[:i])) + subfolders.update(ancestors) + return subfolders + + +class FolderSetsTests(unittest.TestCase): + def test_nested_folders_match_the_previous_subfolder_listing(self): + root = Path("/out") + files = [ + (root / "top.png", 1.0), + (root / "a" / "one.png", 2.0), + (root / "a" / "b" / "c" / "two.png", 3.0), + (root / "d" / "three.png", 4.0), + ] + folders, with_ancestors = folder_sets(root, files) + self.assertEqual(folders, _legacy_subfolders(root, files, False)) + self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True)) + self.assertIn("a/b", with_ancestors) + self.assertNotIn("a/b", folders) + self.assertIsInstance(folders, frozenset) + self.assertIsInstance(with_ancestors, frozenset) + + def test_windows_separators_keep_the_previous_spelling(self): + root = PureWindowsPath("C:/out") + files = [ + (root / "x" / "y" / "z.png", 1.0), + (root / "w.png", 2.0), + ] + folders, with_ancestors = folder_sets(root, files) + self.assertEqual(folders, {"x\\y"}) + self.assertEqual(with_ancestors, {"x\\y", "x"}) + self.assertEqual(folders, _legacy_subfolders(root, files, False)) + self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True)) + + def test_an_empty_listing_has_no_folders(self): + self.assertEqual(folder_sets(Path("/out"), []), (frozenset(), frozenset())) + + def test_folders_come_from_the_path_strings_not_relative_to(self): + # relative_to per file was half of a 20,000-file walk (#348 live test). + root = Path("/out") + files = [(root / "a" / "b" / "one.png", 1.0), (root / "two.png", 2.0)] + with patch.object(PurePath, "relative_to", side_effect=AssertionError): + folders, with_ancestors = folder_sets(root, files) + self.assertEqual(folders, {os.path.join("a", "b")}) + self.assertEqual(with_ancestors, {os.path.join("a", "b"), "a"}) + + def test_a_path_not_spelled_like_the_root_falls_back_to_relative_to(self): + # Windows folds case: the string prefix misses, relative_to matches. + root = PureWindowsPath("C:/Out") + files = [(PureWindowsPath("c:/out/x/y.png"), 1.0)] + folders, with_ancestors = folder_sets(root, files) + self.assertEqual(folders, {"x"}) + self.assertEqual(folders, _legacy_subfolders(root, files, False)) + self.assertEqual(with_ancestors, _legacy_subfolders(root, files, True)) + + +class SameListingTests(unittest.TestCase): + def test_equal_path_and_mtime_lists_are_the_same(self): + old = [(Path("/o/a.png"), 1.0), (Path("/o/b.png"), 2.0)] + self.assertTrue(same_listing(old, list(old))) + self.assertTrue(same_listing(old, old)) + + def test_a_changed_mtime_path_or_length_differs(self): + old = [(Path("/o/a.png"), 1.0), (Path("/o/b.png"), 2.0)] + self.assertFalse(same_listing(old, [(Path("/o/a.png"), 1.5), old[1]])) + self.assertFalse(same_listing(old, [(Path("/o/c.png"), 1.0), old[1]])) + self.assertFalse(same_listing(old, old[:1])) + self.assertFalse(same_listing(None, old)) + + +class WalkRootTests(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.root = Path(self._tmp.name) / "output" + (self.root / "sub").mkdir(parents=True) + + def _scan(self, files): + calls = [] + + def scan(root): + calls.append(root) + return list(files) + + return scan, calls + + def test_returns_the_listing_and_both_folder_sets(self): + files = [(self.root / "sub" / "a.png", 2.0), (self.root / "b.png", 1.0)] + scan, calls = self._scan(files) + listed, folders, with_ancestors = walk_root(scan, self.root, None) + self.assertEqual(calls, [self.root]) + self.assertEqual(listed, files) + self.assertEqual(folders, frozenset({"sub"})) + self.assertEqual(with_ancestors, frozenset({"sub"})) + + def test_an_unchanged_walk_returns_the_previous_list_object(self): + files = [(self.root / "sub" / "a.png", 2.0)] + previous = list(files) + scan, _calls = self._scan(files) + listed, _folders, _ancestors = walk_root(scan, self.root, previous) + self.assertIs(listed, previous) + + def test_a_changed_walk_returns_the_new_list(self): + previous = [(self.root / "sub" / "a.png", 2.0)] + files = [(self.root / "sub" / "a.png", 3.0)] + scan, _calls = self._scan(files) + listed, _folders, _ancestors = walk_root(scan, self.root, previous) + self.assertIsNot(listed, previous) + self.assertEqual(listed, files) + + def test_a_deleted_root_raises_instead_of_listing_nothing(self): + scan, calls = self._scan([]) + missing = self.root / "gone" + with self.assertRaises(OSError): + walk_root(scan, missing, None) + self.assertEqual(calls, []) + + +class RecordTests(unittest.TestCase): + def test_drop_clears_the_listing_and_bumps_the_generation(self): + record = RootListing() + record.files = [(Path("/o/a.png"), 1.0)] + record.folders = frozenset({"x"}) + record.folders_with_ancestors = frozenset({"x"}) + record.walked_at = 5.0 + record.drop() + record.drop() + self.assertIsNone(record.files) + self.assertIsNone(record.folders) + self.assertIsNone(record.folders_with_ancestors) + self.assertIsNone(record.walked_at) + self.assertEqual(record.generation, 2) + + def test_outcomes(self): + error = OSError("gone") + self.assertEqual(WalkOutcome.stored().kind, STORED) + self.assertEqual(WalkOutcome.discarded().kind, DISCARDED) + failed = WalkOutcome.failed(error) + self.assertEqual(failed.kind, FAILED) + self.assertIs(failed.error, error) + + def test_refresh_constants(self): + self.assertEqual(gallery_listing.REFRESH_AFTER, 30) + self.assertEqual(gallery_listing.REFRESH_DELAY, 1.0) + self.assertIsNone(RootListing().rewalk_handle) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_thumbnail_convert.py b/tests/test_thumbnail_convert.py index 8383fd3..c549e8e 100644 --- a/tests/test_thumbnail_convert.py +++ b/tests/test_thumbnail_convert.py @@ -85,6 +85,14 @@ class TestConvertOutdated(ConvertCase): ) self.assertEqual(self.api._jobs.list(), []) + async def test_conversion_keeps_the_listing(self): + self._legacy("a") + await self.client.get("/prompt_manager/images/output?mode=workspace") + files = self.api._root_listings[str(self.output_dir)].files + _summary, job = await self._convert() + self.assertEqual(job.state, "finished") + self.assertIs(self.api._root_listings[str(self.output_dir)].files, files) + async def test_converts_from_the_source_and_removes_the_old_preview(self): old = self._legacy("a") original = (self.output_dir / "a.png").read_bytes() diff --git a/tests/test_thumbnail_jobs.py b/tests/test_thumbnail_jobs.py index fe4fd96..3daed5b 100644 --- a/tests/test_thumbnail_jobs.py +++ b/tests/test_thumbnail_jobs.py @@ -345,6 +345,20 @@ class TestThumbnailJob(ThumbnailJobCase): sizes[quality] = thumb.stat().st_size self.assertLess(sizes["low"], sizes["high"]) + async def test_a_job_keeps_the_listing_and_the_next_page_shows_the_preview(self): + make_png(self.output_dir / "a.png", size=(64, 64)) + url = "/prompt_manager/images/output?mode=workspace" + before = (await (await self.client.get(url)).json())["images"][0] + files = self.api._root_listings[str(self.output_dir)].files + status, body = await self._start(mode="missing") + self.assertEqual(status, 202) + await self._finished(body["job"]["id"]) + # Previews are not part of the listing (#348): no rewalk is forced. + self.assertIs(self.api._root_listings[str(self.output_dir)].files, files) + after = (await (await self.client.get(url)).json())["images"][0] + self.assertNotIn("v=", before["thumbnail_url"]) + self.assertIn("v=", after["thumbnail_url"]) + async def test_no_output_folder_fails_the_job(self): self.api._get_all_output_dirs = lambda: [] _, body = await self._start() diff --git a/tools/vendor/perf/images_load_probe.mjs b/tools/vendor/perf/images_load_probe.mjs new file mode 100644 index 0000000..b448682 --- /dev/null +++ b/tools/vendor/perf/images_load_probe.mjs @@ -0,0 +1,81 @@ +// Maintainer tool, run by hand (never in CI): how long the workspace Images +// page takes to show its first thumbnail (#348, TESTING.md "Images listing +// cache (#348)"). +// +// Once: npm ci --prefix tools/vendor +// npx --prefix tools/vendor playwright install chromium +// Then: node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188 +// +// Prints the time to the first loaded thumbnail, the time to network idle and +// the 10 slowest requests. Each run opens a fresh browser context (no cache). +import { chromium } from '@playwright/test'; + +const SLOWEST = 10; +const TIMEOUT_MS = 120_000; + +function usage() { + console.error('usage: node tools/vendor/perf/images_load_probe.mjs '); + console.error('example: node tools/vendor/perf/images_load_probe.mjs http://127.0.0.1:8188'); + process.exit(2); +} + +function parseBase(argument) { + if (!argument) usage(); + let url; + try { + url = new URL(argument); + } catch { + usage(); + } + if (url.protocol !== 'http:' && url.protocol !== 'https:') usage(); + return url.origin + url.pathname.replace(/\/+$/, ''); +} + +const ms = value => `${Math.round(value)} ms`; + +async function probe(base) { + const browser = await chromium.launch(); + try { + const context = await browser.newContext({ viewport: { width: 1440, height: 900 } }); + const page = await context.newPage(); + const requests = []; + page.on('requestfinished', request => { + const timing = request.timing(); + requests.push({ + url: request.url().replace(base, ''), + status: null, + ms: timing.responseEnd >= 0 ? timing.responseEnd : NaN, + response: request.response(), + }); + }); + const started = Date.now(); + await page.goto(`${base}/prompt_manager/workspace#/images`, { timeout: TIMEOUT_MS }); + // The first Images thumbnail the browser has decoded. + await page.waitForFunction(() => [...document.images].some(image => + image.complete && image.naturalWidth > 0 + && image.src.includes('/prompt_manager/images/serve/')), null, + { timeout: TIMEOUT_MS, polling: 'raf' }); + const firstThumbnail = Date.now() - started; + await page.waitForLoadState('networkidle', { timeout: TIMEOUT_MS }); + const networkIdle = Date.now() - started; + for (const entry of requests) { + entry.status = (await entry.response)?.status() ?? 'n/a'; + } + const slowest = requests.filter(entry => Number.isFinite(entry.ms)) + .sort((a, b) => b.ms - a.ms).slice(0, SLOWEST); + console.log(`first thumbnail: ${ms(firstThumbnail)}`); + console.log(`network idle: ${ms(networkIdle)}`); + console.log(`requests: ${requests.length}`); + console.log(`slowest ${slowest.length}:`); + for (const entry of slowest) { + console.log(` ${ms(entry.ms).padStart(9)} ${String(entry.status).padEnd(3)} ${entry.url}`); + } + } finally { + await browser.close(); + } +} + +probe(parseBase(process.argv[2])).catch(error => { + console.error(`probe failed: ${error.message}`); + process.exit(1); +});