ComfyAssets-ComfyUI_PromptManager/py/api/preview_store.py

"""Where previews live on disk (#320): the one resolver.

Every mapping between an image source (a "root"), one of its media files and
that file's preview on disk goes through this module. URL text stays in
``thumbnail_settings`` (``thumbnail_rel``: ``thumbnails/<rel>_thumb<ext>``,
the shape the serve route, ``inspection-state.mjs`` and ``admin.js`` share);
this module answers where such a preview is on disk.

The store of a source is where its previews live. With the default location
(mode ``source``, "Inside each source") it is ``<source>/thumbnails``, exactly
the folder every caller joined by hand before, so nothing moves.

Writes and the switch: every preview writer takes ``write_guard`` when it
starts (it records the location's ``generation``); the final rename runs
under ``write_lock`` and only when the generation is unchanged and no hold
applies (``thumbnail_encode.replace_unless_held``). Changing the location
bumps the generation under the same lock, so no write that began before can
land in the old store afterwards. Nothing reads or commits config while
``write_lock`` is held. The hold check runs under it too, so hold checks take
no lock (the job registry's ``running()`` is a lock-free snapshot) and
``bump_generation`` is never called with a registry or config lock held.

All functions are blocking: call them off the event loop.
"""

import errno
import logging
import os
import threading
import time
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path

from . import preview_location, thumbnail_settings

try:
    from ...utils.media_folders import PREVIEW_FOLDER
except ImportError:  # tests import the package top-level
    from utils.media_folders import PREVIEW_FOLDER

logger = logging.getLogger("prompt_manager.thumbnails")

SOURCE_MODE = preview_location.SOURCE_MODE
FOLDER_MODE = preview_location.FOLDER_MODE
# The probe checks the folder this often, and a check (or a store call)
# taking longer than the budget marks it unreachable.
PROBE_INTERVAL = 5.0
IO_BUDGET = 2.0
PROBE_THREAD = "pm-preview-probe"


@dataclass(frozen=True)
class Location:
    """One snapshot of the thumbnail location (take one per request or job)."""

    mode: str = SOURCE_MODE
    path: str | None = None
    available: bool = True
    problem: str | None = None
    generation: int = 0


write_lock = threading.Lock()
_generation = 0
# The runtime location: the parsed setting, availability and the marker keys
# last read. Swapped as a whole; readers take one snapshot.
_state_lock = threading.Lock()
_parsed = preview_location.ParsedLocation()
_available = True
_problem = None
_keys = {}
_key_cache = {}
_clock = time.monotonic


def current():
    """The thumbnail location now."""
    parsed = _parsed
    if parsed.mode == SOURCE_MODE:
        return Location(generation=_generation)
    return Location(FOLDER_MODE, parsed.path, _available, _problem, _generation)


def setting():
    """The parsed setting in use (``preview_location.ParsedLocation``)."""
    return _parsed


def generation():
    return _generation


def bump_generation():
    """Invalidate every write that started before (under ``write_lock``)."""
    global _generation
    with write_lock:
        _generation += 1
        return _generation


def _canonical(path):
    """Resolved and case-folded, so paths compare on every OS."""
    return Path(os.path.normcase(os.path.realpath(str(path))))


def _within(path, folder):
    return _canonical(path).is_relative_to(_canonical(folder))


# ── the setting and availability ─────────────────────────────────────


def configure(value, *, context=None, check=True):
    """Use the ``thumbnails.location`` *value* (config.json) from now on.

    A folder location is validated against *context* (``startup`` rules: a
    problem other than an unreachable folder makes it ``invalid``), then
    checked once; the probe keeps checking it. Swapped under ``write_lock``
    with a generation bump, so no write that started before lands.
    """
    global _parsed, _available, _problem, _generation
    parsed = preview_location.parse_location(value)
    problem = parsed.problem
    if parsed.mode == FOLDER_MODE and problem is None and context is not None:
        code = preview_location.location_problem(
            parsed.path, purpose="startup", context=context
        )
        if code == "location-unavailable":
            problem = "unreachable"
        elif code is not None:
            problem = "invalid"
            parsed = preview_location.ParsedLocation(
                mode=FOLDER_MODE,
                path=parsed.path,
                id=parsed.id,
                install=parsed.install,
                cleanup=parsed.cleanup,
                problem="invalid",
            )
    with write_lock:
        with _state_lock:
            _parsed = parsed
            _available = parsed.mode == SOURCE_MODE
            _problem = problem if parsed.mode == FOLDER_MODE else None
            _keys.clear()
            _key_cache.clear()
        _generation += 1
    if parsed.mode == FOLDER_MODE:
        if check and parsed.problem is None:
            refresh()
        _probe.start()
    else:
        _probe.stop()
    return current()


