pdf/gateway/app/services/convert/doc_cache.py

"""Per-conversion cache for expensive PDF reads.

``PdfReader.extract_text`` is the dominant cost of a text-heavy conversion, and
three independent stages need the same text: the document router deciding which
pages need OCR, the layout pipeline building the IDM, and the quality scorer
establishing a source baseline. Each stage opened its own reader, so a 14-page
document ran extraction 42 times.

The cache is bound to a conversion scope, so entries are released when the
conversion ends and two concurrent conversions never share state. Outside a
scope every call falls through to a direct read, which keeps callers total and
makes the cache a pure optimisation.
"""

from __future__ import annotations

import hashlib
import io
from collections import OrderedDict
from collections.abc import Iterator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any

from pypdf import PdfReader

# Hashing the whole payload is cheap next to a single page extraction
# (~1ms/MB against ~130ms/page), and keying on content means a converter that
# handles more than one document inside a scope stays correct.
_DIGEST_BYTES = 8


def _key(data: bytes) -> str:
    return hashlib.blake2b(data, digest_size=_DIGEST_BYTES).hexdigest()


@dataclass
class DocumentCache:
    """Readers and extracted page text for the documents seen in one scope."""

    readers: dict[str, PdfReader] = field(default_factory=dict)
    page_text: dict[tuple[str, int], str] = field(default_factory=dict)
    # Glyph boxes and vector path operators recovered from the content stream
    # when the C++ engine is absent. Tokenising a content stream and resolving
    # its fonts costs about as much as extracting the page's text, and the
    # layout pipeline, table detector and scorer all ask for the same answer.
    geometry: OrderedDict[tuple[str, int], tuple[list, list]] = field(
        default_factory=OrderedDict
    )
    # Rendered page images keyed by (document, page, dpi). A page is rasterised
    # once and reused by layout detection, OCR and image export.
    rasters: OrderedDict[tuple[str, int, int], bytes] = field(default_factory=OrderedDict)
    # The engine document handle, opened at most once per document.
    engines: dict[str, object] = field(default_factory=dict)
    hits: int = 0
    misses: int = 0
    raster_hits: int = 0
    raster_misses: int = 0
    raster_bytes: int = 0

    def reader_for(self, data: bytes) -> PdfReader:
        key = _key(data)
        reader = self.readers.get(key)
        if reader is None:
            reader = PdfReader(io.BytesIO(data), strict=False)
            self.readers[key] = reader
        return reader

    def text_for(self, data: bytes, index: int) -> str:
        key = (_key(data), index)
        cached = self.page_text.get(key)
        if cached is not None:
            self.hits += 1
            return cached
        self.misses += 1
        reader = self.reader_for(data)
        try:
            text = reader.pages[index].extract_text() or ""
        except Exception:
            text = ""
        self.page_text[key] = text
        return text

    def geometry_for(self, data: bytes, index: int, extract) -> tuple[list, list]:
        """``(glyphs, path_ops)`` for one page, computed at most once.

        Bounded to the most recent pages. Every consumer of a page's glyphs —
        the layout pipeline, the table detectors, the scorer — asks for them
        while that page is being built and never again, so holding all 126
        pages' worth to the end of the conversion buys nothing and costs the
        peak. Evicting is always safe: a later ask recomputes.
        """
        key = (_key(data), index)
        cached = self.geometry.get(key)
        if cached is not None:
            self.hits += 1
            # Refresh recency so a page still being worked on is not evicted.
            self.geometry.move_to_end(key)
            return cached
        self.misses += 1
        try:
            result = extract()
        except Exception:
            result = ([], [])
        self.geometry[key] = result
        while len(self.geometry) > MAX_CACHED_GEOMETRY_PAGES:
            self.geometry.popitem(last=False)
        return result

    def raster_for(self, data: bytes, index: int, dpi: int, render) -> bytes:
        """Return a rendered page, computing it at most once per (page, dpi).

        Rasters are the largest thing here by an order of magnitude, and a
        page's raster is wanted by at most three consecutive stages — layout
        detection, OCR, image export. Once the cache is full the oldest page
        is dropped rather than the new one refused, so the pages being worked
        on now are the ones held.
        """
        key = (_key(data), index, int(dpi))
        cached = self.rasters.get(key)
        if cached is not None:
            self.raster_hits += 1
            self.rasters.move_to_end(key)
            return cached
        self.raster_misses += 1
        png = render()
        if len(png) <= MAX_RASTER_CACHE_BYTES:
            self.rasters[key] = png
            self.raster_bytes += len(png)
            self.rasters.move_to_end(key)
            while self.raster_bytes > MAX_RASTER_CACHE_BYTES and len(self.rasters) > 1:
                _evicted, payload = self.rasters.popitem(last=False)
                self.raster_bytes -= len(payload)
        return png

    def engine_for(self, data: bytes, open_document) -> object | None:
        """Open the engine document once per conversion."""
        key = _key(data)
        if key in self.engines:
            return self.engines[key]
        doc = open_document()
        self.engines[key] = doc
        return doc


