diff --git a/cista/app.py b/cista/app.py index a8ff2e3..a2e8897 100644 --- a/cista/app.py +++ b/cista/app.py @@ -19,8 +19,9 @@ from stream_zip import ZIP_AUTO, stream_zip from zstandard import ZstdCompressor from cista import auth, config, preview, session, sso, watching +from cista.preview import shutdown_preview_workers, start_preview_workers from cista.api import bp -from cista.sanic_logging import configure_access_logging, format_access_log +from cista.sanic_logging import configure_access_logging, configure_main_logging, format_access_log from cista.sanic_logging import logger as access_logger from cista.util.apphelpers import handle_sanic_exception @@ -30,6 +31,7 @@ sanic.helpers._ENTITY_HEADERS = frozenset() configure_access_logging() app = Sanic("cista", strict_slashes=True) +configure_main_logging() # Register either SSO proxy or built-in auth routes based on PASKIA_BACKEND_URL if sso.paskia_enabled(): app.blueprint(sso.bp) # SSO proxy for /auth/* routes @@ -47,13 +49,12 @@ setproctitle("cista-main") async def main_start(app): config.load_config() setproctitle(f"cista {config.config.path.name}") - # Small pool for memory-intensive preview generation - preview_workers = max(2, min(8, cpu_count())) app.ctx.threadexec = ThreadPoolExecutor( - max_workers=preview_workers, thread_name_prefix="cista-preview" + max_workers=4, thread_name_prefix="cista-worker" ) # Larger pool for long-running but low-memory zip operations app.ctx.zipexec = ThreadPoolExecutor(max_workers=32, thread_name_prefix="cista-zip") + await start_preview_workers() watching.start(app) @@ -61,6 +62,7 @@ async def main_start(app): @app.before_server_stop async def main_stop(app): watching.stop(app) + await shutdown_preview_workers() app.ctx.threadexec.shutdown() app.ctx.zipexec.shutdown(cancel_futures=True) await sso.close_client() @@ -102,7 +104,8 @@ async def log_access(req, res): if isinstance(qs, bytes): qs = qs.decode(errors="replace") path = f"{path}?{qs}" - line = format_access_log(client, res.status, req.method, host, path, duration_ms) + extra = getattr(req.ctx, "_log_extra", None) + line = format_access_log(client, res.status, req.method, host, path, duration_ms, extra=extra) access_logger.info(line) return res diff --git a/cista/preview.py b/cista/preview.py index 5793006..ba10fc5 100644 --- a/cista/preview.py +++ b/cista/preview.py @@ -2,25 +2,33 @@ import asyncio import gc import io import mimetypes +import os +import struct +import sys import threading import urllib.parse from collections import OrderedDict from dataclasses import dataclass +from multiprocessing import cpu_count from pathlib import PurePosixPath from time import perf_counter from urllib.parse import unquote from wsgiref.handlers import format_date_time +import msgspec + import av import fitz # PyMuPDF import numpy as np import pillow_heif +from blake3 import blake3 from PIL import Image from sanic import Blueprint, empty, raw, redirect from sanic.exceptions import NotFound from sanic.log import logger from cista import auth, config +from cista.preview_worker import PreviewRequest, PreviewResponse from cista.util.filename import sanitize pillow_heif.register_heif_opener() @@ -70,6 +78,194 @@ class PreviewCache: # Global preview cache instance _preview_cache = PreviewCache(capacity=500) +PREVIEW_TIMEOUT = 3.0 # seconds until preview subprocess is killed +PREVIEW_WORKERS = max(2, min(8, cpu_count())) +_active_procs: set[asyncio.subprocess.Process] = set() +_preview_pool = None +_preview_pool_lock = asyncio.Lock() +AVIF_FAST_EFFORT = 0 +FORCE_PIL = os.environ.get("CISTA_PIL") == "1" +WORKER_CHECKSUM_BYTES = 32 +WORKER_MAX_JSON_BYTES = 1_000_000 + + +class WorkerChecksumError(Exception): + """Raised when worker response checksum does not match the packet.""" + + +class WorkerProtocolError(Exception): + """Raised when worker response packet is malformed.""" + + +class _PreviewWorker: + def __init__(self, proc: asyncio.subprocess.Process): + self.proc = proc + + async def request(self, filepath, quality: int, maxsize: int, maxzoom: float): + if self.proc.returncode is not None: + raise WorkerProtocolError("worker already exited") + if self.proc.stdin is None or self.proc.stdout is None: + raise WorkerProtocolError("worker streams not available") + + line = msgspec.json.encode( + PreviewRequest(path=str(filepath), quality=quality, maxsize=maxsize, maxzoom=maxzoom) + ) + b"\n" + self.proc.stdin.write(line) + await self.proc.stdin.drain() + + checksum = await self.proc.stdout.readexactly(WORKER_CHECKSUM_BYTES) + header = await self.proc.stdout.readexactly(8) + json_size, data_size = struct.unpack(" WORKER_MAX_JSON_BYTES: + raise WorkerProtocolError(f"worker JSON too large: {json_size}") + meta_raw = await self.proc.stdout.readexactly(json_size) + payload = await self.proc.stdout.readexactly(data_size) + packet = header + meta_raw + payload + if blake3(packet).digest() != checksum: + raise WorkerChecksumError("worker checksum mismatch") + + resp = msgspec.json.decode(meta_raw, type=PreviewResponse) + if not resp.ok: + raise PreviewError(resp.error or "preview worker error") + return payload or None, resp + + async def kill(self) -> None: + if self.proc.returncode is None: + try: + self.proc.kill() + except ProcessLookupError: + pass + await self.proc.wait() + _active_procs.discard(self.proc) + + +class _PreviewWorkerPool: + def __init__(self, size: int): + self.size = size + self._idle: asyncio.Queue[_PreviewWorker] = asyncio.Queue() + self._workers: set[_PreviewWorker] = set() + self._closed = False + + async def _spawn_worker(self) -> _PreviewWorker: + proc = await asyncio.create_subprocess_exec( + sys.executable, + "-m", + "cista.preview_worker", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, + start_new_session=True, + ) + _active_procs.add(proc) + return _PreviewWorker(proc) + + async def _add_worker(self) -> None: + worker = await self._spawn_worker() + self._workers.add(worker) + await self._idle.put(worker) + + async def _replace_worker(self, worker: _PreviewWorker) -> None: + self._workers.discard(worker) + await worker.kill() + if self._closed: + return + try: + await self._add_worker() + except Exception: + logger.exception("Failed to replace preview worker") + + async def start(self) -> None: + for _ in range(self.size): + await self._add_worker() + + async def run(self, filepath, quality: int, maxsize: int, maxzoom: float): + if self._closed: + raise PreviewError("preview worker pool closed") + worker = await self._idle.get() + replace = False + try: + out, resp = await asyncio.wait_for( + worker.request(filepath, quality, maxsize, maxzoom), + timeout=PREVIEW_TIMEOUT, + ) + return out, resp + except asyncio.TimeoutError: + replace = True + logger.warning( + "Preview timeout (%ds) for %s", int(PREVIEW_TIMEOUT), filepath.name + ) + raise PreviewTimeout(filepath.name) + except WorkerChecksumError: + replace = True + logger.error("Preview checksum mismatch for %s", filepath.name) + raise PreviewError(filepath.name) + except PreviewError: + raise + except ( + WorkerProtocolError, + asyncio.IncompleteReadError, + BrokenPipeError, + ConnectionResetError, + OSError, + ValueError, + msgspec.json.DecodeError, + ) as e: + replace = True + logger.warning("Preview worker protocol failure for %s: %s", filepath.name, e) + raise PreviewError(filepath.name) + finally: + if replace: + await self._replace_worker(worker) + else: + if worker.proc.returncode is None: + await self._idle.put(worker) + else: + await self._replace_worker(worker) + + async def close(self) -> None: + self._closed = True + workers = list(self._workers) + self._workers.clear() + while not self._idle.empty(): + try: + self._idle.get_nowait() + except asyncio.QueueEmpty: + break + await asyncio.gather(*(worker.kill() for worker in workers), return_exceptions=True) + + +async def start_preview_workers() -> None: + """Warm up persistent preview workers during server startup.""" + global _preview_pool + if _preview_pool is not None: + return + async with _preview_pool_lock: + if _preview_pool is not None: + return + pool = _PreviewWorkerPool(PREVIEW_WORKERS) + await pool.start() + _preview_pool = pool + logger.info("Started %d persistent preview workers", PREVIEW_WORKERS) + + +async def shutdown_preview_workers() -> None: + """Kill persistent preview workers (called during server shutdown).""" + global _preview_pool + async with _preview_pool_lock: + pool = _preview_pool + _preview_pool = None + if pool is not None: + await pool.close() + if not _active_procs: + return + for proc in list(_active_procs): + try: + proc.kill() + except ProcessLookupError: + pass + await asyncio.gather(*(proc.wait() for proc in list(_active_procs)), return_exceptions=True) + _active_procs.clear() + @bp.on_request async def verify_preview(request): @@ -77,6 +273,24 @@ async def verify_preview(request): await auth.verify(request) +class PreviewTimeout(Exception): + """Raised when the preview subprocess exceeds PREVIEW_TIMEOUT.""" + + +class PreviewError(Exception): + """Raised when the preview subprocess exits with a non-zero status.""" + + +async def _run_preview_process( + filepath, quality: int, maxsize: int, maxzoom: float +) -> tuple[bytes | None, PreviewResponse | None]: + """Run preview request in a persistent worker process.""" + await start_preview_workers() + if _preview_pool is None: + raise PreviewError(filepath.name) + return await _preview_pool.run(filepath, quality, maxsize, maxzoom) + + # Map EXIF Orientation value to a corresponding PIL transpose EXIF_ORI = { 2: Image.Transpose.FLIP_LEFT_RIGHT, @@ -89,6 +303,19 @@ EXIF_ORI = { } +DOC_PREVIEW_SUFFIXES = {".pdf", ".xps", ".epub", ".mobi"} + + +def is_previewable_path(path) -> bool: + suffix = path.suffix.lower() + if suffix in DOC_PREVIEW_SUFFIXES: + return True + mime_type, _ = mimetypes.guess_type(path.name) + if not mime_type: + return False + return mime_type.startswith("image/") or mime_type.startswith("video/") + + @bp.get("/") async def preview(req, path): """Preview a file""" @@ -102,6 +329,9 @@ async def preview(req, path): except FileNotFoundError: raise NotFound() from None + if not is_previewable_path(filepath): + return empty(415) + etag = config.derived_secret( "preview", rel, stat.st_mtime_ns, quality, maxsize, maxzoom ).hex() @@ -117,21 +347,34 @@ async def preview(req, path): return raw(cached.body, headers=cached.headers) # Generate preview - img = await asyncio.get_event_loop().run_in_executor( - req.app.ctx.threadexec, dispatch, filepath, quality, maxsize, maxzoom - ) + try: + img, preview_resp = await _run_preview_process(filepath, quality, maxsize, maxzoom) + except PreviewTimeout: + return empty(504) + except PreviewError: + return empty(422) + if preview_resp and preview_resp.backend: + load_ms = int(round(preview_resp.load_ms or 0.0)) + process_ms = int(round(preview_resp.process_ms or 0.0)) + save_ms = int(round(preview_resp.save_ms or 0.0)) + req.ctx._log_extra = f"{preview_resp.backend} {load_ms}/{process_ms}/{save_ms} =" if not img: # Preview generation failed, redirect to the file itself return redirect(f"/files/{path}", status=303) # Build headers and cache the full response + preview_mime = ( + preview_resp.mime + if preview_resp is not None and preview_resp.mime is not None + else "image/avif" + ) savename = PurePosixPath(filepath.name).with_suffix(".avif") headers = { "etag": etag, "last-modified": format_date_time(stat.st_mtime), "cache-control": "max-age=604800, immutable" + ("" if config.config.public else ", private"), - "content-type": "image/avif", + "content-type": preview_mime, "content-disposition": f"inline; filename*=UTF-8''{urllib.parse.quote(savename.as_posix())}", } _preview_cache.set(etag, CachedPreview(headers=headers, body=img)) @@ -141,19 +384,27 @@ async def preview(req, path): def dispatch(path, quality, maxsize, maxzoom): try: - if path.suffix.lower() in (".pdf", ".xps", ".epub", ".mobi"): + if path.suffix.lower() in DOC_PREVIEW_SUFFIXES: return process_pdf(path, quality=quality, maxsize=maxsize, maxzoom=maxzoom) - type, _ = mimetypes.guess_type(path.name) - if type and type.startswith("video/"): + mime_type, _ = mimetypes.guess_type(path.name) + if mime_type and mime_type.startswith("video/"): return process_video(path, quality=quality, maxsize=maxsize) - return process_image(path, quality=quality, maxsize=maxsize) + if mime_type and mime_type.startswith("image/"): + return process_image(path, quality=quality, maxsize=maxsize) except ValueError as e: logger.warning(f"Cannot generate preview for {path}: {e}") except Exception as e: logger.exception(f"Error generating preview for {path}: {e}") + return None, PreviewResponse(ok=False) def process_image(path, *, maxsize, quality): + return process_image_with_timing(path, maxsize=maxsize, quality=quality) + + +def process_image_with_timing(path, *, maxsize, quality): + if FORCE_PIL: + return process_image_pillow(path, maxsize=maxsize, quality=quality) try: return process_image_pyvips(path, maxsize=maxsize, quality=quality) except Exception as e: @@ -174,22 +425,27 @@ def process_image_pyvips(path, *, maxsize, quality): img = img.resize(scale) t_save = perf_counter() - ret = img.write_to_buffer(".avif", Q=quality) + ret = img.write_to_buffer( + ".avif", + Q=quality, + effort=AVIF_FAST_EFFORT, + strip=True, + ) t_end = perf_counter() load_ms = (t_proc - t_load) * 1000 proc_ms = (t_save - t_proc) * 1000 save_ms = (t_end - t_save) * 1000 - logger.debug( - "Preview image %s via pyvips: load=%.1fms process=%.1fms save=%.1fms", - path.name, - load_ms, - proc_ms, - save_ms, + return ret, PreviewResponse( + ok=True, + mime="image/avif", + backend="pyvips", + load_ms=round(load_ms, 1), + process_ms=round(proc_ms, 1), + save_ms=round(save_ms, 1), + total_ms=round((t_end - t_load) * 1000, 1), ) - return ret - def process_image_pillow(path, *, maxsize, quality): t_load = perf_counter() @@ -222,16 +478,16 @@ def process_image_pillow(path, *, maxsize, quality): load_ms = (t_proc - t_load) * 1000 proc_ms = (t_save - t_proc) * 1000 save_ms = (t_end - t_save) * 1000 - logger.debug( - "Preview image %s via Pillow: load=%.1fms process=%.1fms save=%.1fms", - path.name, - load_ms, - proc_ms, - save_ms, + return ret, PreviewResponse( + ok=True, + mime="image/avif", + backend="pillow", + load_ms=round(load_ms, 1), + process_ms=round(proc_ms, 1), + save_ms=round(save_ms, 1), + total_ms=round((t_end - t_load) * 1000, 1), ) - return ret - def process_pdf(path, *, maxsize, maxzoom, quality, page_number=0): t_load_start = perf_counter() @@ -249,13 +505,14 @@ def process_pdf(path, *, maxsize, maxzoom, quality, page_number=0): ) t_save_end = perf_counter() - logger.debug( - "Preview pdf %s: load+render=%.1fms save=%.1fms", - path.name, - (t_load_end - t_load_start) * 1000, - (t_save_end - t_save_start) * 1000, + return ret, PreviewResponse( + ok=True, + mime="image/avif", + backend="pdf", + load_ms=round((t_load_end - t_load_start) * 1000, 1), + save_ms=round((t_save_end - t_save_start) * 1000, 1), + total_ms=round((t_save_end - t_load_start) * 1000, 1), ) - return ret def process_video(path, *, maxsize, quality): @@ -268,7 +525,13 @@ def process_video(path, *, maxsize, quality): t_save_start = t_load_start t_save_end = t_load_start with ( - av.open(str(path)) as icontainer, + av.open( + str(path), + options={ + "analyzeduration": "1000000", # 1 second (in microseconds) + "fflags": "fastseek", + }, + ) as icontainer, av.open(imgdata, "w", format="avif") as ocontainer, ): istream = icontainer.streams.video[0] @@ -360,14 +623,16 @@ def process_video(path, *, maxsize, quality): ocontainer.mux(ostream.encode(None)) # Flush the stream t_save_end = perf_counter() - # Capture frame dimensions before cleanup + # Capture result before cleanup ret = imgdata.getvalue() - logger.debug( - "Preview video %s: load+decode=%.1fms save=%.1fms", - path.name, - (t_load_end - t_load_start) * 1000, - (t_save_end - t_save_start) * 1000, + resp = PreviewResponse( + ok=True, + mime="image/avif", + backend="video", + load_ms=round((t_load_end - t_load_start) * 1000, 1), + save_ms=round((t_save_end - t_save_start) * 1000, 1), + total_ms=round((t_save_end - t_load_start) * 1000, 1), ) del imgdata, istream, ostream, icc, occ, frame gc.collect() - return ret + return ret, resp diff --git a/cista/preview_worker.py b/cista/preview_worker.py new file mode 100644 index 0000000..945a432 --- /dev/null +++ b/cista/preview_worker.py @@ -0,0 +1,95 @@ +"""Preview generation worker subprocess. + +Two modes are supported: +1) Legacy one-shot mode: argv has path/quality/maxsize/maxzoom. +2) Long-lived mode: read JSONL commands from stdin and write framed responses. + +Framed response format: + (blake3(packet))(uint32 json size)(uint32 payload size)(json)(binary payload) +where packet = (uint32 json size)(uint32 payload size)(json)(binary payload). +""" +import logging +import struct +import sys +from pathlib import Path + +import msgspec +from blake3 import blake3 + + +class PreviewRequest(msgspec.Struct, omit_defaults=True): + path: str + quality: int + maxsize: int + maxzoom: float + + +class PreviewResponse(msgspec.Struct, omit_defaults=True): + ok: bool + mime: str | None = None + backend: str | None = None + load_ms: float | None = None + process_ms: float | None = None + save_ms: float | None = None + total_ms: float | None = None + error: str | None = None + + +_enc = msgspec.json.Encoder() +_dec_req = msgspec.json.Decoder(PreviewRequest) + + +def _write_response(resp: PreviewResponse, payload: bytes) -> None: + meta_bytes = _enc.encode(resp) + packet = struct.pack(" None: + if len(sys.argv) != 5: + sys.stderr.write( + f"Usage: {sys.argv[0]} \n" + ) + sys.exit(1) + + from cista.preview import dispatch + + path = Path(sys.argv[1]) + quality = int(sys.argv[2]) + maxsize = int(sys.argv[3]) + maxzoom = float(sys.argv[4]) + result, _ = dispatch(path, quality, maxsize, maxzoom) + if result: + sys.stdout.buffer.write(result) + sys.stdout.buffer.flush() + + +def _run_loop() -> None: + from cista.preview import dispatch + + while True: + line = sys.stdin.buffer.readline() + if not line: + return + try: + req = _dec_req.decode(line) + result, resp = dispatch(Path(req.path), req.quality, req.maxsize, req.maxzoom) + _write_response(resp, result or b"") + except Exception as e: + _write_response(PreviewResponse(ok=False, error=str(e)), b"") + + +def main() -> None: + # Configure all log output to stderr before any imports that may emit logs. + logging.basicConfig(stream=sys.stderr, level=logging.INFO) + if len(sys.argv) > 1: + _run_once() + return + _run_loop() + + +if __name__ == "__main__": + main()