remote/remote_fallback.py

"""Remote (HTTP) ChessTB tablebase support for the standard
``chess.chesstb`` module (noobpwnftw's python-chess fork, see
requirements.txt), selected through tablebase_router when
``config.REMOTE_MODE`` is ``"download"``.

Why this exists
----------------
``chess.chesstb.Tablebase`` only ever reads from a local filesystem path:
``_TableFile._open_source()`` resolves a real filesystem path with
``os.open()`` and memory-maps it (``mmap.mmap(..., access=mmap.ACCESS_READ)``),
and ``Tablebase._find()`` always checks candidates on the local filesystem.
``_open_source`` is itself a documented seam another transport can
replace -- ``remote_direct.py`` does exactly that, reading HTTP byte
ranges on demand. This module takes the simpler of the two
approaches: it leaves ``_open_source`` untouched and overrides only
``_find``, described below.

What this module does instead is subclass ``chess.chesstb.Tablebase`` and
override only the file-resolution step (``_find``): the first time a
probe touches a given material, the corresponding remote table file is
fetched in full and written to a local per-process cache directory, and
the resulting local path is handed unchanged to the standard, untouched
``WDLFile``/``DTZFile``/``DTCFile``/``DTMFile``/``DTM50File`` constructors --
which run completely unaware they're not reading from a hand-picked local
directory. Every
other behaviour (probing, index maths, block decoding, ``MissingTableError``,
the ``WIN``/``DRAW``/... constants, ``ProbeResult``) is the standard
module's own, unmodified implementation -- this module only changes *how a
table's bytes reach disk*, never what they mean once there.

A whole table file is downloaded the first time any position touches
that material, rather than only the specific byte ranges (header/index
plus whichever compressed blocks the position's index falls into) a probe
actually needs -- there's no way to get bytes in front of code that only
ever calls ``open()`` on a path other than putting real bytes at that
path first (see ``remote_direct.py`` for the byte-range alternative,
which reads through ``remote_source.RemoteFileView`` instead). Once
cached, later probes against that material are served straight
from local disk with no further network round trip. Concretely this
means: slower first touch of a given material (a full download instead
of a few small ranged fetches), fast disk-speed reads for everything
after, and real, if bounded and temporary, local disk usage. See
``_RemoteDiskCache`` below for the bound (``REMOTE_PAGE_CACHE_BYTES``
doing double duty as an on-disk budget here rather than an in-memory one)
and ``open_tablebase``/``Tablebase.close`` for the "temporary": the cache
lives in a fresh directory created per process and is removed when the
``Tablebase`` is closed (app.py does this at interpreter exit via
``atexit``, same as ``TB.close()``).

This module does not add any locking of its own around opening a
not-yet-seen remote material. ``SourcedTablebase`` holds upstream's
per-kind lock around ``_find`` whenever a download-mode source is
configured, so two threads racing to probe the same never-before-seen
material collapse onto one download. Materials of the same kind serialize
behind each other's download; that is a performance-only characteristic,
not a correctness gap.
"""
from __future__ import annotations

import collections
import os
import shutil
import tempfile
import threading
from typing import Optional

import chess.chesstb as chesstb

from tablebase_opening import block_cache_usage

from . import remote_source

__all__ = ["open_tablebase"]