# Page rasters at 200 dpi run 1-3 MB each. A page's raster is wanted by at
# most three consecutive stages, so a cache large enough to hold a hundred of
# them is a hundred pages of peak bought for nothing. 64 MB covers the working
# set with room to spare and takes a third of a gigabyte off a long scan.
MAX_RASTER_CACHE_BYTES = 64 * 1024 * 1024
# Glyph lists are small next to a raster but not free: a dense page carries a
# few thousand dicts, and a 126-page booklet held every one of them to the end
# of the conversion. Twelve pages is far more than any stage looks back.
MAX_CACHED_GEOMETRY_PAGES = 12

_cache: ContextVar[DocumentCache | None] = ContextVar("convert_doc_cache", default=None)


@contextmanager
def document_cache_scope() -> Iterator[DocumentCache]:
    """Cache PDF reads for the duration of one conversion."""
    cache = DocumentCache()
    token = _cache.set(cache)
    try:
        yield cache
    finally:
        _cache.reset(token)


def current() -> DocumentCache | None:
    return _cache.get()


def get_reader(data: bytes) -> PdfReader:
    """A ``PdfReader`` for *data*, reused within the current conversion."""
    cache = _cache.get()
    if cache is None:
        return PdfReader(io.BytesIO(data), strict=False)
    return cache.reader_for(data)


def page_text(data: bytes, index: int) -> str:
    """Extracted text for one page, computed at most once per conversion."""
    cache = _cache.get()
    if cache is None:
        try:
            reader = PdfReader(io.BytesIO(data), strict=False)
            return reader.pages[index].extract_text() or ""
        except Exception:
            return ""
    return cache.text_for(data, index)


def page_geometry(data: bytes, index: int, extract) -> tuple[list, list]:
    """Glyph boxes and path operators for one page, reused within a conversion."""
    cache = _cache.get()
    if cache is None:
        try:
            return extract()
        except Exception:
            return ([], [])
    return cache.geometry_for(data, index, extract)


def page_raster(data: bytes, index: int, dpi: int, render) -> bytes:
    """Rendered page image, reused across layout detection, OCR and export."""
    cache = _cache.get()
    if cache is None:
        return render()
    return cache.raster_for(data, index, dpi, render)


def engine_document(data: bytes, open_document):
    """Engine handle for *data*, opened at most once per conversion."""
    cache = _cache.get()
    if cache is None:
        return open_document()
    return cache.engine_for(data, open_document)


def page_count(data: bytes) -> int:
    try:
        return len(get_reader(data).pages)
    except Exception:
        return 0


def all_text(data: bytes) -> str:
    """Whole-document plain text, reusing any per-page text already extracted."""
    count = page_count(data)
    return "\n".join(page_text(data, i) for i in range(count))


def stats() -> dict[str, Any]:
    cache = _cache.get()
    if cache is None:
        return {}
    return {
        "documents": len(cache.readers),
        "pages_cached": len(cache.page_text),
        "geometry_cached": len(cache.geometry),
        "hits": cache.hits,
        "misses": cache.misses,
        "rasters_cached": len(cache.rasters),
        "raster_hits": cache.raster_hits,
        "raster_misses": cache.raster_misses,
        "raster_mb": round(cache.raster_bytes / 1024 / 1024, 1),
    }