Spaces:
Running
Running
| """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, | |
| ) | |