pdf/gateway/app/services/render_cache.py

import threading
import time
from dataclasses import dataclass
from enum import Enum
from typing import Any

from app.config import get_settings

RENDERER_VERSION: int = 1


class RenderMode(str, Enum):
    # Current
    NORMAL = "normal"
    PRINT = "print"
    GRAYSCALE = "grayscale"

    # Future (reserved)
    ANNOTATION_ONLY = "annotation_only"
    EDIT_PREVIEW = "edit_preview"
    SELECTION_OVERLAY = "selection_overlay"
    PROOF = "proof"
    HIGH_QUALITY = "high_quality"
    DRAFT = "draft"


@dataclass(frozen=True)
class TileCacheKey:
    doc_hash: str
    renderer_version: int
    page: int
    dpi: int
    zoom: float
    rotation: int
    render_mode: str
    tile_x: float
    tile_y: float
    tile_w: float
    tile_h: float


@dataclass
class CacheStats:
    hits: int = 0
    misses: int = 0
    evictions: int = 0
    current_entries: int = 0
    current_bytes: int = 0
    max_bytes: int = 0
    hit_rate: float = 0.0
    avg_lookup_ns: float = 0.0
    avg_render_time_ns: float = 0.0


class _Node:
    __slots__ = ["key", "value", "size", "prev", "next"]

    def __init__(self, key: TileCacheKey | None, value: bytes | None, size: int):
        self.key = key
        self.value = value
        self.size = size
        self.prev: _Node | None = None
        self.next: _Node | None = None


class TileCache:
    def __init__(self):
        settings = get_settings()
        self.enabled = settings.render_cache_enabled
        self.max_entries = settings.render_cache_max_entries
        self.max_bytes = settings.render_cache_max_bytes
        self.max_entry_bytes = settings.render_cache_max_entry_bytes

        self._lock = threading.Lock()
        self._cache: dict[TileCacheKey, _Node] = {}
        self._head = _Node(None, None, 0)
        self._tail = _Node(None, None, 0)
        self._head.next = self._tail
        self._tail.prev = self._head

        self._current_bytes = 0

        # Stats
        self._hits = 0
        self._misses = 0
        self._evictions = 0
        self._total_lookup_ns = 0
        self._total_render_time_ns = 0
        self._render_time_count = 0

    def _remove(self, node: _Node):
        p = node.prev
        n = node.next
        if p and n:
            p.next = n
            n.prev = p

    def _add_to_front(self, node: _Node):
        first = self._head.next
        if first:
            self._head.next = node
            node.prev = self._head
            node.next = first
            first.prev = node

    def _evict(self):
        last = self._tail.prev
        if last and last != self._head:
            self._remove(last)
            if last.key:
                del self._cache[last.key]
            self._current_bytes -= last.size
            self._evictions += 1

    def get(self, key: TileCacheKey) -> bytes | None:
        if not self.enabled:
            return None

        start_time = time.perf_counter_ns()
        with self._lock:
            node = self._cache.get(key)
            if node:
                self._hits += 1
                self._remove(node)
                self._add_to_front(node)
                res = node.value
            else:
                self._misses += 1
                res = None

            lookup_time = time.perf_counter_ns() - start_time
            self._total_lookup_ns += lookup_time

            return res

    def put(self, key: TileCacheKey, data: bytes) -> None:
        if not self.enabled:
            return

        size = len(data)
        if size > self.max_entry_bytes:
            # Too large to cache
            return

        with self._lock:
            if key in self._cache:
                node = self._cache[key]
                self._current_bytes -= node.size
                node.value = data
                node.size = size
                self._current_bytes += size
                self._remove(node)
                self._add_to_front(node)
            else:
                new_node = _Node(key, data, size)
                self._cache[key] = new_node
                self._add_to_front(new_node)
                self._current_bytes += size

            # Evict if over limits
            while len(self._cache) > self.max_entries or (self.max_bytes > 0 and self._current_bytes > self.max_bytes):
                self._evict()

    def invalidate_doc(self, doc_hash: str) -> None:
        with self._lock:
            keys_to_remove = [k for k in self._cache.keys() if k.doc_hash == doc_hash]
            for k in keys_to_remove:
                node = self._cache[k]
                self._remove(node)
                self._current_bytes -= node.size
                del self._cache[k]

    def invalidate_renderer_version(self, old_version: int) -> None:
        with self._lock:
            keys_to_remove = [k for k in self._cache.keys() if k.renderer_version == old_version]
            for k in keys_to_remove:
                node = self._cache[k]
                self._remove(node)
                self._current_bytes -= node.size
                del self._cache[k]

    def clear(self) -> None:
        with self._lock:
            self._cache.clear()
            self._head.next = self._tail
            self._tail.prev = self._head
            self._current_bytes = 0

    def stats(self) -> CacheStats:
        with self._lock:
            hit_rate = 0.0
            total_reqs = self._hits + self._misses
            if total_reqs > 0:
                hit_rate = self._hits / total_reqs

            avg_lookup = 0.0
            if total_reqs > 0:
                avg_lookup = self._total_lookup_ns / total_reqs

            avg_render = 0.0
            if self._render_time_count > 0:
                avg_render = self._total_render_time_ns / self._render_time_count

            return CacheStats(
                hits=self._hits,
                misses=self._misses,
                evictions=self._evictions,
                current_entries=len(self._cache),
                current_bytes=self._current_bytes,
                max_bytes=self.max_bytes,
                hit_rate=hit_rate,
                avg_lookup_ns=avg_lookup,
                avg_render_time_ns=avg_render,
            )

    def record_render_time(self, elapsed_ns: int):
        with self._lock:
            self._total_render_time_ns += elapsed_ns
            self._render_time_count += 1


tile_cache = TileCache()