data/cache.py

"""Byte-bounded disk cache and HTTP range reader.

Shared verbatim with the OCR half (06-multilingual/ocr): the caching problem is identical once the unit
is "a parquet row group", and a divergent copy would be a second thing to get right.
"""

import fcntl
import hashlib
import io
import os
import threading
import time
import uuid
from collections import Counter
from contextlib import contextmanager
from pathlib import Path

import requests


class DiskCache:
    def __init__(self, directory, max_bytes, max_entry_bytes=None):
        self.directory = Path(directory)
        self.directory.mkdir(parents=True, exist_ok=True)
        self.max_bytes = int(max_bytes)
        self.max_entry_bytes = min(int(max_entry_bytes or max_bytes), self.max_bytes)
        if self.max_entry_bytes < 1:
            raise ValueError("Cache budgets must be positive")
        self.counters = Counter()
        self._stats_lock = threading.Lock()
        # Recover staging space after a process was killed before its finally block.
        # A live writer owns this stripe, including writers in another process.
        for path in self.directory.glob("*.part"):
            stripe = path.name[:3]
            with (self.directory / f"stripe-{stripe}.lock").open("a+b") as lock:
                try:
                    fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
                except BlockingIOError:
                    continue
                path.unlink(missing_ok=True)

    def increment(self, key, value=1):
        with self._stats_lock:
            self.counters[key] += value

    def stats(self):
        with self._stats_lock:
            result = dict(self.counters)
        files = list(self.directory.glob("*.bin"))
        sizes = []
        for path in files:
            try:
                sizes.append(path.stat().st_size)
            except FileNotFoundError:
                pass
        return {
            **result,
            "bytes": sum(sizes),
            "entries": len(sizes),
            "max_bytes": self.max_bytes,
            "max_entry_bytes": self.max_entry_bytes,
        }

    @contextmanager
    def lease(self, key, loader):
        digest = hashlib.sha256(key.encode()).hexdigest()
        path = self.directory / f"{digest}.bin"
        # Fixed lock stripes avoid millions of permanent lock files over full epochs.
        # A stripe covers both loading and reading; collisions only serialize work.
        stripe = digest[:3]
        with (self.directory / f"stripe-{stripe}.lock").open("a+b") as lock:
            fcntl.flock(lock, fcntl.LOCK_EX)
            if path.exists():
                self.increment("hits")
            else:
                self.increment("misses")
                temporary = self.directory / f"{digest}.{uuid.uuid4().hex}.part"
                try:
                    loader(temporary)
                    size = temporary.stat().st_size
                    if size > self.max_entry_bytes:
                        raise ValueError(
                            f"Cache entry {size} exceeds {self.max_entry_bytes} byte limit"
                        )
                    with (self.directory / ".eviction.lock").open("a+b") as eviction:
                        fcntl.flock(eviction, fcntl.LOCK_EX)
                        self._make_room(size, stripe)
                        os.replace(temporary, path)
                    self.increment("loads")
                    self.increment("loaded_bytes", size)
                except BaseException:
                    self.increment("load_failures")
                    raise
                finally:
                    temporary.unlink(missing_ok=True)
            os.utime(path, None)
            yield path

    def _make_room(self, needed, owned_stripe):
        entries = [
            (p.stat().st_mtime_ns, p.stat().st_size, p)
            for p in self.directory.glob("*.bin")
        ]
        used = sum(size for _, size, _ in entries)
        for _, size, path in sorted(entries):
            if used + needed <= self.max_bytes:
                return
            stripe = path.stem[:3]
            if stripe == owned_stripe:
                # Our exclusive stripe already proves no other reader holds this entry.
                path.unlink(missing_ok=True)
                used -= size
                self.increment("evictions")
                continue
            with (self.directory / f"stripe-{stripe}.lock").open("a+b") as lock:
                try:
                    fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
                except BlockingIOError:
                    continue
                path.unlink(missing_ok=True)
                used -= size
                self.increment("evictions")
        if used + needed > self.max_bytes:
            raise RuntimeError(
                "Cache budget is occupied by active readers; reduce prefetch or increase cache bytes"
            )


class RangeReader(io.RawIOBase):
    """Seekable HTTP object: require exact 206 responses and meter actual payload bytes."""

    def __init__(self, url, size, *, headers=None, on_read=None, timeout=120):
        super().__init__()
        self.url, self.size, self.position = url, size, 0
        self.headers = headers or {}
        self.on_read = on_read or (lambda amount: None)
        self.timeout = timeout
        self.session = requests.Session()
        self._lock = threading.RLock()

    def readable(self):
        return True

    def seekable(self):
        return True

    def tell(self):
        return self.position

    def seek(self, offset, whence=0):
        with self._lock:
            position = (
                offset
                if whence == 0
                else self.position + offset
                if whence == 1
                else self.size + offset
            )
            if position < 0 or whence not in (0, 1, 2):
                raise ValueError("Invalid seek")
            self.position = position
            return position

    def read(self, size=-1):
        with self._lock:
            start = self.position
            stop = self.size if size < 0 else min(start + size, self.size)
            if start >= stop:
                return b""
            for attempt in range(3):
                try:
                    with self.session.get(
                        self.url,
                        headers={**self.headers, "Range": f"bytes={start}-{stop - 1}"},
                        timeout=(15, self.timeout),
                        stream=True,
                    ) as response:
                        response.raise_for_status()
                        if response.status_code != 206:
                            raise ValueError(
                                "Storage ignored byte range; refusing a whole-shard download"
                            )
                        if (
                            response.headers.get("Content-Range")
                            != f"bytes {start}-{stop - 1}/{self.size}"
                        ):
                            raise ValueError(
                                "Storage returned an unexpected byte range or object size"
                            )
                        data = bytearray()
                        for chunk in response.iter_content(1024 * 1024):
                            self.on_read(len(chunk))
                            data.extend(chunk)
                            if len(data) > stop - start:
                                raise ValueError(
                                    "Storage returned more bytes than requested"
                                )
                        if len(data) != stop - start:
                            raise requests.ConnectionError(
                                "Incomplete byte-range response"
                            )
                    self.position = stop
                    return bytes(data)
                except (requests.ConnectionError, requests.Timeout):
                    if attempt == 2:
                        raise
                    time.sleep(0.25 * (attempt + 1))

    def readinto(self, buffer):
        value = self.read(len(buffer))
        buffer[: len(value)] = value
        return len(value)

    def close(self):
        self.session.close()
        super().close()