def refresh():
    """Check the folder now (bounded by IO_BUDGET); returns availability."""
    parsed = _parsed
    if parsed.mode != FOLDER_MODE or parsed.problem is not None:
        return parsed.mode == SOURCE_MODE
    try:
        available, problem, marker = _probe.check(parsed.path, parsed.id)
    except (OSError, TimeoutError, RuntimeError):
        available, problem, marker = False, "unreachable", None
    _set_state(parsed, available, problem, marker)
    return available


def _set_state(parsed, available, problem, marker):
    global _available, _problem
    with _state_lock:
        if _parsed is not parsed:
            return  # the setting changed meanwhile
        _available, _problem = available, problem
        keys = dict((marker or {}).get("keys") or {})
        if keys != _keys:
            _keys.clear()
            _keys.update(keys)
            _key_cache.clear()


def note_failure():
    """A store call failed or stalled: unreachable until the probe finds the
    folder again (requests do no store I/O meanwhile)."""
    global _available, _problem
    if _parsed.mode != FOLDER_MODE:
        return
    with _state_lock:
        _available, _problem = False, "unreachable"
    _probe.wake()


@contextmanager
def store_io():
    """Around a store call in folder mode: an I/O error (other than a file
    that is not there) or a call over IO_BUDGET marks the folder
    unreachable."""
    if _parsed.mode != FOLDER_MODE:
        yield
        return
    started = _clock()
    try:
        yield
    except OSError as error:
        if not isinstance(error, (FileNotFoundError, NotADirectoryError)):
            note_failure()
        raise
    finally:
        if _clock() - started > IO_BUDGET:
            note_failure()


class _Probe:
    """The background check of the thumbnail folder (one thread, one check
    at a time; a stuck check is left to finish on its own thread)."""

    def __init__(self):
        self._lock = threading.Lock()
        self._wake = threading.Event()
        self._stop = threading.Event()
        self._thread = None
        self._checking = None

    def check(self, path, expected_id):
        """``folder_state`` within IO_BUDGET; raises TimeoutError. A check
        already running (the probe's, or a stuck one) is joined, never run
        twice at once."""
        with self._lock:
            pending = self._checking
            if pending is None or pending[0] != (path, expected_id):
                pending = self._start_check(path, expected_id)
        key, done, result = pending
        if not done.wait(IO_BUDGET):
            raise TimeoutError("the thumbnail folder did not answer in time")
        with self._lock:
            if self._checking is pending:
                self._checking = None
        if "error" in result:
            raise result["error"]
        return result["value"]

    def _start_check(self, path, expected_id):
        """Start one check on its own thread (lock held)."""
        result = {}
        done = threading.Event()

        def run():
            try:
                result["value"] = preview_location.folder_state(path, expected_id)
            except BaseException as error:  # handed to the caller
                result["error"] = error
            finally:
                done.set()

        pending = ((path, expected_id), done, result)
        self._checking = pending
        threading.Thread(target=run, name=PROBE_THREAD + "-check", daemon=True).start()
        return pending

    def start(self):
        with self._lock:
            if self._thread is not None and self._thread.is_alive():
                return
            self._stop.clear()
            thread = threading.Thread(target=self._run, name=PROBE_THREAD, daemon=True)
            try:
                thread.start()
            except RuntimeError:
                return  # refused at interpreter exit: no probe
            self._thread = thread

    def _run(self):
        while not self._stop.is_set():
            self._wake.wait(PROBE_INTERVAL)
            self._wake.clear()
            if self._stop.is_set():
                return
            try:
                refresh()
            except Exception:  # one failed check never stops the probe
                logger.debug("Thumbnail folder check failed", exc_info=True)

    def wake(self):
        self._wake.set()

    def stop(self, timeout=2.0):
        """Stop the probe; True when its thread ended."""
        with self._lock:
            thread, self._thread = self._thread, None
        self._stop.set()
        self._wake.set()
        if thread is None or thread is threading.current_thread():
            return True
        thread.join(timeout)
        return not thread.is_alive()