class _RemoteDiskCache:
    """Soft-budgeted LRU of whole downloaded table files on local disk,
    keyed by the remote-relative path each was fetched from.

    Structurally mirrors one shard of :class:`remote_source._PageCache`
    one layer up the stack: same OrderedDict-as-LRU-ledger, same soft
    byte budget, same per-key lock so concurrent probe threads landing on
    the same not-yet-cached material download it once rather than once
    each -- whole files here instead of fixed-size pages, since
    that's the unit the standard backend's file constructors can consume.

    A freshly recorded (or freshly touched, via ``has()``) entry always
    ends up at the end of the LRU order, and eviction always pops from
    the front -- so the file ``_find()`` resolved a path for can
    never be the one an eviction happening on another thread picks next;
    only entries stale relative to it can be. That closes the same race
    :meth:`remote_source._PageCacheShard.put` avoids for in-memory pages: a
    long-untouched file could still be unlinked out from under a reader
    holding an OS-level file descriptor already open on it, but on POSIX
    an unlinked-but-open file stays readable to that descriptor until
    it is closed, so an in-flight ``open().read()`` (the only operation the
    standard ``WDLFile``/``DTZFile``/``DTCFile``/``DTMFile``/``DTM50File``
    ever do) is unaffected either way.
    """

    def __init__(self, root: str, max_bytes: int) -> None:
        self.root = root
        self.max_bytes = max_bytes
        self._sizes: "collections.OrderedDict[str, int]" = collections.OrderedDict()
        self._cur_bytes = 0
        self._lock = threading.Lock()
        self._entry_locks: dict[str, threading.Lock] = {}
        self._meta_lock = threading.Lock()
        # Guards the install-and-record tail of _ensure_cached() below
        # (the os.replace() + record() that actually mutates `root`'s
        # contents and this cache's bookkeeping together) against
        # clear()'s rmtree/recreate running at the same moment. A
        # per-entry lock (see lock_for()) already serializes concurrent
        # downloads of the same rel_path; this instead protects the
        # publish step against a whole-cache clear(). Reentrant because
        # clear() below nests other locks (_lock, _meta_lock) inside it.
        self.lifecycle_lock = threading.RLock()
        # Counts downloads that currently have an open, not-yet-published
        # temp file on disk. clear() waits for this to hit zero before
        # rmtree-ing `root` -- see begin_download()/end_download() and
        # clear() below for why: unlike lifecycle_lock, which only needs
        # to cover the brief publish step, a download's temp file is
        # vulnerable for its entire (potentially slow) transfer, not
        # that tail. On POSIX, rmtree-ing a directory containing a
        # file that's still open only unlinks its name -- the data stays
        # readable/writable through the already-open descriptor -- but
        # once that descriptor closes (as it does at the end of
        # _ensure_cached()'s ``with open(tmp_path, "wb") as f:`` block,
        # before publish), an unlinked file's data is actually freed, so
        # the publish step's os.replace(tmp_path, ...) would fail with
        # FileNotFoundError for a missing source path
        # anywhere on disk. Deferring clear() instead of racing it avoids
        # that outright, and -- unlike holding lifecycle_lock for the
        # whole transfer -- doesn't serialize unrelated concurrent
        # downloads against each other, only against clear() itself.
        self._active_downloads = 0
        self._no_active_downloads = threading.Condition(self.lifecycle_lock)

    def begin_download(self) -> None:
        """Mark one download as having an in-flight temp file. Must be
        paired with a later :meth:`end_download`, typically via
        ``try/finally``."""
        with self.lifecycle_lock:
            self._active_downloads += 1

    def end_download(self) -> None:
        with self.lifecycle_lock:
            self._active_downloads -= 1
            if self._active_downloads <= 0:
                self._no_active_downloads.notify_all()

    def lock_for(self, rel_path: str) -> threading.Lock:
        lk = self._entry_locks.get(rel_path)
        if lk is not None:
            return lk
        with self._meta_lock:
            lk = self._entry_locks.get(rel_path)
            if lk is None:
                lk = threading.Lock()
                self._entry_locks[rel_path] = lk
            return lk

    def local_path(self, rel_path: str) -> str:
        return os.path.join(self.root, rel_path)

    def has(self, rel_path: str) -> bool:
        """True if `rel_path` is already cached on disk -- and, if so,
        mark it most-recently-used so a same-moment eviction on another
        thread cannot acquire it (see class docstring)."""
        with self._lock:
            if rel_path in self._sizes:
                self._sizes.move_to_end(rel_path)
                return True
        return False

    def record(self, rel_path: str, size: int) -> None:
        with self._lock:
            old = self._sizes.pop(rel_path, None)
            if old is not None:
                self._cur_bytes -= old
            self._sizes[rel_path] = size
            self._cur_bytes += size
            # Keep the newly added file (it is at the end); never evict fully
            # empty, mirroring remote_source._PageCacheShard.put / chesstb.py's
            # own _BlockCache.record.
            while self._cur_bytes > self.max_bytes and len(self._sizes) > 1:
                ev_path, ev_size = self._sizes.popitem(last=False)
                self._cur_bytes -= ev_size
                try:
                    os.remove(self.local_path(ev_path))
                except OSError:
                    pass

    def clear(self) -> None:
        with self.lifecycle_lock:
            while self._active_downloads > 0:
                self._no_active_downloads.wait()
            with self._lock:
                self._sizes.clear()
                self._cur_bytes = 0
            with self._meta_lock:
                self._entry_locks.clear()
            shutil.rmtree(self.root, ignore_errors=True)
            os.makedirs(self.root, exist_ok=True)

    def stats(self) -> dict:
        with self._lock:
            return {
                "cached_files": len(self._sizes),
                "cur_bytes":    self._cur_bytes,
                "max_bytes":    self.max_bytes,
            }


