import atexit
import os
import re
import shutil
import logging
import importlib.metadata
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import List, Optional, Union
try:
import pysciqlop_cache as sc
except ImportError: # pragma: no cover - platform-specific (WASM has no wheel)
# No compiled backend (e.g. WASM/Pyodide): fall back to a no-op cache so
# importing Speasy still works; caching is simply disabled.
from . import _noop_cache as sc
from . import _noop_cache
from .version import str_to_version, version_to_str, Version
from speasy.config import cache as cache_cfg
cache_version = str_to_version("3.0")
log = logging.getLogger(__name__)
def _release_at_exit(owner, attr: str) -> None:
"""Swap ``owner.<attr>`` (a native store) for a no-op one when the interpreter exits.
Module-level singletons are not guaranteed to be destroyed before the
extension module is torn down (e.g. when Python is embedded and the host
still holds references at ``Py_Finalize``), which leaves the SQLite database
open and makes nanobind report leaked instances. A no-op stand-in keeps late
callers (other ``atexit`` hooks) working.
"""
atexit.register(setattr, owner, attr, _noop_cache.Cache())
def _leading_version_ints(v: str) -> List[int]:
out = []
for part in re.split(r"[.+-]", v)[:3]:
if not part.isdigit():
break
out.append(int(part))
return out
def _version_to_epoch(v: str) -> int:
"""Pack a version string's leading "major.minor.patch" into a single
monotonically increasing int, ignoring any dev/local suffix -- e.g.
``"1.8.1.dev3+g1a2b3c4"`` and ``"1.8.1"`` both pack to 1_008_001.
"""
parts = (_leading_version_ints(v) + [0, 0, 0])[:3]
return parts[0] * 1_000_000 + parts[1] * 1_000 + parts[2]
try:
_CURRENT_CACHE_EPOCH = _version_to_epoch(importlib.metadata.version("speasy"))
except Exception:
_CURRENT_CACHE_EPOCH = 0
[docs]
class CacheItem:
def __init__(self, data, version, lifetime=None):
self.data = data
self.version = version
if lifetime is not None and isinstance(lifetime, (float, int)):
lifetime = timedelta(seconds=lifetime)
self.lifetime = lifetime
self.created = datetime.now(tz=timezone.utc)
self.cache_epoch = _CURRENT_CACHE_EPOCH
[docs]
def bump_creation_time(self) -> "CacheItem":
self.created = datetime.now(tz=timezone.utc)
return self
def __setstate__(self, state):
self.data = state["data"]
self.version = state["version"]
self.lifetime = state.get("lifetime", None)
self.created = state.get("created", datetime.now(tz=timezone.utc))
self.cache_epoch = state.get("cache_epoch", 0)
[docs]
def is_expired(self) -> bool:
if isinstance(self.lifetime, timedelta):
return datetime.now(tz=timezone.utc) >= (self.created + self.lifetime)
else:
return False
def _is_already_sciqlop_cache_layout(p: Path) -> bool:
return (p / "sciqlop-cache.db").is_file() or bool(list(p.glob("*/sciqlop-cache.db")))
def _is_legacy_diskcache_layout(p: Path) -> bool:
if not p.exists():
return False
# A directory already on the new backend can still have a stray ``cache.db``
# left over from an older pysciqlop-cache release (before that file was
# renamed to ``sciqlop-cache.db``). Presence of the new backend's own file is
# definitive: never treat such a directory as needing (re-)migration.
if _is_already_sciqlop_cache_layout(p):
return False
return (p / "cache.db").is_file() or bool(list(p.glob("*/cache.db")))
def _restore_legacy_cache(p: Path, backup: Path) -> None:
if p.exists():
shutil.rmtree(str(p))
os.rename(str(backup), str(p))
def _merge_stray_legacy_entries(p: Path) -> bool:
"""A ``cache.db`` sitting next to an already-migrated ``sciqlop-cache.db`` most
often means an older Speasy version (which knows nothing about sciqlop-cache)
was used since the last migration and wrote its own fresh diskcache into the
same directory. Merge those entries directly into the *existing* live cache in
place. Returns True if anything was merged.
This is deliberately not ``_migrate_legacy_diskcache``'s rename-the-whole-
directory-then-recreate approach: that would sweep the live sciqlop-cache.db
into the backup too, and migrate() only ever reads from diskcache, so the
live data would be orphaned in the backup, never restored.
"""
legacy_db = p / "cache.db"
if not legacy_db.is_file():
return False
import diskcache as dc
try:
legacy = dc.Cache(str(p))
keys = list(legacy)
except Exception:
log.exception(
f"Found {legacy_db} alongside the live cache, but it doesn't look like "
f"a valid diskcache (probably harmless leftover debris); leaving it and "
f"the live cache untouched."
)
return False
if not keys:
legacy.close()
return False
log.warning(
f"Detected {len(keys)} diskcache entries alongside the live cache at {p} "
f"(likely written by an older Speasy version used since the last "
f"migration); merging them into the live cache."
)
live = sc.Cache(cache_path=str(p), max_size=cache_cfg.size())
merged = 0
for key in keys:
try:
value = legacy.get(key)
except Exception:
log.exception(f"Could not read legacy entry {key!r} at {p}, skipping")
continue
if value is None:
continue
try:
live.set(key, value)
merged += 1
except Exception:
log.exception(f"Failed to merge legacy entry {key!r} into the live cache at {p}")
legacy.close()
for suffix in ("", "-wal", "-shm", "-journal"):
stray = Path(f"{legacy_db}{suffix}")
if stray.exists():
stray.unlink()
log.info(f"Merged {merged} new entries from a legacy diskcache found at {p} into the live cache.")
return True
def _migrate_legacy_diskcache(full_path: str, move: bool = False) -> bool:
"""Detect a legacy diskcache layout at ``full_path`` and migrate it to
sciqlop-cache format. Returns True if a migration was performed.
The legacy cache is renamed to ``<full_path>.diskcache.backup`` and the new
cache is written directly at ``full_path`` (sciqlop-cache references external
value files by absolute path, so the destination cannot be relocated after
writing). On any failure the legacy cache is restored.
For large caches this can take minutes — a one-time cost on first launch
after upgrading. The legacy backup is kept so the user can verify and
delete it manually, unless ``move`` is set.
Parameters
----------
move : bool
If True, delete each legacy entry as soon as it's migrated instead of
preserving the whole legacy cache as a rollback backup. Uses much less
peak disk space during migration, at the cost of not being able to fall
back to the legacy cache afterwards. See the ``migrate_by_moving`` config
entry (``[CACHE] migrate_by_moving`` / ``SPEASY_CACHE_MIGRATE_BY_MOVING``)
for the user-facing toggle.
Note: if migration itself fails partway through (see the ``except
Exception`` branch below), any entries already dropped from the legacy
cache before the failure are only recoverable from the (discarded) new
cache, not from the restored legacy one -- a narrow gap inherent to
deleting-as-you-go, unlike the default (non-move) mode where the legacy
cache is fully intact until migration succeeds outright.
"""
p = Path(full_path)
if _is_already_sciqlop_cache_layout(p):
# Already migrated -- but a cache.db here too means a pre-1.8 Speasy
# version was used again since, and wrote fresh legacy data alongside the
# live cache. Merge it in rather than silently discarding it.
return _merge_stray_legacy_entries(p)
if not _is_legacy_diskcache_layout(p):
return False
backup = Path(f"{p}.diskcache.backup")
if backup.exists():
log.warning(
f"Legacy cache backup already present at {backup}; "
f"skipping auto-migration. Move or remove it to retry."
)
return False
try:
from pysciqlop_cache.migrate import migrate
except ImportError as e:
log.exception(
f"Detected legacy diskcache layout at {p} but cannot migrate "
f"(missing dependency: {e}). Install diskcache once to allow "
f"migration, or delete {p} to start fresh."
)
return False
log.warning(
f"Detected legacy diskcache layout at {p}; migrating to sciqlop-cache. "
f"This is a one-time operation and may take several minutes for large caches."
)
# sciqlop-cache stores large values as external files referenced by
# absolute path, so a migrated cache cannot be relocated afterwards.
# Rename the legacy cache out of the way, then migrate straight into the
# final path; restore the legacy cache on any failure.
os.rename(str(p), str(backup))
try:
result = migrate(str(backup), str(p), drop=move)
except ImportError as e:
_restore_legacy_cache(p, backup)
log.exception(
f"Detected legacy diskcache layout at {p} but cannot migrate "
f"(missing dependency: {e}). Install diskcache once to allow "
f"migration, or delete {p} to start fresh."
)
return False
except Exception as e:
# Includes a legacy entry that fails to deserialize (e.g. pickled by an
# incompatible library version) -- migrate() has no per-entry error
# containment, so one bad entry aborts the whole migration. Never fatal:
# the legacy cache still works fine as a fallback, that's the entire point
# of not deleting it until migration succeeds.
_restore_legacy_cache(p, backup)
log.exception(
f"Migration of legacy cache at {p} failed ({e}); left the legacy cache "
f"in place, Speasy will keep using it and retry migration next time. "
f"Delete {p} to start fresh instead."
)
return False
if move:
log.info(
f"Migration complete: {result['migrated']} entries moved in "
f"{result['elapsed_secs']}s. Legacy entries were deleted as they "
f"migrated, so {backup} should already be empty."
)
else:
log.info(
f"Migration complete: {result['migrated']} entries in "
f"{result['elapsed_secs']}s. Legacy cache preserved at {backup}; "
f"delete it once you've verified the new cache works."
)
return True
def _open_or_recover(ctor, full_path: str, label: str):
"""Call ``ctor()`` to open the cache/index database at ``full_path``. If it
raises -- the database is corrupted, which can happen after a move-mode
migration interrupted mid-way (see ``migrate_by_moving``) or from
unrelated disk corruption -- delete the broken directory and retry once
with a fresh one.
Both :class:`Cache` and ``SpeasyIndex`` are constructed at Speasy import
time, so letting this exception propagate would crash every ``import
speasy`` until the user manually found and deleted the offending
directory. Degrading to "start over" is preferable: it's just a cache,
and a corrupted database has no rollback value worth keeping around --
unlike a legacy migration backup, there's nothing recoverable in it.
"""
try:
return ctor()
except Exception:
p = Path(full_path)
if not p.exists():
raise
log.exception(
f"Could not open the {label} at {full_path} (likely corrupted); deleting it and "
f"starting fresh so Speasy can still start."
)
shutil.rmtree(str(p))
return ctor()
def _warn_if_backup_present(full_path: str) -> None:
"""Log a reminder if a migration backup still sits next to ``full_path``.
Called on every construction of a cache/index built at this path, so the
reminder keeps showing up (once per import) until the user actually
deletes the backup with :func:`delete_migration_backups`.
"""
backup = Path(f"{full_path}.diskcache.backup")
if backup.exists():
log.warning(
f"A legacy cache backup from migrating to sciqlop-cache is still "
f"present at {backup}. Once you've verified the new cache works, "
f"delete it with speasy.core.cache.delete_migration_backups()."
)
def _known_migratable_paths() -> List[str]:
"""Paths of every cache/index Speasy builds that could carry a legacy
diskcache layout -- i.e. every call site of :func:`_migrate_legacy_diskcache`.
"""
from speasy.config import index as index_cfg
return [f"{cache_cfg.path()}/Cache", index_cfg.path()]
[docs]
def migration_backups() -> List[str]:
"""List legacy diskcache backups left over from migrating to sciqlop-cache.
Returns
-------
List[str]
Paths to backup directories that currently exist on disk.
"""
return [
backup for backup in (f"{p}.diskcache.backup" for p in _known_migratable_paths())
if os.path.isdir(backup)
]
[docs]
def delete_migration_backups() -> List[str]:
"""Delete legacy diskcache backups left over from migrating to sciqlop-cache.
Safe to call at any time: only removes the ``<path>.diskcache.backup``
directories created by the one-time migration, never a live cache.
Returns
-------
List[str]
Paths that were actually deleted.
"""
deleted = []
for backup in migration_backups():
shutil.rmtree(backup)
deleted.append(backup)
log.info(f"Deleted migration backup at {backup}")
return deleted
[docs]
class Cache:
__slots__ = ["cache_file", "_data", "cache_type"]
def __init__(self, cache_path: str = "", cache_type: str = "Cache"):
full_path = f"{cache_path}/{cache_type}"
_migrate_legacy_diskcache(full_path, move=cache_cfg.migrate_by_moving())
_warn_if_backup_present(full_path)
if cache_type == "Fanout":
self._data = _open_or_recover(
lambda: sc.FanoutCache(cache_path=full_path, shard_count=8, max_size=cache_cfg.size(),
serializer=sc.PickleOOBSerializer()),
full_path, "cache")
elif cache_type == "Cache":
self._data = _open_or_recover(
lambda: sc.Cache(cache_path=full_path, max_size=cache_cfg.size(), serializer=sc.PickleOOBSerializer()),
full_path, "cache")
else:
raise ValueError(f"Unimplemented cache type: {cache_type}")
self.cache_type = cache_type
self._data.reset_stats()
if self.version < cache_version:
self._data.clear()
self.version = cache_version
@property
def version(self):
return str_to_version(self._data.get("cache/version", default="0.0.0"))
@version.setter
def version(self, v: Union[str, Version]):
self._data["cache/version"] = v if type(v) is str else version_to_str(v)
[docs]
def disk_size(self):
return self._data.volume()
[docs]
def stats(self):
s = self._data.stats()
return {
"hit": s["hits"],
"misses": s["misses"],
}
def __len__(self):
return len(self._data)
[docs]
def keys(self):
return list(self._data)
def __contains__(self, item):
return item in self._data
def __getitem__(self, key):
return self._data[key]
def __setitem__(self, key, value):
self._data[key] = value
[docs]
def set(self, key, value, expire=None, tag: Optional[str] = None):
self._data.set(key, value, expire=expire, tag=tag)
[docs]
def add(self, key, value, expire=None, tag: Optional[str] = None):
return self._data.add(key, value, expire=expire, tag=tag)
[docs]
def get(self, key, default_value=None):
try:
return self._data.get(key, default_value)
except Exception:
# An entry that fails to deserialize (e.g. pickled by an incompatible
# library version, such as a different numpy major version) is not
# readable by this process -- treat it as a miss instead of crashing.
log.warning(f"Cache entry {key!r} could not be loaded (possibly written by an "
f"incompatible library version); treating it as a cache miss.",
exc_info=True)
return default_value
[docs]
def incr(self, key, delta=1, default=0):
return self._data.incr(key, delta, default=default)
[docs]
def drop(self, key):
self._data.delete(key)
[docs]
def evict_tag(self, tag: str):
return self._data.evict_tag(tag)
[docs]
def drop_matching_entries(self, pattern: Union[str, re.Pattern]):
"""Drop all cache entries that match a given pattern
Parameters
----------
pattern : str or re.Pattern
The pattern to match cache keys against
"""
if isinstance(pattern, str):
pattern = re.compile(pattern)
for key in filter(pattern.match, self.keys()):
log.debug(f"Dropping cache entry {key}")
self.drop(key)
[docs]
@contextmanager
def transact(self, key: Optional[str] = None):
"""Open a transaction context.
For a single Cache the ``key`` argument is ignored. For a FanoutCache
the transaction is per-shard and ``key`` selects the shard — atomicity
only applies within that shard, so all operations in the block must
target the same shard key.
"""
if self.cache_type == "Fanout":
if key is None:
raise ValueError(
"FanoutCache.transact requires a key to select a shard"
)
with self._data.transact(key):
yield
else:
with self._data.transact():
yield
[docs]
@contextmanager
def lock(self, key: str):
lock = sc.Lock(self._data, key)
lock.acquire()
try:
yield lock
finally:
lock.release()