_probe = _Probe()


def stop_probe(timeout=2.0):
    """Stop the folder probe (the package teardown, tests)."""
    return _probe.stop(timeout)


# ── stores ───────────────────────────────────────────────────────────


def store_container(root, location=None):
    """The folder a store of *root* must stay inside: the source itself, or
    the thumbnail folder in folder mode."""
    location = location or current()
    if location.mode == FOLDER_MODE:
        return Path(location.path) if location.path else None
    return Path(root)


def _identity(root):
    return os.path.normcase(os.path.realpath(str(root)))


def _lookup_key(root):
    spelling = str(root)
    with _state_lock:
        if spelling in _key_cache:
            return _key_cache[spelling]
        keys = dict(_keys)
    key = preview_location.find_key(
        keys,
        spelling,
        _identity(root),
        preview_location.file_identity(root),
    )
    if key is not None:
        with _state_lock:
            _key_cache[spelling] = key
    return key


def store_dir(root, location=None):
    """The store of *root*: ``<source>/thumbnails``, or in folder mode
    ``<folder>/<key>`` (None while the folder is unavailable or before the
    source has a key: a reader then finds no preview)."""
    location = location or current()
    if location.mode == SOURCE_MODE:
        return Path(root) / PREVIEW_FOLDER
    if not location.available or not location.path:
        return None
    key = _lookup_key(root)
    return None if key is None else Path(location.path) / key


def ensure_store(root, location=None, *, create=False):
    """``store_dir`` for a writer: in folder mode a source without a key gets
    one (recorded in the marker); with *create* its key folder is made (never
    the thumbnail folder itself). None while the folder is unavailable."""
    location = location or current()
    store = store_dir(root, location)
    if location.mode == SOURCE_MODE or (store is None and not location.available):
        return store
    if store is None:
        store = _assign_key(root, location)
        if store is None:
            return None
    if create:
        try:
            with store_io():
                os.mkdir(store)
        except FileExistsError:
            pass
        except OSError:
            return None
    return store


def _assign_key(root, location):
    parsed = _parsed
    if location.path != parsed.path:
        return None
    spelling = str(root)
    identity = _identity(root)
    dev_ino = preview_location.file_identity(root)
    assigned = {}

    def add(marker):
        keys = marker["keys"]
        key = preview_location.find_key(keys, spelling, identity, dev_ino)
        if key is None:
            taken = {k: v.get("identity") for k, v in keys.items()}
            key = preview_location.source_key(spelling, identity, taken)
            keys[key] = {
                "path": spelling,
                "identity": identity,
                "dev_ino": dev_ino,
                "install": parsed.install,
            }
        assigned["key"] = key

    try:
        with store_io(), write_lock:
            marker = preview_location.write_marker(location.path, add)
    except (OSError, ValueError):
        return None
    _set_state(parsed, True, None, marker)
    return Path(location.path) / assigned["key"]


def make_parent(dst, location=None):
    """Make the folders a preview *dst* is written into. In folder mode only
    folders below the thumbnail folder are made, never the folder itself (a
    missing folder raises FileNotFoundError: it may be an unplugged drive)."""
    location = location or current()
    dst = Path(dst)
    if location.mode == SOURCE_MODE or not location.path:
        dst.parent.mkdir(parents=True, exist_ok=True)
        return
    folder = Path(location.path)
    try:
        relative = dst.parent.relative_to(folder)
    except ValueError:
        dst.parent.mkdir(parents=True, exist_ok=True)  # a store given by a test
        return
    with store_io():
        if not location.available or not folder.is_dir():
            raise FileNotFoundError(errno.ENOENT, "thumbnail folder unavailable")
        current_folder = folder
        for part in relative.parts:
            current_folder = current_folder / part
            try:
                os.mkdir(current_folder)
            except FileExistsError:
                pass