class _RemoteCachingTablebase(chesstb.Tablebase):
    """A standard ``chess.chesstb.Tablebase`` whose table files are
    fetched from a remote ``http(s)://`` base URL and disk-cached on
    first touch, rather than read from a local directory. See this
    module's docstring for the full design and trade-offs.
    """

    def __init__(self, base_url: str, *, block_cache_bytes: int,
                 remote_page_cache_bytes: int, remote_page_size: int,
                 remote_timeout: float, remote_max_retries: int,
                 remote_pool_maxsize: int) -> None:
        self._client = remote_source.RemoteHTTPClient(
            base_url, timeout=remote_timeout, max_retries=remote_max_retries,
            pool_maxsize=remote_pool_maxsize,
        )
        self._cache_root = tempfile.mkdtemp(prefix="chesstb_remote_cache_")
        self._disk_cache = _RemoteDiskCache(self._cache_root, remote_page_cache_bytes)
        self._download_chunk = remote_page_size
        self._size_cache: dict[str, Optional[int]] = {}
        self._size_lock = threading.Lock()
        # Base Tablebase.__init__ calls self.add_directory(base_url) --
        # overridden to a no-op below, since _find here resolves directly
        # against "<kind>/" and never falls back to a bare-root candidate
        # the way local _find does -- the published remote layout always
        # has the kind subdirectories, per README.md's "Getting the
        # tablebase files".
        super().__init__(base_url, block_cache_bytes=block_cache_bytes)

    def add_directory(self, base_dir: str) -> None:
        """No-op -- see the comment in __init__ above."""

    # --- table file resolution: fetch-and-cache-to-disk, then delegate ---

    def _find(self, kind: str, name: str, ext: str) -> Optional[str]:
        rel_path = f"{kind}/{name}{ext}"
        return self._ensure_cached(rel_path)

    def _remote_size(self, rel_path: str) -> Optional[int]:
        with self._size_lock:
            if rel_path in self._size_cache:
                return self._size_cache[rel_path]
        size = self._client.head_size(rel_path)  # may raise RemoteSourceError
        with self._size_lock:
            self._size_cache[rel_path] = size
        return size

    def _ensure_cached(self, rel_path: str) -> Optional[str]:
        if self._disk_cache.has(rel_path):
            return self._disk_cache.local_path(rel_path)
        with self._disk_cache.lock_for(rel_path):
            if self._disk_cache.has(rel_path):
                return self._disk_cache.local_path(rel_path)
            size = self._remote_size(rel_path)
            if size is None:
                return None
            local_path = self._disk_cache.local_path(rel_path)
            tmp_path = f"{local_path}.part"
            # begin_download()/end_download() make a concurrent
            # /admin/cache/clear() wait for this download rather than
            # racing it -- see _RemoteDiskCache's docstring on
            # _active_downloads for why racing it (rather than
            # deferring) isn't safe to do here. This does not serialize
            # concurrent downloads of different tables against each
            # other -- only clear() waits, not other downloads -- so
            # probe concurrency (see config.py's PROBE_THREADS) is
            # unaffected; the per-path lock_for(rel_path) above already
            # prevents two threads from downloading the same rel_path at
            # once, which is the actual duplicate-work hazard here.
            self._disk_cache.begin_download()
            try:
                os.makedirs(os.path.dirname(local_path), exist_ok=True)
                try:
                    with open(tmp_path, "wb") as f:
                        offset = 0
                        while offset < size:
                            remaining = size - offset
                            chunk = self._client.get_range(
                                rel_path, offset, min(self._download_chunk, remaining)
                            )
                            if not chunk:
                                raise remote_source.RemoteSourceError(
                                    f"Incomplete download for {rel_path}: received "
                                    f"{offset} of {size} bytes"
                                )
                            if len(chunk) > remaining:
                                raise remote_source.RemoteSourceError(
                                    f"Oversized range response for {rel_path}: "
                                    f"{len(chunk)} bytes requested max {remaining}"
                                )
                            f.write(chunk)
                            offset += len(chunk)
                    # Publication -- the only part that actually needs
                    # `root` to still exist and mean what it did when this
                    # method started -- is what lifecycle_lock protects.
                    # Only reached once the loop above has written exactly
                    # `size` bytes -- an atomic rename protects
                    # visibility, not content correctness, so a short
                    # download must never reach this line.
                    with self._disk_cache.lifecycle_lock:
                        os.replace(tmp_path, local_path)
                        self._disk_cache.record(rel_path, size)
                except BaseException:
                    try:
                        os.remove(tmp_path)
                    except OSError:
                        pass
                    raise
            finally:
                self._disk_cache.end_download()
            return local_path

    # --- lifecycle / admin surfaces (reached through SourcedTablebase's
    #     clear_caches()/cache_stats(), both getattr-guarded so this stays
    #     optional) ---

    def close(self) -> None:
        try:
            super().close()
        finally:
            try:
                self._client.close()
            finally:
                shutil.rmtree(self._cache_root, ignore_errors=True)

    def clear_caches(self) -> None:
        """Drop the decoded-block cache and the on-disk downloaded-file
        cache, without forgetting which tables are open -- re-opening
        tables here would defeat the point of a cache-clear (a fresh
        Tablebase would re-download every table on first touch)."""
        self._block_cache.clear()
        self._disk_cache.clear()

    def cache_stats(self) -> dict:
        blocks, block_bytes = block_cache_usage(self)
        return {
            "block_cache_blocks": blocks,
            "block_cache_bytes":  block_bytes,
            "remote_disk_cache":  self._disk_cache.stats(),
            "remote_http":        self._client.stats(),
        }


def open_tablebase(directory: str, *,
                    block_cache_bytes: int = chesstb.DEFAULT_BLOCK_CACHE_BYTES,
                    remote_page_cache_bytes: int = remote_source.DEFAULT_PAGE_CACHE_BYTES,
                    remote_page_size: int = remote_source.DEFAULT_PAGE_SIZE,
                    remote_timeout: float = remote_source.DEFAULT_TIMEOUT,
                    remote_max_retries: int = remote_source.DEFAULT_MAX_RETRIES,
                    remote_pool_maxsize: int = remote_source.DEFAULT_POOL_MAXSIZE,
                    ) -> _RemoteCachingTablebase:
    """Open a remote ChessTB base URL, downloading and disk-caching tables
    on first touch. Signature-compatible with ``remote_direct.open_tablebase``,
    so tablebase_router.open_one selects either without special-casing.
    """
    return _RemoteCachingTablebase(
        directory,
        block_cache_bytes=block_cache_bytes,
        remote_page_cache_bytes=remote_page_cache_bytes,
        remote_page_size=remote_page_size,
        remote_timeout=remote_timeout,
        remote_max_retries=remote_max_retries,
        remote_pool_maxsize=remote_pool_maxsize,
    )