def given_store(root, store):
    """A caller's explicit store argument: the default spelling
    (``<root>/thumbnails``) means the resolver's store; another folder (only
    tests pass one) is used as given."""
    store = Path(store)
    if store == Path(root) / PREVIEW_FOLDER:
        return store_dir(root)
    return store


def preview_path(root, media_rel, ext=None, location=None, *, create=False):
    """Where the preview of *media_rel* (relative to *root*) is, or is to be
    written (*create*: a writer, which gives a new source its key), in *ext*
    (the extension written now when None)."""
    if create:
        store = ensure_store(root, location)
    else:
        store = store_dir(root, location)
    if store is None:
        return None
    if ext is None:
        from .thumbnail_encode import writer_settings

        ext = writer_settings().extension
    return store / (thumbnail_settings.preview_stem(media_rel) + ext)


def request_path(root, parts, location=None):
    """The disk path of the URL parts ``("thumbnails", ...)``, or None when
    *parts* is not under an exact ``thumbnails`` first part or the store has
    no place for it. The text is checked before (``is_preview_text``)."""
    parts = tuple(parts)
    if len(parts) < 2 or parts[0] != PREVIEW_FOLDER:
        return None
    store = store_dir(root, location)
    if store is None:
        return None
    return store.joinpath(*parts[1:])


def _no_link_below(store, path):
    """No folder from *store* down to *path* (and *store* itself) is a link,
    junction or reparse point (lexical walk, one lstat each)."""
    store = Path(os.path.abspath(store))
    path = Path(os.path.abspath(path))
    try:
        relative = path.relative_to(store)
    except ValueError:
        return False
    current = store
    for part in ("",) + relative.parts:
        current = current / part if part else current
        if preview_location.is_link_like(current):
            return False
    return True


def owns(root, path, location=None):
    """Whether *path* lies inside the store of *root*, resolved: the store
    stays inside its container and no link takes *path* out of the store.
    In folder mode no link, junction or reparse point may sit anywhere from
    the store down to *path*."""
    location = location or current()
    store = store_dir(root, location)
    if store is None:
        return False
    container = store_container(root, location)
    if not (_within(store, container) and _within(path, store)):
        return False
    if location.mode == SOURCE_MODE:
        return True
    try:
        return _no_link_below(store, path)
    except OSError:
        note_failure()
        return False


# ── preview text ─────────────────────────────────────────────────────


def is_preview_text(filepath):
    """URL text naming a preview: ``thumbnails/<...>/<name>_thumb<ext>``.

    Text only, before any join: no backslash, drive or stream colon, and no
    empty, ``.`` or ``..`` segment; ``thumbnails`` exactly as the first part
    and nowhere else among the folders; the extension, in any case, one a
    preview may carry (``X_thumb.JPG`` of an older release included).
    """
    if not filepath or "\\" in filepath or ":" in filepath:
        return False
    parts = filepath.split("/")
    if len(parts) < 2 or parts[0] != PREVIEW_FOLDER:
        return False
    if any(part in ("", ".", "..") for part in parts):
        return False
    if PREVIEW_FOLDER in parts[1:-1]:
        return False
    stem, extension = os.path.splitext(parts[-1])
    return (
        stem.endswith(thumbnail_settings.THUMB_SUFFIX)
        and len(stem) > len(thumbnail_settings.THUMB_SUFFIX)
        and extension.lower() in thumbnail_settings.THUMBNAIL_EXTENSIONS
    )


# ── writes ───────────────────────────────────────────────────────────


class _WriteGuard:
    """``guard()`` for one preview write: true while the location has not
    changed since the write started and *hold* (if any) allows it."""

    __slots__ = ("generation", "_hold")

    def __init__(self, hold=None):
        self.generation = _generation
        self._hold = hold

    def __call__(self):
        if self.generation != _generation:
            return False
        return self._hold is None or bool(self._hold())


def write_guard(hold=None):
    """Take when a preview write starts; hand to its final rename."""
    if isinstance(hold, _WriteGuard):
        return hold
    return _WriteGuard(hold)


@contextmanager
def committed(guard=None):
    """Around a preview's final rename: holds ``write_lock`` and yields
    whether the write may land (``guard()``, re-checked under the lock)."""
    with write_lock:
        yield guard is None or bool(guard())