diff --git a/AGENTS.md b/AGENTS.md index 24c8bec..45d421d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -11,7 +11,7 @@ cashet is a content-addressable compute cache with git semantics. Python functio - Install with HTTP server: `uv sync --extra server` - Install with all extras: `uv sync --all-extras` - Run tests: `uv run pytest tests/ -v` -- Redis tests auto-detect: skipped if Redis isn't reachable at `localhost:6379`. Set `CASHET_REDIS=1` to skip auto-detection and force-enable. +- Redis tests target `redis://localhost:6379/15` by default and FLUSH that database between tests. Override with `CASHET_TEST_REDIS_URL` (point it at a dedicated instance or database, never one holding data you care about). Skipped automatically if the target isn't reachable; set `CASHET_REDIS=1` to skip auto-detection and force-enable. - Start Redis for full suite: `docker run -d --name cashet-redis -p 6379:6379 redis:7-alpine` - Lint: `uv run ruff check src/ tests/` - Type check: `uv run pyright src/` @@ -39,13 +39,15 @@ Protocol-based dependency injection via three pluggable protocols in `src/cashet - **Store** — metadata + blob storage. Default: `SQLiteStore` in `store.py`. Also: `RedisStore` and `AsyncRedisStore` in `redis_store.py` - **AsyncStore** — async protocol for IO-bound backends (Redis, S3, HTTP). Defined in `protocols.py` - **Executor** — runs functions. Default: `LocalExecutor` in `executor.py` (sync). Async variant: `AsyncLocalExecutor` in `async_executor.py` -- **Serializer** — serialize/deserialize results. Default: `PickleSerializer` in `hashing.py` +- **Serializer** — serialize/deserialize results. Default: `PickleSerializer` in `serializers.py` (re-exported from `hashing.py` for compatibility) + +Internal module layout: `store.py` holds only the public SQLite store classes; the engine lives in `_sqlite_core.py`, schema and migrations in `_sqlite_schema.py`, cross-process lock plumbing in `_locks.py`. `redis_store.py` holds the Redis store classes; the key scheme, commit codec, and index pipeline commands live in `_redis_codec.py`. Hash-prefix validation shared by both stores is in `_ids.py`. Data flow (sync): `Client.submit()` → `build_task_def()` hashes function source + args → `LocalExecutor.submit()` checks cache → runs if needed → `build_commit()` creates commit with parent lineage → blobs stored via `Store.put_blob()` with zlib compression (256B threshold). Async data flow: `AsyncClient.submit()` → `build_task_def()` → `AsyncLocalExecutor.submit()` checks cache via async store → runs function in thread via `asyncio.to_thread()` → stores result via async `AsyncStore.put_blob()`. Heartbeat and locking are asyncio-native. -HTTP server: `client.serve(host, port, require_token=None)` exposes endpoints over HTTP. When `require_token` is set, all requests must include `Authorization: Bearer `. Async server (`AsyncClient.serve()`) uses `create_async_app()` with native async handlers. Sync server (`Client.serve()`) uses `create_app()` which wraps sync Client calls in `asyncio.to_thread`. +HTTP server: `client.serve(host, port, require_token=None)` exposes endpoints over HTTP. When `require_token` is set, all requests must include `Authorization: Bearer `. One handler set serves both clients through an ops adapter in `server.py`: `create_async_app()` (from `AsyncClient.serve()`) calls `AsyncClient` natively, `create_app()` (from `Client.serve()`) wraps sync Client calls in `asyncio.to_thread`. ## Key design decisions @@ -57,7 +59,7 @@ HTTP server: `client.serve(host, port, require_token=None)` exposes endpoints ov - `AsyncResultRef` is the async counterpart for `AsyncClient`. - IO-bound layers (Store, HTTP) are async. CPU-bound execution stays on threads via `asyncio.to_thread()`. - Redis blob data keys: `cashet:blob:data:{hash}`. Ref counters: `cashet:blob:ref:{hash}`. Per-fingerprint locks: `cashet:lock:{fingerprint}`. -- Redis `delete_commit` uses a Lua script (`_DECR_DELETE_SCRIPT`) for atomic DECR + conditional DELETE of blob keys when ref count reaches zero, avoiding a TOCTOU race. +- Redis `delete_commit` uses a Lua script (`DECR_DELETE_SCRIPT` in `_redis_codec.py`) for atomic DECR + conditional DELETE of blob keys when ref count reaches zero, avoiding a TOCTOU race. - Sync `Client` and async `AsyncClient` share common helpers extracted to `_client_base.py` (store dir resolution, task config, diff). ## Code style diff --git a/CHANGELOG.md b/CHANGELOG.md index ea11a8d..8ad3cb8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,90 @@ # Changelog +## 0.5.0 - 12.7.2026. + +Experimental correctness and performance release. Cache hits are lock-free +and roughly 10x faster; several data-integrity and cross-store consistency +bugs are fixed. Read the Notes before upgrading shared stores. + +### Fixed +- Blobs are written atomically (temp file, then rename). A crash mid-write + could leave a truncated file at the content-addressed path that the + deduplication check then trusted forever, permanently poisoning that hash. + `get_blob` also deletes a corrupt blob file so the next `put_blob` can store + the content again. +- Redis fingerprint lookups return the newest completed commit, matching + SQLite. The per-fingerprint index was ordered by expiry with arbitrary tie + breaking, so after a force re-run later calls could keep serving the older + result. Legacy index entries are rescored in place on first lookup. +- Garbage collection removes TTL-expired commits on both stores instead of + keeping them until they age past the access-time cutoff. Running and pending + commits are exempt from expiry eviction: their TTL clock starts at claim + time, so a task outliving its TTL must not be deleted mid-execution. +- A store error during a heartbeat renewal no longer kills the lease loop + (which let another worker reclaim and double-run a live task) and can no + longer surface as an exception that destroys an already-computed result. + Failed renewals retry on a quarter of the normal interval, so a single + transient error cannot leave the claim stale before the next attempt. +- Set and frozenset arguments are ordered by their stable serialized form + when hashing. Ordering by raw repr embedded memory addresses, so the same + logical set could hash differently in another process. +- Per-fingerprint lock state no longer grows without bound: SQLite fingerprint + locks are striped across 256 paths, and Redis locks are created per + acquisition instead of being cached per fingerprint. + +### Performance +Measured with `benchmarks/bench_hot_path.py` on one machine (medians, +0.4.5 vs 0.5.0): +- sync cache hit: 4.0 ms to 0.35 ms +- async cache hit: 3.9 ms to 0.21 ms +- hashing (`build_task_def`): 256 us to 10 us +- cache miss (run + store): 8.2 ms to 2.7 ms + +- Cache hits take no fingerprint lock, never rewrite the commit, and skip the + redundant parent query on the claim path. SQLite hits inside the access-bump + window perform no writes at all; Redis hits still update the shared access + index with one ZADD, which single-instance Redis absorbs cheaply. +- `last_accessed_at` is bumped at most once per hour per commit. Eviction + cutoffs are measured in days, so LRU eviction behavior is unchanged. +- SQLite connections use `synchronous=NORMAL` under WAL: a power loss can + drop the newest commits, which a compute cache recomputes; the database + itself cannot corrupt. +- Function source resolution and AST canonicalization are memoized per + function object. Globals, defaults, and closures are still hashed live, so + mutated module state keeps invalidating as before. + +### Changed +- A cache hit no longer rewrites the stored commit's status to `cached`; + stored commits keep `completed` and `cashet log` shows them as such. +- The HTTP server returns 422 with the final exception line (no traceback, + no server file paths) when a submitted task raises. 500 is reserved for + actual server errors. +- CLI commands that operate on an existing store exit non-zero when the store + directory does not exist instead of silently creating an empty one; + `import` and `serve` still create it. A group-level `--store-dir` option + overrides `$CASHET_DIR` and the `./.cashet` default. +- Redis tests target `redis://localhost:6379/15` by default and honor + `CASHET_TEST_REDIS_URL`. The suite flushes the target database, which + previously defaulted to a shared database 0. +- Internal modules reorganized; public import paths are unchanged. + Serializers moved from `cashet.hashing` to `cashet.serializers` (the old + imports still work via re-exports). `store.py` keeps the public SQLite + classes and delegates to `_sqlite_core`, `_sqlite_schema`, and `_locks`; + `redis_store.py` keeps the Redis classes and delegates its key scheme and + codec to `_redis_codec`. The HTTP server's duplicated sync and async + handlers were unified into one set behind an ops adapter. + +### Notes +- Argument hashes change for sets containing custom objects or single-element + tuples; affected entries recompute on first access. +- The Redis per-fingerprint index is rescored lazily. Entries written by + older versions with a TTL can still order incorrectly until rewritten, and + commits written before 0.5.0 are absent from the new expiry index, so they + only age out by access time. +- Blob writes create short-lived `*.tmp` files inside `objects/`. A failed + write removes its temp file, storage stats never count them, and garbage + collection sweeps any left behind by a crash once they are an hour old. + ## 0.4.5 - 21.6.2026. ### Fixed diff --git a/README.md b/README.md index f89050d..e5288bc 100644 --- a/README.md +++ b/README.md @@ -234,9 +234,10 @@ cashet export backup.tar.gz # export commits and blobs to an archive cashet import backup.tar.gz # import from an archive cashet stats # storage statistics cashet serve --host 127.0.0.1 --port 8000 +cashet --store-dir /path/to/store log # any command against another store ``` -`cashet show`, `cashet get`, and `cashet rm` exit with a non-zero status when the commit is not found, so they compose in scripts. +`cashet show`, `cashet get`, and `cashet rm` exit with a non-zero status when the commit is not found, so they compose in scripts. Commands that operate on an existing store exit non-zero when the store directory does not exist instead of creating an empty one; `import` and `serve` create it. `--store-dir` overrides `$CASHET_DIR` and the `./.cashet` default. ## Python API @@ -362,7 +363,7 @@ client.serve(host="127.0.0.1", port=8000, require_token="secret123") | GET | `/stats` | Storage statistics | | POST | `/gc` | Run garbage collection | -When `require_token` is set, every request needs an `Authorization: Bearer ` header. Request bodies are size-limited (default 500 MB). Use `AsyncClient.serve()` for the native async server. +When `require_token` is set, every request needs an `Authorization: Bearer ` header. Request bodies are size-limited (default 500 MB). A task that raises returns 422 with the exception line; 500 means the server itself failed. Use `AsyncClient.serve()` for the native async server. > **Security.** `/submit` does not execute client-supplied Python by default. The legacy `func_source`, `func_b64`, `args_b64`, and `kwargs_b64` payloads require both `allow_remote_code=True` and a non-empty token (`cashet serve --require-token secret123 --allow-remote-code`). That mode deserializes and runs arbitrary Python, so expose it only to trusted clients. @@ -447,7 +448,7 @@ Small results (under 1 KB) are stored inline in `meta.db` to avoid inode overhea ### Concurrency -cashet is safe across threads, processes, and machines that share one store. Concurrent submissions of the same uncached task are deduplicated: the function runs exactly once and all callers get the same result. Cross-process claims use file locks (SQLite) or per-fingerprint Redis locks, with a heartbeat lease so a crashed worker's claim is reclaimed (default 5 minutes, configurable via `LocalExecutor(running_ttl=...)`). +cashet is safe across threads, processes, and machines that share one store. Concurrent submissions of the same uncached task are deduplicated: the function runs exactly once and all callers get the same result. Cross-process claims use file locks (SQLite) or per-fingerprint Redis locks, with a heartbeat lease so a crashed worker's claim is reclaimed (default 5 minutes, configurable via `LocalExecutor(running_ttl=...)`). Cache hits take no locks. On SQLite a steady-state hit performs no writes at all; on Redis a hit updates the shared access index with a single ZADD, which is what lets every worker's hits feed one LRU. ### Notebooks @@ -462,6 +463,7 @@ cashet resolves function source through `inspect.getsource`, then `dill.source.g Upgrading across a hash-format change does not corrupt anything; old entries simply miss and recompute on first access. To start clean, run `cashet clear` or point at a fresh store directory. +- **0.4.5 to 0.5.0:** Argument hashes change for sets containing custom objects or single-element tuples; those entries recompute on first access. Cache hits no longer rewrite commit status to `cached`. The Redis fingerprint index is rescored to created_at ordering lazily on lookup; pre-0.5.0 TTL entries may order incorrectly until rewritten and are collected by access age only. SQLite opens with `synchronous=NORMAL`. The HTTP server returns 422 instead of 500 for task failures. CLI read commands require an existing store directory. - **0.4.4 to 0.4.5:** Hashing fixes (slotted-object state, referenced global containers, `ast.unparse` canonicalization) change function and argument cache keys, so results cached by earlier versions recompute on first access. The Redis tag index key scheme changed and is not migrated; rewrite affected commits to rebuild tag indexes. `import_archive` now returns `ImportResult(imported, skipped)` instead of a bare count. - **0.3.x to 0.4.0:** Added per-entry TTL and tag-based invalidation. SQLite auto-migrates; old caches stay readable. - **0.3.0 to 0.3.1:** Redis blob keys renamed to `cashet:blob:data:{hash}` and stats backfilled once. Clear Redis caches before upgrading if you rely on long-lived reuse. diff --git a/benchmarks/bench_hot_path.py b/benchmarks/bench_hot_path.py new file mode 100644 index 0000000..1fe2b54 --- /dev/null +++ b/benchmarks/bench_hot_path.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +import asyncio +import statistics +import tempfile +import time +from collections.abc import Callable +from pathlib import Path +from typing import Any + +from cashet import Client +from cashet.async_client import AsyncClient +from cashet.hashing import build_task_def + +ROUNDS = 300 + + +def transform(data: list[int], scale: float = 2.0) -> list[float]: + return [x * scale for x in data] + + +def _timed(fn: Callable[[], Any], rounds: int = ROUNDS) -> tuple[float, float]: + fn() + samples: list[float] = [] + for _ in range(rounds): + start = time.perf_counter() + fn() + samples.append((time.perf_counter() - start) * 1e6) + return statistics.median(samples), min(samples) + + +def _report(label: str, median_us: float, min_us: float) -> None: + print(f"{label:<28} median {median_us:9.1f} us min {min_us:9.1f} us") + + +def bench_sync(root: Path) -> None: + with Client(store_dir=root / "sync") as client: + client.submit(transform, [1, 2, 3]).load() + _report( + "hash (build_task_def)", + *_timed(lambda: build_task_def(transform, ([1, 2, 3],), {})), + ) + _report("sync hit (submit)", *_timed(lambda: client.submit(transform, [1, 2, 3]))) + _report( + "sync hit (submit + load)", + *_timed(lambda: client.submit(transform, [1, 2, 3]).load()), + ) + + +def bench_async(root: Path) -> None: + async def run() -> None: + client = AsyncClient(store_dir=root / "async") + ref = await client.submit(transform, [1, 2, 3]) + await ref.load() + + samples: list[float] = [] + for _ in range(ROUNDS): + start = time.perf_counter() + await client.submit(transform, [1, 2, 3]) + samples.append((time.perf_counter() - start) * 1e6) + _report("async hit (submit)", statistics.median(samples), min(samples)) + await client.close() + + asyncio.run(run()) + + +def bench_miss(root: Path) -> None: + with Client(store_dir=root / "miss") as client: + samples: list[float] = [] + for i in range(50): + start = time.perf_counter() + client.submit(transform, [i]) + samples.append((time.perf_counter() - start) * 1e6) + _report("sync miss (run + store)", statistics.median(samples), min(samples)) + + +def main() -> None: + root = Path(tempfile.mkdtemp(prefix="cashet-bench-")) + print(f"store: {root} rounds: {ROUNDS}") + bench_sync(root) + bench_async(root) + bench_miss(root) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 36692e5..61a504f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "cashet" -version = "0.4.5" +version = "0.5.0" description = "A Python memoization cache with Redis, async support, DAG pipelines, and an HTTP server" readme = "README.md" license = "MIT" diff --git a/src/cashet/__init__.py b/src/cashet/__init__.py index 83a7cd4..026be5b 100644 --- a/src/cashet/__init__.py +++ b/src/cashet/__init__.py @@ -6,15 +6,15 @@ from cashet.client import Client from cashet.dag import AsyncResultRef, ResultRef, TaskRef from cashet.executor import LocalExecutor -from cashet.hashing import ( - ClosureWarning, +from cashet.hashing import ClosureWarning +from cashet.models import Commit, ObjectRef, TaskDef, TaskError, TaskStatus +from cashet.protocols import AsyncStore, Executor, Store +from cashet.serializers import ( JsonSerializer, PickleSerializer, SafePickleSerializer, Serializer, ) -from cashet.models import Commit, ObjectRef, TaskDef, TaskError, TaskStatus -from cashet.protocols import AsyncStore, Executor, Store from cashet.store import AsyncSQLiteStore, SQLiteStore try: diff --git a/src/cashet/_batch.py b/src/cashet/_batch.py index 9a1f790..4246ac6 100644 --- a/src/cashet/_batch.py +++ b/src/cashet/_batch.py @@ -7,9 +7,10 @@ from cashet._client_base import resolve_task_config from cashet.dag import AsyncResultRef, TaskRef -from cashet.hashing import Serializer, build_task_def +from cashet.hashing import build_task_def from cashet.models import TaskError from cashet.protocols import AsyncStore +from cashet.serializers import Serializer BatchKey = int | str NormalizedTask = tuple[ diff --git a/src/cashet/_client_base.py b/src/cashet/_client_base.py index e63b5f2..f952df0 100644 --- a/src/cashet/_client_base.py +++ b/src/cashet/_client_base.py @@ -6,9 +6,9 @@ from pathlib import Path from typing import Any -from cashet.hashing import Serializer from cashet.models import Commit, TaskStatus from cashet.protocols import AsyncStore +from cashet.serializers import Serializer _DEFAULT_STORE_DIR = ".cashet" _CASHET_DIR_ENV = "CASHET_DIR" diff --git a/src/cashet/_ids.py b/src/cashet/_ids.py new file mode 100644 index 0000000..f4e09d4 --- /dev/null +++ b/src/cashet/_ids.py @@ -0,0 +1,7 @@ +from __future__ import annotations + + +def normalize_hash_prefix(hash: str) -> str | None: + if 0 < len(hash) <= 64 and all(c in "0123456789abcdefABCDEF" for c in hash): + return hash.lower() + return None diff --git a/src/cashet/_locks.py b/src/cashet/_locks.py new file mode 100644 index 0000000..5b505c8 --- /dev/null +++ b/src/cashet/_locks.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import asyncio +import threading +from dataclasses import dataclass +from typing import Any + + +@dataclass +class SQLiteLockState: + thread_lock: threading.Lock + file_lock: Any + + +SQLITE_LOCKS: dict[str, SQLiteLockState] = {} +SQLITE_LOCKS_GUARD = threading.Lock() + + +def sqlite_lock_state(lock_path: str) -> SQLiteLockState: + with SQLITE_LOCKS_GUARD: + state = SQLITE_LOCKS.get(lock_path) + if state is None: + from filelock import FileLock + + state = SQLiteLockState( + thread_lock=threading.Lock(), + file_lock=FileLock(lock_path, timeout=30, thread_local=False), + ) + SQLITE_LOCKS[lock_path] = state + return state + + +class SQLiteFingerprintLock: + def __init__(self, lock_path: str) -> None: + self._state = sqlite_lock_state(lock_path) + + async def __aenter__(self) -> None: + await asyncio.to_thread(self._state.thread_lock.acquire) + try: + await asyncio.to_thread(self._state.file_lock.acquire) + except Exception: + self._state.thread_lock.release() + raise + + async def __aexit__(self, *args: Any) -> None: + try: + await asyncio.to_thread(self._state.file_lock.release) + finally: + self._state.thread_lock.release() diff --git a/src/cashet/_redis_codec.py b/src/cashet/_redis_codec.py new file mode 100644 index 0000000..a282a1d --- /dev/null +++ b/src/cashet/_redis_codec.py @@ -0,0 +1,287 @@ +from __future__ import annotations + +import base64 +import json +from datetime import UTC, datetime, timedelta +from typing import Any + +from cashet.models import Commit, ObjectRef, StorageTier, TaskDef, TaskStatus + +DECR_DELETE_SCRIPT = """ + local ref_key = KEYS[1] + local blob_key = KEYS[2] + local stats_key = KEYS[3] + if redis.call('EXISTS', ref_key) == 0 then + return 0 + end + local count = redis.call('DECR', ref_key) + if count <= 0 then + local existed = redis.call('EXISTS', blob_key) + local bytes = 0 + if existed == 1 then + bytes = redis.call('STRLEN', blob_key) + end + redis.call('DEL', blob_key, ref_key) + if existed == 1 and redis.call('HGET', stats_key, 'ready') == '1' then + redis.call('HINCRBY', stats_key, 'objects', -1) + redis.call('HINCRBY', stats_key, 'bytes', -bytes) + end + return bytes + end + return 0 +""" + + +def commit_key(hash: str) -> str: + return f"cashet:commit:{hash}" + + +def blob_key(hash: str) -> str: + return f"cashet:blob:data:{hash}" + + +def fp_key(fingerprint: str) -> str: + return f"cashet:index:fingerprint:{fingerprint}" + + +def running_key(fingerprint: str) -> str: + return f"cashet:index:running:{fingerprint}" + + +def func_key(func_name: str) -> str: + return f"cashet:index:func:{func_name}" + + +def status_key(status: str) -> str: + return f"cashet:index:status:{status}" + + +def tag_key(key: str) -> str: + return f"cashet:tagk:{key}" + + +def tag_value_key(key: str, value: str) -> str: + # Length-prefix the key so a ':' inside a tag key or value can never make + # two distinct (key, value) pairs collide onto the same set. + return f"cashet:tagv:{len(key)}:{key}:{value}" + + +def access_key() -> str: + return "cashet:index:last_accessed" + + +def expires_key() -> str: + return "cashet:index:expires" + + +def blob_stats_key() -> str: + return "cashet:stats:blob" + + +def blob_stats_lock_key() -> str: + return "cashet:stats:blob:lock" + + +def stats_ready(raw: Any) -> bool: + if isinstance(raw, bytes): + return raw == b"1" + return raw == "1" or raw == 1 + + +def encode_commit(commit: Commit) -> bytes: + d: dict[str, Any] = { + "hash": commit.hash, + "fingerprint": commit.fingerprint, + "func_name": commit.task_def.func_name, + "func_hash": commit.task_def.func_hash, + "args_hash": commit.task_def.args_hash, + "args_snapshot_b64": base64.b64encode(commit.task_def.args_snapshot).decode(), + "func_source": commit.task_def.func_source, + "dep_versions": commit.task_def.dep_versions, + "cache": commit.task_def.cache, + "retries": commit.task_def.retries, + "force": commit.task_def.force, + "timeout_seconds": ( + commit.task_def.timeout.total_seconds() if commit.task_def.timeout else None + ), + "ttl_seconds": ( + commit.task_def.ttl.total_seconds() if commit.task_def.ttl else None + ), + "expires_at": commit.expires_at.isoformat() if commit.expires_at else None, + "input_refs": [ + {"hash": r.hash, "size": r.size, "tier": r.tier.value} for r in commit.input_refs + ], + "output_hash": commit.output_ref.hash if commit.output_ref else None, + "output_size": commit.output_ref.size if commit.output_ref else None, + "output_tier": commit.output_ref.tier.value if commit.output_ref else None, + "parent_hash": commit.parent_hash, + "status": commit.status.value, + "error": commit.error, + "tags": commit.tags, + "created_at": commit.created_at.isoformat(), + "claimed_at": commit.claimed_at.isoformat(), + "last_accessed_at": datetime.now(UTC).isoformat(), + } + return json.dumps(d, separators=(",", ":")).encode() + + +def decode_commit(data: bytes) -> Commit: + d = json.loads(data) + task_def = TaskDef( + func_hash=d["func_hash"], + func_name=d["func_name"], + func_source=d.get("func_source", ""), + args_hash=d["args_hash"], + args_snapshot=base64.b64decode(d.get("args_snapshot_b64", "")), + dep_versions=d.get("dep_versions", {}), + cache=d.get("cache", True), + tags=d.get("tags", {}), + retries=d.get("retries", 0), + force=d.get("force", False), + timeout=( + timedelta(seconds=d["timeout_seconds"]) + if d.get("timeout_seconds") is not None + else None + ), + ttl=( + timedelta(seconds=d["ttl_seconds"]) + if d.get("ttl_seconds") is not None + else None + ), + ) + input_refs = [ + ObjectRef( + hash=r["hash"], + size=r.get("size", 0), + tier=StorageTier(r.get("tier", "blob")), + ) + for r in d.get("input_refs", []) + ] + output_ref = None + if d.get("output_hash"): + output_ref = ObjectRef( + hash=d["output_hash"], + size=d.get("output_size", 0), + tier=StorageTier(d.get("output_tier", "blob")), + ) + created_at = ( + datetime.fromisoformat(d["created_at"]) if "created_at" in d else datetime.now(UTC) + ) + claimed_at = ( + datetime.fromisoformat(d["claimed_at"]) if "claimed_at" in d else datetime.now(UTC) + ) + expires_at = ( + datetime.fromisoformat(d["expires_at"]) if d.get("expires_at") is not None else None + ) + return Commit( + hash=d["hash"], + task_def=task_def, + input_refs=input_refs, + output_ref=output_ref, + parent_hash=d.get("parent_hash"), + status=TaskStatus(d["status"]), + created_at=created_at, + claimed_at=claimed_at, + error=d.get("error"), + tags=d.get("tags", {}), + expires_at=expires_at, + ) + + +def decode_hash(raw: Any) -> str: + return raw.decode() if isinstance(raw, bytes) else raw + + +def commit_access_timestamp(data: bytes) -> float: + d = json.loads(data) + value = d.get("last_accessed_at") or d.get("created_at") + if value is None: + return datetime.now(UTC).timestamp() + return datetime.fromisoformat(value).timestamp() + + +def blob_ref_key(blob_hash: str) -> str: + return f"cashet:blob:ref:{blob_hash}" + + +def blob_hashes(commit: Commit) -> set[str]: + hashes: set[str] = set() + if commit.output_ref: + hashes.add(commit.output_ref.hash) + for ref in commit.input_refs: + hashes.add(ref.hash) + return hashes + + +def matches_tags(commit: Commit, tags: dict[str, str | None]) -> bool: + for key, val in tags.items(): + if val is None: + if key not in commit.tags: + return False + else: + if commit.tags.get(key) != val: + return False + return True + + +def stats_dict(total: int, completed: int, blob_count: int, blob_bytes: int) -> dict[str, int]: + return { + "total_commits": total, + "completed_commits": completed, + "stored_objects": blob_count, + "disk_bytes": blob_bytes, + "blob_objects": blob_count, + "blob_bytes": blob_bytes, + "inline_objects": 0, + "inline_bytes": 0, + } + + +def index_commit_commands(pipe: Any, commit: Commit) -> None: + pipe.get(commit_key(commit.hash)) + pipe.set(commit_key(commit.hash), encode_commit(commit)) + ts = commit.created_at.timestamp() + pipe.zadd("cashet:index:all", {commit.hash: ts}) + pipe.zadd(fp_key(commit.fingerprint), {commit.hash: ts}) + # Index expiry only once the commit is terminal: expires_at is stamped at + # claim time, so a task outliving its TTL is expired while RUNNING, and + # expiry eviction must not delete a live claim out from under its worker. + if commit.expires_at is not None and commit.status not in ( + TaskStatus.RUNNING, + TaskStatus.PENDING, + ): + pipe.zadd(expires_key(), {commit.hash: commit.expires_at.timestamp()}) + else: + pipe.zrem(expires_key(), commit.hash) + pipe.zadd(func_key(commit.task_def.func_name), {commit.hash: ts}) + now_ts = datetime.now(UTC).timestamp() + pipe.zadd(access_key(), {commit.hash: now_ts}) + for status in TaskStatus: + pipe.srem(status_key(status.value), commit.hash) + pipe.sadd(status_key(commit.status.value), commit.hash) + if commit.status == TaskStatus.RUNNING: + pipe.sadd(running_key(commit.fingerprint), commit.hash) + else: + pipe.srem(running_key(commit.fingerprint), commit.hash) + for key, val in commit.tags.items(): + pipe.sadd(tag_key(key), commit.hash) + pipe.sadd(tag_value_key(key, val), commit.hash) + + +def remove_commit_index_commands(pipe: Any, commit: Commit, resolved_hash: str) -> None: + pipe.delete(commit_key(resolved_hash)) + pipe.zrem("cashet:index:all", resolved_hash) + pipe.zrem(fp_key(commit.fingerprint), resolved_hash) + pipe.zrem(expires_key(), resolved_hash) + pipe.srem(running_key(commit.fingerprint), resolved_hash) + pipe.zrem(func_key(commit.task_def.func_name), resolved_hash) + pipe.zrem(access_key(), resolved_hash) + for status in TaskStatus: + pipe.srem(status_key(status.value), resolved_hash) + for key, val in commit.tags.items(): + pipe.srem(tag_key(key), resolved_hash) + pipe.srem(tag_value_key(key, val), resolved_hash) + + +def commit_hash_from_key(raw: Any) -> str: + return decode_hash(raw).split(":")[-1] diff --git a/src/cashet/_sqlite_core.py b/src/cashet/_sqlite_core.py new file mode 100644 index 0000000..9762a9f --- /dev/null +++ b/src/cashet/_sqlite_core.py @@ -0,0 +1,652 @@ +from __future__ import annotations + +import hashlib +import json +import logging +import os +import sqlite3 +import threading +import time +import zlib +from datetime import UTC, datetime, timedelta +from pathlib import Path +from typing import Any + +from cashet._ids import normalize_hash_prefix +from cashet._sqlite_schema import ensure_schema, put_commit_row, row_to_commit +from cashet.models import Commit, ObjectRef, StorageTier, TaskStatus + +_BLOB_COMPRESS_THRESHOLD = 256 +_INLINE_THRESHOLD = 1024 +_ACCESS_BUMP_GRANULARITY = timedelta(hours=1) +_TMP_SWEEP_AGE_SECONDS = 3600 + +logger = logging.getLogger("cashet") + + +class SQLiteStoreCore: + def __init__(self, root: Path) -> None: + self.root = root + self.objects_dir = root / "objects" + self.db_path = root / "meta.db" + self.objects_dir.mkdir(parents=True, exist_ok=True) + self._tls = threading.local() + logger.debug("initializing sqlite store path=%s", str(self.db_path)) + ensure_schema(self._connect()) + + def _connect(self, *, immediate: bool = False) -> sqlite3.Connection: + conn: sqlite3.Connection | None = getattr(self._tls, "conn", None) + try: + if conn is not None: + conn.execute("SELECT 1") + else: + conn = None + except sqlite3.Error: + conn = None + if conn is None: + conn = sqlite3.connect(str(self.db_path), isolation_level=None) + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA auto_vacuum=INCREMENTAL") + conn.execute("PRAGMA journal_mode=WAL") + # NORMAL under WAL cannot corrupt the database; a power loss may + # drop the most recent commits, which a compute cache can recompute. + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute("PRAGMA busy_timeout=5000") + self._tls.conn = conn + if immediate: + conn.execute("BEGIN IMMEDIATE") + return conn + + def close(self) -> None: + logger.debug("closing sqlite store path=%s", str(self.db_path)) + conn: sqlite3.Connection | None = getattr(self._tls, "conn", None) + if conn is not None: + conn.close() + self._tls.conn = None + + def put_blob(self, data: bytes) -> ObjectRef: + content_hash = hashlib.sha256(data).hexdigest() + if len(data) < _INLINE_THRESHOLD: + conn = self._connect() + conn.execute( + "INSERT OR IGNORE INTO inline_objects (hash, data) VALUES (?, ?)", + (content_hash, data), + ) + logger.info( + "blob stored hash=%s size=%d tier=inline", + content_hash[:12], + len(data), + ) + return ObjectRef(hash=content_hash, size=len(data), tier=StorageTier.INLINE) + prefix = content_hash[:2] + suffix = content_hash[2:] + obj_path = self.objects_dir / prefix / suffix + if obj_path.exists(): + logger.debug( + "blob deduplicated hash=%s size=%d", + content_hash[:12], + len(data), + ) + return ObjectRef(hash=content_hash, size=len(data), tier=StorageTier.BLOB) + obj_path.parent.mkdir(parents=True, exist_ok=True) + stored = data + if len(data) >= _BLOB_COMPRESS_THRESHOLD: + compressed = zlib.compress(data, level=6) + if len(compressed) < len(data): + stored = compressed + # Write-then-rename so a crash mid-write can never leave a partial file + # at the content-addressed path, where the exists() dedup check above + # would treat it as the real blob forever. + tmp_path = obj_path.with_name( + f"{obj_path.name}.{os.getpid()}.{threading.get_ident()}.tmp" + ) + try: + tmp_path.write_bytes(stored) + os.replace(tmp_path, obj_path) + except BaseException: + tmp_path.unlink(missing_ok=True) + raise + logger.info( + "blob stored hash=%s size=%d tier=blob compressed=%s", + content_hash[:12], + len(data), + "true" if stored is not data else "false", + ) + return ObjectRef(hash=content_hash, size=len(data), tier=StorageTier.BLOB) + + def get_blob(self, ref: ObjectRef) -> bytes: + if ref.tier == StorageTier.INLINE: + conn = self._connect() + row = conn.execute( + "SELECT data FROM inline_objects WHERE hash = ?", (ref.hash,) + ).fetchone() + if row is None: + logger.warning("inline blob not found hash=%s", ref.hash[:12]) + raise ValueError(f"Inline blob {ref.hash} not found") + data = row["data"] + if hashlib.sha256(data).hexdigest() != ref.hash: + logger.error( + "inline blob integrity check failed hash=%s", + ref.hash[:12], + ) + raise ValueError(f"Inline blob {ref.hash} integrity check failed") + return data + prefix = ref.hash[:2] + suffix = ref.hash[2:] + obj_path = self.objects_dir / prefix / suffix + raw = obj_path.read_bytes() + try: + decompressed = zlib.decompress(raw) + if hashlib.sha256(decompressed).hexdigest() == ref.hash: + return decompressed + except zlib.error: + pass + if hashlib.sha256(raw).hexdigest() != ref.hash: + # The path is derived from the content hash, so a mismatched file is + # corrupt; removing it lets the next put_blob store it again. + obj_path.unlink(missing_ok=True) + logger.error( + "corrupt blob removed hash=%s", + ref.hash[:12], + ) + raise ValueError(f"Blob {ref.hash} integrity check failed") + return raw + + def blob_exists(self, hash: str) -> bool: + if (self.objects_dir / hash[:2] / hash[2:]).exists(): + return True + conn = self._connect() + row = conn.execute( + "SELECT 1 FROM inline_objects WHERE hash = ?", (hash,) + ).fetchone() + return row is not None + + def put_commit(self, commit: Commit) -> None: + conn = self._connect(immediate=True) + now = datetime.now(UTC).isoformat() + try: + put_commit_row(conn, commit, now) + conn.execute("COMMIT") + except Exception: + conn.execute("ROLLBACK") + raise + + def find_by_fingerprint(self, fingerprint: str) -> Commit | None: + conn = self._connect() + now = datetime.now(UTC) + now_iso = now.isoformat() + row = conn.execute( + """SELECT * FROM commits + WHERE fingerprint = ? AND status IN ('completed', 'cached') + AND (expires_at IS NULL OR expires_at > ?) + ORDER BY created_at DESC + LIMIT 1""", + (fingerprint, now_iso), + ).fetchone() + if row is None: + return None + # Bump at most once per granularity window: eviction cutoffs are + # measured in days, so hour-precision LRU keeps steady-state cache hits + # free of write statements. Decide from the row already in hand (even a + # zero-row UPDATE would take the writer lock); the SQL predicate stays + # as the guard against a concurrent bump. + threshold = (now - _ACCESS_BUMP_GRANULARITY).isoformat() + last_accessed = row["last_accessed_at"] + if last_accessed is None or last_accessed < threshold: + try: + conn.execute( + "UPDATE commits SET last_accessed_at = ? WHERE hash = ? " + "AND (last_accessed_at IS NULL OR last_accessed_at < ?)", + (now_iso, row["hash"], threshold), + ) + except sqlite3.OperationalError: + # Access-time bump only feeds LRU ordering; never fail a cache + # hit because a concurrent writer holds the lock past + # busy_timeout. + logger.debug( + "last_accessed_at bump skipped (db locked) hash=%s", row["hash"][:12] + ) + return row_to_commit(row) + + def find_running_by_fingerprint(self, fingerprint: str) -> Commit | None: + conn = self._connect() + row = conn.execute( + """SELECT * FROM commits + WHERE fingerprint = ? AND status = 'running' + ORDER BY claimed_at DESC LIMIT 1""", + (fingerprint,), + ).fetchone() + if row is None: + return None + return row_to_commit(row) + + def get_commit(self, hash: str) -> Commit | None: + normalized = normalize_hash_prefix(hash) + if normalized is None: + return None + hash = normalized + conn = self._connect() + if len(hash) < 64: + rows = conn.execute( + "SELECT * FROM commits WHERE hash LIKE ?", (hash + "%",) + ).fetchall() + if not rows: + return None + if len(rows) > 1: + matches = ", ".join(r["hash"][:12] for r in rows) + raise ValueError( + f"Ambiguous prefix {hash[:12]} matches {len(rows)} commits: {matches}" + ) + return row_to_commit(rows[0]) + row = conn.execute("SELECT * FROM commits WHERE hash = ?", (hash,)).fetchone() + if row is None: + return None + return row_to_commit(row) + + def list_commits( + self, + func_name: str | None = None, + limit: int = 50, + status: TaskStatus | None = None, + tags: dict[str, str | None] | None = None, + ) -> list[Commit]: + conn = self._connect() + query = "SELECT * FROM commits WHERE 1=1" + params: list[Any] = [] + if func_name: + query += " AND func_name = ?" + params.append(func_name) + if status: + query += " AND status = ?" + params.append(status.value) + if tags: + for key, val in tags.items(): + if val is None: + query += " AND json_extract(tags, ?) IS NOT NULL" + params.append(f"$.{key}") + else: + query += " AND json_extract(tags, ?) = ?" + params.append(f"$.{key}") + params.append(val) + query += " ORDER BY created_at DESC LIMIT ?" + params.append(limit) + rows = conn.execute(query, params).fetchall() + return [row_to_commit(r) for r in rows] + + def get_history(self, hash: str) -> list[Commit]: + normalized = normalize_hash_prefix(hash) + if normalized is None: + return [] + hash = normalized + conn = self._connect() + if len(hash) < 64: + rows = conn.execute( + "SELECT * FROM commits WHERE hash LIKE ?", (hash + "%",) + ).fetchall() + if not rows: + return [] + if len(rows) > 1: + matches = ", ".join(r["hash"][:12] for r in rows) + raise ValueError( + f"Ambiguous prefix {hash[:12]} matches {len(rows)} commits: {matches}" + ) + commit = row_to_commit(rows[0]) + else: + row = conn.execute("SELECT * FROM commits WHERE hash = ?", (hash,)).fetchone() + if row is None: + return [] + commit = row_to_commit(row) + fingerprint = commit.fingerprint + rows = conn.execute( + """SELECT * FROM commits + WHERE fingerprint = ? AND status IN ('completed', 'cached') + ORDER BY created_at ASC""", + (fingerprint,), + ).fetchall() + now = datetime.now(UTC) + result: list[Commit] = [] + for r in rows: + c = row_to_commit(r) + if c.expires_at is not None and c.expires_at <= now: + continue + result.append(c) + return result + + def stats(self) -> dict[str, int]: + conn = self._connect() + total = conn.execute("SELECT COUNT(*) FROM commits").fetchone()[0] + completed = conn.execute( + "SELECT COUNT(*) FROM commits WHERE status IN ('completed', 'cached')" + ).fetchone()[0] + obj_count, total_bytes = self._blob_storage_totals() + inline_count, inline_bytes = self._inline_storage_totals(conn) + return { + "total_commits": total, + "completed_commits": completed, + "stored_objects": obj_count + inline_count, + "disk_bytes": total_bytes + inline_bytes, + "blob_objects": obj_count, + "blob_bytes": total_bytes, + "inline_objects": inline_count, + "inline_bytes": inline_bytes, + } + + def _blob_storage_totals(self) -> tuple[int, int]: + obj_count = 0 + total_bytes = 0 + for p in self.objects_dir.iterdir(): + if p.is_dir(): + for f in p.iterdir(): + # A crashed writer's leftover temp file is not a stored + # object; counting it would inflate size-based eviction's + # target with bytes commit GC can never reclaim. + if f.is_file() and not f.name.endswith(".tmp"): + obj_count += 1 + total_bytes += f.stat().st_size + return obj_count, total_bytes + + def _sweep_stale_tmp_files(self) -> None: + # Age-gated so a concurrent writer's in-flight temp file survives; a + # healthy write lives milliseconds, an hour-old one is crash debris. + cutoff = time.time() - _TMP_SWEEP_AGE_SECONDS + for prefix_dir in self.objects_dir.iterdir(): + if not prefix_dir.is_dir(): + continue + for f in prefix_dir.glob("*.tmp"): + try: + if f.stat().st_mtime < cutoff: + f.unlink(missing_ok=True) + except OSError: + continue + + def _inline_storage_totals(self, conn: sqlite3.Connection) -> tuple[int, int]: + inline_row = conn.execute( + "SELECT COUNT(*) AS cnt, COALESCE(SUM(LENGTH(data)), 0) AS bytes FROM inline_objects" + ).fetchone() + inline_count = inline_row["cnt"] if inline_row else 0 + inline_bytes = inline_row["bytes"] if inline_row else 0 + return inline_count, inline_bytes + + def _storage_bytes(self, conn: sqlite3.Connection) -> int: + return self._blob_storage_totals()[1] + self._inline_storage_totals(conn)[1] + + def evict( + self, older_than: datetime, max_size_bytes: int | None = None + ) -> int: + self._sweep_stale_tmp_files() + conn = self._connect(immediate=True) + orphans: list[str] = [] + # The expiry term must never delete a live claim: expires_at is stamped + # at claim time, so a task outliving its TTL is expired while RUNNING, + # and deleting it would let a second submission double-run it. Access + # age stays status-blind so a crashed worker's old claim still ages out. + evictable = ( + "(last_accessed_at < ? OR (expires_at IS NOT NULL AND expires_at <= ? " + "AND status NOT IN ('running', 'pending')))" + ) + try: + params = (older_than.isoformat(), datetime.now(UTC).isoformat()) + # The OR on expires_at defeats the last_accessed index, so evaluate + # the predicate in one scan and delete by hash. + rows = conn.execute( + f"SELECT hash, output_hash, input_refs FROM commits WHERE {evictable}", + params, + ).fetchall() + evicted_hashes = [r[0] for r in rows] + candidates: set[str] = set() + for r in rows: + if r[1]: + candidates.add(r[1]) + if r[2]: + candidates.update(json.loads(r[2])) + deleted = len(evicted_hashes) + if evicted_hashes: + placeholders = ", ".join("?" for _ in evicted_hashes) + conn.execute( + f"UPDATE commits SET parent_hash = NULL WHERE parent_hash IN ({placeholders})", + evicted_hashes, + ) + conn.execute( + f"DELETE FROM commits WHERE hash IN ({placeholders})", evicted_hashes + ) + if candidates: + orphans = self._find_orphan_objects(conn, candidates) + conn.execute("COMMIT") + except Exception: + conn.execute("ROLLBACK") + raise + if orphans: + logger.info( + "orphan objects cleaned count=%d", + len(orphans), + ) + self._delete_orphan_objects(conn, orphans) + + if max_size_bytes is not None: + current_bytes = self._storage_bytes(conn) + if current_bytes > max_size_bytes: + deleted += self._evict_to_size(current_bytes, max_size_bytes) + + if deleted: + logger.info( + "eviction complete deleted=%d reason=%s", + deleted, + "size_limit" if max_size_bytes is not None else "ttl", + ) + try: + self._vacuum() + except sqlite3.OperationalError: + # Reclaiming space is best-effort; a concurrent writer holding the + # lock must not undo an eviction whose deletes already committed. + logger.debug("vacuum skipped (db busy) after eviction") + else: + logger.debug("eviction found no candidates") + + return deleted + + def _vacuum(self) -> None: + conn = self._connect() + mode = conn.execute("PRAGMA auto_vacuum").fetchone()[0] + if mode == 2: + conn.execute("PRAGMA incremental_vacuum") + else: + conn.execute("VACUUM") + + def _evict_to_size(self, current_bytes: int, max_size_bytes: int) -> int: + pending_orphans: list[str] = [] + deleted = 0 + conn = self._connect(immediate=True) + try: + ref_counts = self._object_ref_counts(conn) + rows = conn.execute( + "SELECT * FROM commits ORDER BY last_accessed_at ASC" + ).fetchall() + for target in rows: + if current_bytes <= max_size_bytes: + break + refs = self._row_object_refs(target) + conn.execute( + "UPDATE commits SET parent_hash = NULL WHERE parent_hash = ?", + (target["hash"],), + ) + conn.execute("DELETE FROM commits WHERE hash = ?", (target["hash"],)) + deleted += 1 + for obj_hash in refs: + count = ref_counts.get(obj_hash, 0) + if count <= 0: + continue + if count == 1: + ref_counts.pop(obj_hash, None) + if obj_hash not in pending_orphans: + current_bytes -= self._object_storage_size(conn, obj_hash) + pending_orphans.append(obj_hash) + else: + ref_counts[obj_hash] = count - 1 + conn.execute("COMMIT") + except Exception: + conn.execute("ROLLBACK") + raise + if pending_orphans: + logger.info( + "orphan objects cleaned count=%d", + len(pending_orphans), + ) + self._delete_orphan_objects(conn, pending_orphans) + return deleted + + def _delete_commit_body(self, conn: sqlite3.Connection, hash: str) -> tuple[bool, list[str]]: + if len(hash) < 64: + rows = conn.execute( + "SELECT * FROM commits WHERE hash LIKE ?", (hash + "%",) + ).fetchall() + if not rows: + return False, [] + if len(rows) > 1: + matches = ", ".join(r["hash"][:12] for r in rows) + raise ValueError( + f"Ambiguous prefix {hash[:12]} matches {len(rows)} commits: {matches}" + ) + target = rows[0] + else: + target = conn.execute( + "SELECT * FROM commits WHERE hash = ?", (hash,) + ).fetchone() + if target is None: + return False, [] + + candidates = set(self._row_object_refs(target)) + conn.execute( + "UPDATE commits SET parent_hash = NULL WHERE parent_hash = ?", + (target["hash"],), + ) + conn.execute("DELETE FROM commits WHERE hash = ?", (target["hash"],)) + + orphans: list[str] = [] + if candidates: + orphans = self._find_orphan_objects(conn, candidates) + return True, orphans + + def delete_commit(self, hash: str) -> bool: + normalized = normalize_hash_prefix(hash) + if normalized is None: + return False + hash = normalized + conn = self._connect(immediate=True) + try: + success, orphans = self._delete_commit_body(conn, hash) + if not success: + # Avoid leaving a dangling transaction that poisons the next write. + conn.execute("ROLLBACK") + return False + conn.execute("COMMIT") + except Exception: + conn.execute("ROLLBACK") + raise + if orphans: + logger.info( + "orphan objects cleaned count=%d", + len(orphans), + ) + self._delete_orphan_objects(conn, orphans) + logger.debug("commit deleted hash=%s", hash[:12]) + return True + + def delete_by_tags(self, tags: dict[str, str | None]) -> int: + conn = self._connect(immediate=True) + query = "SELECT hash, output_hash, input_refs FROM commits WHERE 1=1" + params: list[Any] = [] + for key, val in tags.items(): + if val is None: + query += " AND json_extract(tags, ?) IS NOT NULL" + params.append(f"$.{key}") + else: + query += " AND json_extract(tags, ?) = ?" + params.append(f"$.{key}") + params.append(val) + rows = conn.execute(query, params).fetchall() + if not rows: + conn.execute("ROLLBACK") + return 0 + hashes = [r[0] for r in rows] + candidates: set[str] = set() + for r in rows: + if r[1]: + candidates.add(r[1]) + if r[2]: + for h in json.loads(r[2]): + candidates.add(h) + try: + placeholders = ", ".join("?" for _ in hashes) + conn.execute( + f"UPDATE commits SET parent_hash = NULL WHERE parent_hash IN ({placeholders})", + hashes, + ) + conn.execute( + f"DELETE FROM commits WHERE hash IN ({placeholders})", hashes + ) + deleted = len(hashes) + all_orphans = self._find_orphan_objects(conn, candidates) if candidates else [] + conn.execute("COMMIT") + except Exception: + conn.execute("ROLLBACK") + raise + if all_orphans: + logger.info("orphan objects cleaned count=%d", len(all_orphans)) + self._delete_orphan_objects(conn, all_orphans) + return deleted + + def _row_object_refs(self, row: sqlite3.Row) -> list[str]: + refs: list[str] = [] + if row["output_hash"]: + refs.append(row["output_hash"]) + if row["input_refs"]: + refs.extend(set(json.loads(row["input_refs"]))) + return refs + + def _object_ref_counts(self, conn: sqlite3.Connection) -> dict[str, int]: + counts: dict[str, int] = {} + for row in conn.execute("SELECT output_hash FROM commits WHERE output_hash IS NOT NULL"): + counts[row[0]] = counts.get(row[0], 0) + 1 + for row in conn.execute("SELECT input_refs FROM commits WHERE input_refs IS NOT NULL"): + for h in set(json.loads(row[0])): + counts[h] = counts.get(h, 0) + 1 + return counts + + def _object_storage_size(self, conn: sqlite3.Connection, obj_hash: str) -> int: + size = 0 + blob_path = self.objects_dir / obj_hash[:2] / obj_hash[2:] + if blob_path.exists(): + size += blob_path.stat().st_size + row = conn.execute( + "SELECT LENGTH(data) AS size FROM inline_objects WHERE hash = ?", (obj_hash,) + ).fetchone() + if row is not None: + size += row["size"] or 0 + return size + + def _find_orphan_objects(self, conn: sqlite3.Connection, candidates: set[str]) -> list[str]: + still_output: set[str] = set() + for row in conn.execute("SELECT output_hash FROM commits WHERE output_hash IS NOT NULL"): + still_output.add(row[0]) + still_input: set[str] = set() + for row in conn.execute("SELECT input_refs FROM commits WHERE input_refs IS NOT NULL"): + still_input.update(json.loads(row[0])) + return [h for h in candidates if h not in still_output and h not in still_input] + + def _delete_orphan_objects(self, conn: sqlite3.Connection, orphans: list[str]) -> int: + freed = 0 + for obj_hash in orphans: + blob_path = self.objects_dir / obj_hash[:2] / obj_hash[2:] + if blob_path.exists(): + freed += blob_path.stat().st_size + blob_path.unlink() + row = conn.execute( + "SELECT LENGTH(data) AS size FROM inline_objects WHERE hash = ?", (obj_hash,) + ).fetchone() + if row is not None: + freed += row["size"] or 0 + conn.execute("DELETE FROM inline_objects WHERE hash = ?", (obj_hash,)) + for prefix_dir in list(self.objects_dir.iterdir()): + if prefix_dir.is_dir() and not any(prefix_dir.iterdir()): + prefix_dir.rmdir() + return freed diff --git a/src/cashet/_sqlite_schema.py b/src/cashet/_sqlite_schema.py new file mode 100644 index 0000000..ea0e025 --- /dev/null +++ b/src/cashet/_sqlite_schema.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import json +import sqlite3 +from datetime import datetime, timedelta + +from cashet.models import Commit, ObjectRef, StorageTier, TaskDef, TaskStatus + + +def ensure_schema(conn: sqlite3.Connection) -> None: + conn.execute(""" + CREATE TABLE IF NOT EXISTS commits ( + hash TEXT PRIMARY KEY, + fingerprint TEXT NOT NULL, + func_name TEXT NOT NULL, + func_hash TEXT NOT NULL, + args_hash TEXT NOT NULL, + args_snapshot BLOB, + func_source TEXT, + dep_versions TEXT, + cache INTEGER NOT NULL DEFAULT 1, + retries INTEGER NOT NULL DEFAULT 0, + force INTEGER NOT NULL DEFAULT 0, + timeout_seconds REAL, + ttl_seconds REAL, + input_refs TEXT, + output_hash TEXT, + output_size INTEGER, + output_tier TEXT, + parent_hash TEXT, + status TEXT NOT NULL DEFAULT 'pending', + error TEXT, + tags TEXT, + created_at TEXT NOT NULL, + claimed_at TEXT NOT NULL, + last_accessed_at TEXT, + expires_at TEXT, + FOREIGN KEY (parent_hash) REFERENCES commits(hash) + ) + """) + conn.execute("CREATE INDEX IF NOT EXISTS idx_fingerprint ON commits(fingerprint)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_func_name ON commits(func_name)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_created_at ON commits(created_at)") + _migrate_last_accessed_at(conn) + _migrate_retries(conn) + _migrate_task_options(conn) + _migrate_expires_at(conn) + _migrate_ttl(conn) + _migrate_claimed_at(conn) + conn.execute(""" + CREATE TABLE IF NOT EXISTS inline_objects ( + hash TEXT PRIMARY KEY, + data BLOB NOT NULL + ) + """) + + +def _migrate_last_accessed_at(conn: sqlite3.Connection) -> None: + col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] + if "last_accessed_at" not in col_names: + conn.execute( + "ALTER TABLE commits ADD COLUMN last_accessed_at TEXT DEFAULT NULL" + ) + conn.execute( + "UPDATE commits SET last_accessed_at = created_at WHERE last_accessed_at IS NULL" + ) + conn.execute( + "CREATE INDEX IF NOT EXISTS idx_last_accessed_at ON commits(last_accessed_at)" + ) + + +def _migrate_retries(conn: sqlite3.Connection) -> None: + col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] + if "retries" not in col_names: + conn.execute( + "ALTER TABLE commits ADD COLUMN retries INTEGER NOT NULL DEFAULT 0" + ) + + +def _migrate_task_options(conn: sqlite3.Connection) -> None: + col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] + if "force" not in col_names: + conn.execute( + "ALTER TABLE commits ADD COLUMN force INTEGER NOT NULL DEFAULT 0" + ) + if "timeout_seconds" not in col_names: + conn.execute("ALTER TABLE commits ADD COLUMN timeout_seconds REAL") + + +def _migrate_expires_at(conn: sqlite3.Connection) -> None: + col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] + if "expires_at" not in col_names: + conn.execute("ALTER TABLE commits ADD COLUMN expires_at TEXT") + + +def _migrate_ttl(conn: sqlite3.Connection) -> None: + col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] + if "ttl_seconds" not in col_names: + conn.execute("ALTER TABLE commits ADD COLUMN ttl_seconds REAL") + + +def _migrate_claimed_at(conn: sqlite3.Connection) -> None: + col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] + if "claimed_at" not in col_names: + conn.execute( + "ALTER TABLE commits ADD COLUMN claimed_at TEXT NOT NULL DEFAULT ''" + ) + conn.execute( + "UPDATE commits SET claimed_at = created_at WHERE claimed_at = ''" + ) + + +def put_commit_row(conn: sqlite3.Connection, commit: Commit, accessed_at: str) -> None: + output_hash = commit.output_ref.hash if commit.output_ref else None + output_size = commit.output_ref.size if commit.output_ref else None + output_tier = commit.output_ref.tier.value if commit.output_ref else None + timeout_seconds = ( + commit.task_def.timeout.total_seconds() if commit.task_def.timeout else None + ) + conn.execute( + """INSERT OR REPLACE INTO commits + (hash, fingerprint, func_name, func_hash, args_hash, args_snapshot, + func_source, dep_versions, cache, retries, force, timeout_seconds, + ttl_seconds, input_refs, output_hash, output_size, output_tier, parent_hash, + status, error, tags, created_at, + claimed_at, last_accessed_at, expires_at) + VALUES ( + ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? + )""", + ( + commit.hash, + commit.fingerprint, + commit.task_def.func_name, + commit.task_def.func_hash, + commit.task_def.args_hash, + commit.task_def.args_snapshot, + commit.task_def.func_source, + json.dumps(commit.task_def.dep_versions), + int(commit.task_def.cache), + commit.task_def.retries, + int(commit.task_def.force), + timeout_seconds, + commit.task_def.ttl.total_seconds() if commit.task_def.ttl else None, + json.dumps([r.hash for r in commit.input_refs]), + output_hash, + output_size, + output_tier, + commit.parent_hash, + commit.status.value, + commit.error, + json.dumps(commit.tags), + commit.created_at.isoformat(), + commit.claimed_at.isoformat(), + accessed_at, + commit.expires_at.isoformat() if commit.expires_at else None, + ), + ) + + +def row_to_commit(row: sqlite3.Row) -> Commit: + output_ref = None + if row["output_hash"]: + tier = StorageTier(row["output_tier"]) if row["output_tier"] else StorageTier.BLOB + output_ref = ObjectRef( + hash=row["output_hash"], + size=row["output_size"] or 0, + tier=tier, + ) + input_refs: list[ObjectRef] = [] + if row["input_refs"]: + for h in json.loads(row["input_refs"]): + input_refs.append(ObjectRef(hash=h)) + dep_versions: dict[str, str] = {} + if row["dep_versions"]: + dep_versions = json.loads(row["dep_versions"]) + tags: dict[str, str] = {} + if row["tags"]: + tags = json.loads(row["tags"]) + row_keys = set(row.keys()) + timeout_seconds = row["timeout_seconds"] if "timeout_seconds" in row_keys else None + ttl_seconds = row["ttl_seconds"] if "ttl_seconds" in row_keys else None + task_def = TaskDef( + func_hash=row["func_hash"], + func_name=row["func_name"], + func_source=row["func_source"] or "", + args_hash=row["args_hash"], + args_snapshot=row["args_snapshot"] or b"", + dep_versions=dep_versions, + cache=bool(row["cache"]), + tags=tags, + retries=row["retries"] if "retries" in row_keys else 0, + force=bool(row["force"]) if "force" in row_keys else False, + timeout=timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None, + ttl=timedelta(seconds=ttl_seconds) if ttl_seconds is not None else None, + ) + created_at = row["created_at"] + if isinstance(created_at, str): + created_at = datetime.fromisoformat(created_at) + claimed_at = row["claimed_at"] + if isinstance(claimed_at, str): + claimed_at = datetime.fromisoformat(claimed_at) + expires_at = row["expires_at"] + if isinstance(expires_at, str): + expires_at = datetime.fromisoformat(expires_at) + return Commit( + hash=row["hash"], + task_def=task_def, + input_refs=input_refs, + output_ref=output_ref, + parent_hash=row["parent_hash"], + status=TaskStatus(row["status"]), + created_at=created_at, + claimed_at=claimed_at, + error=row["error"], + tags=tags, + expires_at=expires_at, + ) diff --git a/src/cashet/async_client.py b/src/cashet/async_client.py index deb8696..02ba28f 100644 --- a/src/cashet/async_client.py +++ b/src/cashet/async_client.py @@ -24,9 +24,10 @@ from cashet._export import ImportResult, export_store, import_store from cashet.async_executor import AsyncLocalExecutor from cashet.dag import AsyncResultRef -from cashet.hashing import PickleSerializer, Serializer, build_task_def, warn_default_pickle +from cashet.hashing import build_task_def from cashet.models import Commit, TaskError, TaskStatus from cashet.protocols import AsyncExecutor, AsyncStore +from cashet.serializers import PickleSerializer, Serializer, warn_default_pickle from cashet.store import AsyncSQLiteStore T = TypeVar("T") diff --git a/src/cashet/async_executor.py b/src/cashet/async_executor.py index 8a975dd..e95c9fa 100644 --- a/src/cashet/async_executor.py +++ b/src/cashet/async_executor.py @@ -14,12 +14,11 @@ from cashet.dag import ( build_commit, find_existing_commit, - find_parent_hash, resolve_input_refs, ) -from cashet.hashing import Serializer from cashet.models import Commit, ObjectRef, TaskDef, TaskStatus from cashet.protocols import AsyncStore +from cashet.serializers import Serializer _DEFAULT_RUNNING_TTL = timedelta(seconds=300) _lock_cache: weakref.WeakKeyDictionary[Any, asyncio.Lock] = weakref.WeakKeyDictionary() @@ -103,28 +102,28 @@ async def submit( ) -> tuple[Commit, bool]: fp = task_def.fingerprint func_name = task_def.func_name + # Completed commits are immutable, so a hit needs no cross-process lock + # and no write; the claim loop below re-checks under the lock on a miss. if not task_def.force: - async with _async_store_lock(store, fp): - existing = await find_existing_commit(store, task_def) - if existing is not None: - existing.status = TaskStatus.CACHED - await store.put_commit(existing) - logger.info( - "task cached fingerprint=%s func=%s commit=%s", - fp, - func_name, - existing.hash[:12], - ) - return existing, True + existing = await find_existing_commit(store, task_def) + if existing is not None: + existing.status = TaskStatus.CACHED + logger.info( + "task cached fingerprint=%s func=%s commit=%s", + fp, + func_name, + existing.hash[:12], + ) + return existing, True while True: async with _async_store_lock(store, fp): - if not task_def.force: - existing = await find_existing_commit(store, task_def) - if existing is not None: - existing.status = TaskStatus.CACHED - await store.put_commit(existing) - return existing, True + # One lookup serves both purposes: it is the cached result when + # allowed to reuse it, and the parent for the new claim otherwise. + latest = await store.find_by_fingerprint(fp) + if task_def.cache and not task_def.force and latest is not None: + latest.status = TaskStatus.CACHED + return latest, True claim = await store.find_running_by_fingerprint(task_def.fingerprint) if claim is not None: @@ -151,8 +150,11 @@ async def submit( break else: input_refs = resolve_input_refs(args, kwargs) - parent_hash = await find_parent_hash(store, task_def) - claim = build_commit(task_def, input_refs, parent_hash=parent_hash) + claim = build_commit( + task_def, + input_refs, + parent_hash=latest.hash if latest else None, + ) claim.status = TaskStatus.RUNNING await store.put_commit(claim) logger.debug( @@ -205,19 +207,34 @@ async def _execute( async def _heartbeat() -> None: interval = self._running_ttl.total_seconds() / 2 + # A failed renewal leaves the claim halfway to stale; waiting a full + # interval would retry exactly at the staleness boundary, where any + # delay lets another worker reclaim a live task. Retry fast instead. + retry_interval = interval / 4 + wait = interval while True: try: - await asyncio.wait_for(stop_event.wait(), timeout=interval) + await asyncio.wait_for(stop_event.wait(), timeout=wait) break except TimeoutError: pass if commit.status != TaskStatus.RUNNING: break - async with _async_store_lock(store, commit.task_def.fingerprint): - if commit.status != TaskStatus.RUNNING: - break - commit.claimed_at = datetime.now(UTC) - await store.put_commit(commit) + try: + async with _async_store_lock(store, commit.task_def.fingerprint): + if commit.status != TaskStatus.RUNNING: + break + commit.claimed_at = datetime.now(UTC) + await store.put_commit(commit) + wait = interval + except Exception: + wait = retry_interval + logger.warning( + "heartbeat renewal failed fingerprint=%s commit=%s", + commit.task_def.fingerprint, + commit.hash[:12], + exc_info=True, + ) heartbeat_task = asyncio.create_task(_heartbeat()) @@ -286,7 +303,24 @@ async def _heartbeat() -> None: commit.error, ) stop_event.set() - await asyncio.wait_for(heartbeat_task, timeout=self._running_ttl.total_seconds() * 2) + # The task's outcome is already decided; a misbehaving heartbeat must + # not raise past this point and destroy the result. + try: + await asyncio.wait_for( + heartbeat_task, timeout=self._running_ttl.total_seconds() * 2 + ) + except TimeoutError: + logger.warning( + "heartbeat did not stop in time fingerprint=%s commit=%s", + commit.task_def.fingerprint, + commit.hash[:12], + ) + except Exception: + logger.exception( + "heartbeat task failed fingerprint=%s commit=%s", + commit.task_def.fingerprint, + commit.hash[:12], + ) return commit diff --git a/src/cashet/cli.py b/src/cashet/cli.py index b19f3d2..6cbbc97 100644 --- a/src/cashet/cli.py +++ b/src/cashet/cli.py @@ -8,13 +8,20 @@ from rich.syntax import Syntax from rich.table import Table +from cashet._client_base import resolve_store_dir from cashet.client import Client console = Console() -def _client() -> Client: - return Client() +def _client(require_store: bool = True) -> Client: + ctx = click.get_current_context(silent=True) + store_dir_opt = ctx.obj if ctx is not None else None + store_dir = resolve_store_dir(store_dir_opt, None) + if require_store and not store_dir.exists(): + console.print(f"[red]No cashet store found at {store_dir}.[/red]") + raise SystemExit(1) + return Client(store_dir=store_dir) def _parse_tags(tags: tuple[str, ...]) -> dict[str, str | None] | None: @@ -35,9 +42,15 @@ def _parse_tags(tags: tuple[str, ...]) -> dict[str, str | None] | None: @click.group() -def main() -> None: +@click.option( + "--store-dir", + default=None, + help="Store directory (defaults to $CASHET_DIR or ./.cashet)", +) +@click.pass_context +def main(ctx: click.Context, store_dir: str | None) -> None: """cashet — content-addressable compute cache with git semantics""" - pass + ctx.obj = store_dir @main.command("serve") @@ -53,7 +66,7 @@ def serve_cmd( host: str, port: int, require_token: str | None, allow_remote_code: bool ) -> None: """Start the HTTP server""" - client = _client() + client = _client(require_store=False) try: client.serve( host=host, @@ -355,7 +368,7 @@ def import_archive(path: str) -> None: """Import commits and blobs from a tar.gz archive""" import tarfile - client = _client() + client = _client(require_store=False) try: result = client.import_archive(path) except (OSError, tarfile.TarError, ValueError) as e: diff --git a/src/cashet/client.py b/src/cashet/client.py index c8d5e29..bacf836 100644 --- a/src/cashet/client.py +++ b/src/cashet/client.py @@ -23,9 +23,9 @@ from cashet.async_executor import AsyncLocalExecutor from cashet.dag import ResultRef from cashet.executor import LocalExecutor -from cashet.hashing import Serializer from cashet.models import Commit, TaskStatus from cashet.protocols import AsyncExecutor, AsyncStore, Executor, Store +from cashet.serializers import Serializer from cashet.store import AsyncSQLiteStore, SQLiteStore T = TypeVar("T") diff --git a/src/cashet/dag.py b/src/cashet/dag.py index 6a9711c..69cd073 100644 --- a/src/cashet/dag.py +++ b/src/cashet/dag.py @@ -6,9 +6,10 @@ from datetime import UTC, datetime from typing import Any, Generic, TypeVar -from cashet.hashing import Serializer, object_state +from cashet.hashing import object_state from cashet.models import Commit, ObjectRef, TaskDef, TaskStatus from cashet.protocols import AsyncStore +from cashet.serializers import Serializer T = TypeVar("T") @@ -174,13 +175,6 @@ async def find_existing_commit(store: AsyncStore, task_def: TaskDef) -> Commit | return await store.find_by_fingerprint(task_def.fingerprint) -async def find_parent_hash(store: AsyncStore, task_def: TaskDef) -> str | None: - existing = await store.find_by_fingerprint(task_def.fingerprint) - if existing is not None: - return existing.hash - return None - - def compute_commit_hash( task_def: TaskDef, input_refs: list[ObjectRef], diff --git a/src/cashet/executor.py b/src/cashet/executor.py index 5b9631b..8ae848a 100644 --- a/src/cashet/executor.py +++ b/src/cashet/executor.py @@ -6,9 +6,9 @@ from cashet._runner import BlockingAsyncRunner from cashet.adapters import SyncStoreAdapter from cashet.async_executor import AsyncLocalExecutor -from cashet.hashing import Serializer from cashet.models import Commit, TaskDef from cashet.protocols import Store +from cashet.serializers import Serializer class LocalExecutor: diff --git a/src/cashet/hashing.py b/src/cashet/hashing.py index ffb893b..4deb04c 100644 --- a/src/cashet/hashing.py +++ b/src/cashet/hashing.py @@ -1,6 +1,7 @@ from __future__ import annotations import ast +import contextlib import datetime as _datetime import hashlib import inspect @@ -11,11 +12,17 @@ import textwrap import types import warnings +import weakref from datetime import timedelta from functools import lru_cache -from typing import Any, Protocol, runtime_checkable +from typing import Any from cashet.models import TaskDef +from cashet.serializers import JsonSerializer as JsonSerializer +from cashet.serializers import PickleSerializer as PickleSerializer +from cashet.serializers import SafePickleSerializer as SafePickleSerializer +from cashet.serializers import Serializer as Serializer +from cashet.serializers import warn_default_pickle as warn_default_pickle class ClosureWarning(UserWarning): @@ -26,136 +33,6 @@ class UnhashableArgWarning(UserWarning): pass -_pickle_warning_issued = False - - -def warn_default_pickle() -> None: - global _pickle_warning_issued - if _pickle_warning_issued: - return - _pickle_warning_issued = True - import warnings - - warnings.warn( - "Using PickleSerializer by default — arbitrary code execution risk on " - "untrusted cached results. Pass serializer=SafePickleSerializer() " - "for safer deserialization.", - stacklevel=3, - ) - - -@runtime_checkable -class Serializer(Protocol): - def dumps(self, obj: Any) -> bytes: ... - def loads(self, data: bytes) -> Any: ... - - -class PickleSerializer: - def __init__(self, protocol: int | None = None) -> None: - import pickle - - self._pickle = pickle - self._protocol = protocol or pickle.HIGHEST_PROTOCOL - - def dumps(self, obj: Any) -> bytes: - return self._pickle.dumps(obj, protocol=self._protocol) - - def loads(self, data: bytes) -> Any: - return self._pickle.loads(data) - - -class JsonSerializer: - def dumps(self, obj: Any) -> bytes: - import json - - return json.dumps(obj, default=str, sort_keys=True).encode() - - def loads(self, data: bytes) -> Any: - import json - - return json.loads(data) - - -class SafePickleSerializer: - _cached_allowlist: list[type] | None = None - - def __init__(self, extra_classes: list[type] | None = None) -> None: - import pickle - - self._pickle = pickle - self._allowed: dict[str, type] = {} - for cls in self._default_allowlist(): - key = f"{cls.__module__}.{cls.__qualname__}" - self._allowed[key] = cls - if extra_classes: - for cls in extra_classes: - key = f"{cls.__module__}.{cls.__qualname__}" - self._allowed[key] = cls - - def dumps(self, obj: Any) -> bytes: - return self._pickle.dumps(obj, protocol=self._pickle.HIGHEST_PROTOCOL) - - def loads(self, data: bytes) -> Any: - import io - import pickle - - allowed = self._allowed - blocked_msg = " — not in allowlist. Pass it via SafePickleSerializer(extra_classes=[...])." - - class _RestrictedUnpickler(pickle.Unpickler): - def find_class(self, module: str, name: str) -> Any: # type: ignore[override] - key = f"{module}.{name}" - if key in allowed: - return allowed[key] - raise pickle.UnpicklingError(f"Blocked class {key}{blocked_msg}") - - return _RestrictedUnpickler(io.BytesIO(data)).load() - - @classmethod - def _default_allowlist(cls) -> list[type]: - if cls._cached_allowlist is not None: - return cls._cached_allowlist - import collections - import datetime - - types_list: list[type] = [ - type(None), - bool, - int, - float, - str, - bytes, - bytearray, - list, - dict, - tuple, - set, - frozenset, - slice, - range, - complex, - object, - type, - datetime.datetime, - datetime.date, - datetime.timedelta, - datetime.time, - datetime.timezone, - collections.OrderedDict, - collections.defaultdict, - collections.Counter, - collections.deque, - ] - try: - import numpy # pyright: ignore[reportMissingImports] - - types_list.append(numpy.ndarray) # type: ignore[attr-defined] - except ImportError: - pass - cls._cached_allowlist = types_list - return types_list - - def _normalize_source(source: str) -> str: return textwrap.dedent(source).strip() @@ -171,7 +48,16 @@ def _bytecode_source(func: types.FunctionType) -> str: ) +# Keyed by the function object itself: a redefinition (new cell, reloaded +# module) is a new object and misses, while repeat submissions of the same +# function skip the source lookup and its file IO entirely. +_source_cache: weakref.WeakKeyDictionary[types.FunctionType, str] = weakref.WeakKeyDictionary() + + def get_func_source(func: types.FunctionType) -> str: + cached = _source_cache.get(func) + if cached is not None: + return cached try: source = inspect.getsource(func) except OSError: @@ -184,7 +70,10 @@ def get_func_source(func: types.FunctionType) -> str: pass if source is None: source = _bytecode_source(func) - return _normalize_source(source) + normalized = _normalize_source(source) + with contextlib.suppress(TypeError): + _source_cache[func] = normalized + return normalized def get_dep_versions(func: types.FunctionType) -> dict[str, str]: @@ -207,6 +96,7 @@ def hash_source(source: str) -> str: return hashlib.sha256(source.encode()).hexdigest() +@lru_cache(maxsize=1024) def _ast_canonical(source: str) -> str: # ast.unparse normalizes whitespace and comments like ast.dump but, being # source text rather than the internal AST repr, stays stable across Python @@ -237,15 +127,20 @@ def _strip_docstrings(node: ast.AST) -> None: child.body = [ast.Pass()] -def _is_stdlib_or_site_path(path: str) -> bool: - resolved = os.path.abspath(path) +@lru_cache(maxsize=1) +def _stdlib_and_site_prefixes() -> tuple[str, ...]: stdlib_path = os.path.abspath(os.path.dirname(os.__file__)) site_paths = [os.path.abspath(p) for p in site.getsitepackages() if p] user_site = site.getusersitepackages() if user_site: site_paths.append(os.path.abspath(user_site)) - excluded = [stdlib_path, *site_paths] - for prefix in excluded: + return (stdlib_path, *site_paths) + + +@lru_cache(maxsize=4096) +def _is_stdlib_or_site_path(path: str) -> bool: + resolved = os.path.abspath(path) + for prefix in _stdlib_and_site_prefixes(): try: if os.path.commonpath([resolved, prefix]) == prefix: return True @@ -425,6 +320,17 @@ def object_state(obj: Any) -> dict[str, Any] | None: return state or None +def _stable_item_reprs(items: Any, _visited: set[int]) -> list[str]: + # Set ordering must come from the items' stable serialized form; raw repr + # can embed memory addresses, which reorder across processes. + reprs: list[str] = [] + for item in items: + sub = io.StringIO() + _stable_repr_to(sub, item, _visited) + reprs.append(sub.getvalue()) + return reprs + + def _stable_repr_to( buf: io.StringIO, obj: Any, _visited: set[int] | None = None ) -> None: @@ -456,12 +362,7 @@ def _stable_repr_to( return _visited.add(obj_id) buf.write("{") - first = True - for item in sorted(obj, key=repr): - if not first: - buf.write(", ") - first = False - _stable_repr_to(buf, item, _visited) + buf.write(", ".join(sorted(_stable_item_reprs(obj, _visited)))) buf.write("}") _visited.discard(obj_id) elif isinstance(obj, frozenset): @@ -471,12 +372,7 @@ def _stable_repr_to( return _visited.add(obj_id) buf.write("frozenset({") - first = True - for item in sorted(obj, key=repr): - if not first: - buf.write(", ") - first = False - _stable_repr_to(buf, item, _visited) + buf.write(", ".join(sorted(_stable_item_reprs(obj, _visited)))) buf.write("})") _visited.discard(obj_id) elif isinstance(obj, dict): diff --git a/src/cashet/redis_store.py b/src/cashet/redis_store.py index 334230e..a8c693e 100644 --- a/src/cashet/redis_store.py +++ b/src/cashet/redis_store.py @@ -1,291 +1,44 @@ from __future__ import annotations -import base64 import hashlib -import json import logging -from datetime import UTC, datetime, timedelta +from datetime import UTC, datetime from typing import Any from redis.exceptions import WatchError +from cashet._ids import normalize_hash_prefix +from cashet._redis_codec import ( + DECR_DELETE_SCRIPT, + access_key, + blob_hashes, + blob_key, + blob_ref_key, + blob_stats_key, + blob_stats_lock_key, + commit_access_timestamp, + commit_hash_from_key, + commit_key, + decode_commit, + decode_hash, + expires_key, + fp_key, + func_key, + index_commit_commands, + matches_tags, + remove_commit_index_commands, + running_key, + stats_dict, + stats_ready, + status_key, + tag_key, + tag_value_key, +) from cashet._runner import BlockingAsyncRunner -from cashet.models import Commit, ObjectRef, StorageTier, TaskDef, TaskStatus +from cashet.models import Commit, ObjectRef, StorageTier, TaskStatus logger = logging.getLogger("cashet") -_DECR_DELETE_SCRIPT = """ - local ref_key = KEYS[1] - local blob_key = KEYS[2] - local stats_key = KEYS[3] - if redis.call('EXISTS', ref_key) == 0 then - return 0 - end - local count = redis.call('DECR', ref_key) - if count <= 0 then - local existed = redis.call('EXISTS', blob_key) - local bytes = 0 - if existed == 1 then - bytes = redis.call('STRLEN', blob_key) - end - redis.call('DEL', blob_key, ref_key) - if existed == 1 and redis.call('HGET', stats_key, 'ready') == '1' then - redis.call('HINCRBY', stats_key, 'objects', -1) - redis.call('HINCRBY', stats_key, 'bytes', -bytes) - end - return bytes - end - return 0 -""" - - -def _commit_key(hash: str) -> str: - return f"cashet:commit:{hash}" - - -def _normalize_hash_prefix(hash: str) -> str | None: - if 0 < len(hash) <= 64 and all(c in "0123456789abcdefABCDEF" for c in hash): - return hash.lower() - return None - - -def _blob_key(hash: str) -> str: - return f"cashet:blob:data:{hash}" - - -def _fp_key(fingerprint: str) -> str: - return f"cashet:index:fingerprint:{fingerprint}" - - -def _running_key(fingerprint: str) -> str: - return f"cashet:index:running:{fingerprint}" - - -def _func_key(func_name: str) -> str: - return f"cashet:index:func:{func_name}" - - -def _status_key(status: str) -> str: - return f"cashet:index:status:{status}" - - -def _tag_key(key: str) -> str: - return f"cashet:tagk:{key}" - - -def _tag_value_key(key: str, value: str) -> str: - # Length-prefix the key so a ':' inside a tag key or value can never make - # two distinct (key, value) pairs collide onto the same set. - return f"cashet:tagv:{len(key)}:{key}:{value}" - - -def _access_key() -> str: - return "cashet:index:last_accessed" - - -def _blob_stats_key() -> str: - return "cashet:stats:blob" - - -def _blob_stats_lock_key() -> str: - return "cashet:stats:blob:lock" - - -def _stats_ready(raw: Any) -> bool: - if isinstance(raw, bytes): - return raw == b"1" - return raw == "1" or raw == 1 - - -def _encode_commit(commit: Commit) -> bytes: - d: dict[str, Any] = { - "hash": commit.hash, - "fingerprint": commit.fingerprint, - "func_name": commit.task_def.func_name, - "func_hash": commit.task_def.func_hash, - "args_hash": commit.task_def.args_hash, - "args_snapshot_b64": base64.b64encode(commit.task_def.args_snapshot).decode(), - "func_source": commit.task_def.func_source, - "dep_versions": commit.task_def.dep_versions, - "cache": commit.task_def.cache, - "retries": commit.task_def.retries, - "force": commit.task_def.force, - "timeout_seconds": ( - commit.task_def.timeout.total_seconds() if commit.task_def.timeout else None - ), - "ttl_seconds": ( - commit.task_def.ttl.total_seconds() if commit.task_def.ttl else None - ), - "expires_at": commit.expires_at.isoformat() if commit.expires_at else None, - "input_refs": [ - {"hash": r.hash, "size": r.size, "tier": r.tier.value} for r in commit.input_refs - ], - "output_hash": commit.output_ref.hash if commit.output_ref else None, - "output_size": commit.output_ref.size if commit.output_ref else None, - "output_tier": commit.output_ref.tier.value if commit.output_ref else None, - "parent_hash": commit.parent_hash, - "status": commit.status.value, - "error": commit.error, - "tags": commit.tags, - "created_at": commit.created_at.isoformat(), - "claimed_at": commit.claimed_at.isoformat(), - "last_accessed_at": datetime.now(UTC).isoformat(), - } - return json.dumps(d, separators=(",", ":")).encode() - - -def _decode_commit(data: bytes) -> Commit: - d = json.loads(data) - task_def = TaskDef( - func_hash=d["func_hash"], - func_name=d["func_name"], - func_source=d.get("func_source", ""), - args_hash=d["args_hash"], - args_snapshot=base64.b64decode(d.get("args_snapshot_b64", "")), - dep_versions=d.get("dep_versions", {}), - cache=d.get("cache", True), - tags=d.get("tags", {}), - retries=d.get("retries", 0), - force=d.get("force", False), - timeout=( - timedelta(seconds=d["timeout_seconds"]) - if d.get("timeout_seconds") is not None - else None - ), - ttl=( - timedelta(seconds=d["ttl_seconds"]) - if d.get("ttl_seconds") is not None - else None - ), - ) - input_refs = [ - ObjectRef( - hash=r["hash"], - size=r.get("size", 0), - tier=StorageTier(r.get("tier", "blob")), - ) - for r in d.get("input_refs", []) - ] - output_ref = None - if d.get("output_hash"): - output_ref = ObjectRef( - hash=d["output_hash"], - size=d.get("output_size", 0), - tier=StorageTier(d.get("output_tier", "blob")), - ) - created_at = ( - datetime.fromisoformat(d["created_at"]) if "created_at" in d else datetime.now(UTC) - ) - claimed_at = ( - datetime.fromisoformat(d["claimed_at"]) if "claimed_at" in d else datetime.now(UTC) - ) - expires_at = ( - datetime.fromisoformat(d["expires_at"]) if d.get("expires_at") is not None else None - ) - return Commit( - hash=d["hash"], - task_def=task_def, - input_refs=input_refs, - output_ref=output_ref, - parent_hash=d.get("parent_hash"), - status=TaskStatus(d["status"]), - created_at=created_at, - claimed_at=claimed_at, - error=d.get("error"), - tags=d.get("tags", {}), - expires_at=expires_at, - ) - - -def _decode_hash(raw: Any) -> str: - return raw.decode() if isinstance(raw, bytes) else raw - - -def _commit_access_timestamp(data: bytes) -> float: - d = json.loads(data) - value = d.get("last_accessed_at") or d.get("created_at") - if value is None: - return datetime.now(UTC).timestamp() - return datetime.fromisoformat(value).timestamp() - - -def _blob_ref_key(blob_hash: str) -> str: - return f"cashet:blob:ref:{blob_hash}" - - -def _blob_hashes(commit: Commit) -> set[str]: - hashes: set[str] = set() - if commit.output_ref: - hashes.add(commit.output_ref.hash) - for ref in commit.input_refs: - hashes.add(ref.hash) - return hashes - - -def _matches_tags(commit: Commit, tags: dict[str, str | None]) -> bool: - for key, val in tags.items(): - if val is None: - if key not in commit.tags: - return False - else: - if commit.tags.get(key) != val: - return False - return True - - -def _stats_dict(total: int, completed: int, blob_count: int, blob_bytes: int) -> dict[str, int]: - return { - "total_commits": total, - "completed_commits": completed, - "stored_objects": blob_count, - "disk_bytes": blob_bytes, - "blob_objects": blob_count, - "blob_bytes": blob_bytes, - "inline_objects": 0, - "inline_bytes": 0, - } - - -def _index_commit_commands(pipe: Any, commit: Commit) -> None: - pipe.get(_commit_key(commit.hash)) - pipe.set(_commit_key(commit.hash), _encode_commit(commit)) - ts = commit.created_at.timestamp() - pipe.zadd("cashet:index:all", {commit.hash: ts}) - expires_ts = commit.expires_at.timestamp() if commit.expires_at else float("inf") - pipe.zadd(_fp_key(commit.fingerprint), {commit.hash: expires_ts}) - pipe.zadd(_func_key(commit.task_def.func_name), {commit.hash: ts}) - now_ts = datetime.now(UTC).timestamp() - pipe.zadd(_access_key(), {commit.hash: now_ts}) - for status in TaskStatus: - pipe.srem(_status_key(status.value), commit.hash) - pipe.sadd(_status_key(commit.status.value), commit.hash) - if commit.status == TaskStatus.RUNNING: - pipe.sadd(_running_key(commit.fingerprint), commit.hash) - else: - pipe.srem(_running_key(commit.fingerprint), commit.hash) - for key, val in commit.tags.items(): - pipe.sadd(_tag_key(key), commit.hash) - pipe.sadd(_tag_value_key(key, val), commit.hash) - - -def _remove_commit_index_commands(pipe: Any, commit: Commit, resolved_hash: str) -> None: - pipe.delete(_commit_key(resolved_hash)) - pipe.zrem("cashet:index:all", resolved_hash) - pipe.zrem(_fp_key(commit.fingerprint), resolved_hash) - pipe.srem(_running_key(commit.fingerprint), resolved_hash) - pipe.zrem(_func_key(commit.task_def.func_name), resolved_hash) - pipe.zrem(_access_key(), resolved_hash) - for status in TaskStatus: - pipe.srem(_status_key(status.value), resolved_hash) - for key, val in commit.tags.items(): - pipe.srem(_tag_key(key), resolved_hash) - pipe.srem(_tag_value_key(key, val), resolved_hash) - - -def _commit_hash_from_key(raw: Any) -> str: - key_str = raw.decode() if isinstance(raw, bytes) else str(raw) - return key_str.split(":")[-1] - class AsyncRedisStore: def __init__( @@ -295,22 +48,20 @@ def __init__( self._redis: Any = aioredis.from_url(redis_url) self._lock_timeout = lock_timeout - self._async_locks: dict[str, Any] = {} def _fingerprint_lock(self, fingerprint: str) -> Any: - lock = self._async_locks.get(fingerprint) - if lock is None: - lock = self._redis.lock( - f"cashet:lock:{fingerprint}", - timeout=self._lock_timeout, - blocking_timeout=10, - ) - self._async_locks[fingerprint] = lock - return lock + # A fresh Lock per acquisition: redis-py Lock objects carry per-holder + # token state, so caching one per fingerprint leaked memory and risked + # token confusion when two coroutines reused the same instance. + return self._redis.lock( + f"cashet:lock:{fingerprint}", + timeout=self._lock_timeout, + blocking_timeout=10, + ) async def put_blob(self, data: bytes) -> ObjectRef: content_hash = hashlib.sha256(data).hexdigest() - key = _blob_key(content_hash) + key = blob_key(content_hash) if await self._blob_stats_ready(): stored = await self._redis.set(key, data, nx=True) if stored: @@ -336,7 +87,7 @@ async def put_blob(self, data: bytes) -> ObjectRef: return ObjectRef(hash=content_hash, size=len(data), tier=StorageTier.BLOB) async def get_blob(self, ref: ObjectRef) -> bytes: - data = await self._redis.get(_blob_key(ref.hash)) + data = await self._redis.get(blob_key(ref.hash)) if data is None: logger.warning("blob not found hash=%s", ref.hash[:12]) raise ValueError(f"Blob {ref.hash} not found") @@ -348,21 +99,21 @@ async def get_blob(self, ref: ObjectRef) -> bytes: return data async def put_commit(self, commit: Commit) -> None: - ck = _commit_key(commit.hash) - new_hashes = _blob_hashes(commit) + ck = commit_key(commit.hash) + new_hashes = blob_hashes(commit) while True: async with self._redis.pipeline(transaction=True) as pipe: await pipe.watch(ck) existing_raw = await pipe.get(ck) pipe.multi() - _index_commit_commands(pipe, commit) + index_commit_commands(pipe, commit) if existing_raw is None: for h in new_hashes: - pipe.incr(_blob_ref_key(h)) + pipe.incr(blob_ref_key(h)) else: - old_hashes = _blob_hashes(_decode_commit(existing_raw)) + old_hashes = blob_hashes(decode_commit(existing_raw)) for h in new_hashes - old_hashes: - pipe.incr(_blob_ref_key(h)) + pipe.incr(blob_ref_key(h)) try: await pipe.execute() return @@ -370,48 +121,72 @@ async def put_commit(self, commit: Commit) -> None: continue async def get_commit(self, hash: str) -> Commit | None: - normalized = _normalize_hash_prefix(hash) + normalized = normalize_hash_prefix(hash) if normalized is None: return None hash = normalized if len(hash) < 64: matches: list[str] = [] - async for key in self._redis.scan_iter(match=_commit_key(hash) + "*", count=100): + async for key in self._redis.scan_iter(match=commit_key(hash) + "*", count=100): matches.append(key) if not matches: return None if len(matches) > 1: hashes: list[str] = [] for m in matches: - hashes.append(_commit_hash_from_key(m)) + hashes.append(commit_hash_from_key(m)) matches_str = ", ".join(h[:12] for h in hashes) raise ValueError( f"Ambiguous prefix {hash[:12]} matches {len(matches)} commits: {matches_str}" ) data = await self._redis.get(matches[0]) else: - data = await self._redis.get(_commit_key(hash)) + data = await self._redis.get(commit_key(hash)) if data is None: return None - return _decode_commit(data) + return decode_commit(data) async def find_by_fingerprint(self, fingerprint: str) -> Commit | None: - now_ts = datetime.now(UTC).timestamp() - hashes = await self._redis.zrevrangebyscore( - _fp_key(fingerprint), max="+inf", min=f"({now_ts}" - ) - for h in hashes: - h_str = h.decode() if isinstance(h, bytes) else h + # Scored by created_at, so descending order is recency order and the + # newest completed commit wins, matching SQLiteStore. + entries = await self._redis.zrevrange(fp_key(fingerprint), 0, -1, withscores=True) + if any(score == float("inf") for _, score in entries): + candidates = await self._rescore_legacy_fingerprint_index(fingerprint, entries) + else: + candidates = [decode_hash(h) for h, _score in entries] + now = datetime.now(UTC) + for h_str in candidates: commit = await self.get_commit(h_str) - if commit is not None and commit.status in (TaskStatus.COMPLETED, TaskStatus.CACHED): - if commit.expires_at is not None and commit.expires_at <= datetime.now(UTC): - continue - await self._touch_commit(h_str) - return commit + if commit is None: + continue + if commit.status not in (TaskStatus.COMPLETED, TaskStatus.CACHED): + continue + if commit.expires_at is not None and commit.expires_at <= now: + continue + await self._touch_commit(h_str) + return commit return None + async def _rescore_legacy_fingerprint_index( + self, fingerprint: str, entries: list[tuple[Any, float]] + ) -> list[str]: + # Indexes written before 0.5.0 scored this zset by expiry (infinity for + # commits without a TTL), which would shadow every newer commit; rescore + # by created_at so recency ordering holds for pre-upgrade entries, and + # return the rescored recency order so the lookup needs no re-read. + hashes = [decode_hash(h) for h, _score in entries] + bodies = await self._redis.mget([commit_key(h) for h in hashes]) + rescored: dict[str, float] = {} + for h_str, data in zip(hashes, bodies, strict=True): + if data is None: + continue + rescored[h_str] = decode_commit(data).created_at.timestamp() + if rescored: + await self._redis.zadd(fp_key(fingerprint), rescored) + return sorted(rescored, key=lambda h: (rescored[h], h), reverse=True) + async def find_running_by_fingerprint(self, fingerprint: str) -> Commit | None: - hashes = await self._redis.smembers(_running_key(fingerprint)) + hashes = await self._redis.smembers(running_key(fingerprint)) for h in hashes: h_str = h.decode() if isinstance(h, bytes) else h commit = await self.get_commit(h_str) @@ -427,7 +202,7 @@ async def list_commits( tags: dict[str, str | None] | None = None, ) -> list[Commit]: if func_name: - hashes = await self._redis.zrevrange(_func_key(func_name), 0, -1) + hashes = await self._redis.zrevrange(func_key(func_name), 0, -1) else: hashes = await self._redis.zrevrange("cashet:index:all", 0, -1) commits: list[Commit] = [] @@ -438,7 +213,7 @@ async def list_commits( continue if status is not None and commit.status != status: continue - if tags is not None and not _matches_tags(commit, tags): + if tags is not None and not matches_tags(commit, tags): continue commits.append(commit) if len(commits) >= limit: @@ -446,7 +221,7 @@ async def list_commits( return commits async def get_history(self, hash: str) -> list[Commit]: - normalized = _normalize_hash_prefix(hash) + normalized = normalize_hash_prefix(hash) if normalized is None: return [] hash = normalized @@ -454,7 +229,7 @@ async def get_history(self, hash: str) -> list[Commit]: if commit is None: return [] fingerprint = commit.fingerprint - hashes = await self._redis.zrange(_fp_key(fingerprint), 0, -1) + hashes = await self._redis.zrange(fp_key(fingerprint), 0, -1) now = datetime.now(UTC) commits: list[Commit] = [] for h in hashes: @@ -470,24 +245,24 @@ async def stats(self) -> dict[str, int]: total = await self._redis.zcard("cashet:index:all") completed = 0 for status in ("completed", "cached"): - completed += await self._redis.scard(_status_key(status)) + completed += await self._redis.scard(status_key(status)) blob_count, blob_bytes = await self._blob_storage_totals() - return _stats_dict(total, completed, blob_count, blob_bytes) + return stats_dict(total, completed, blob_count, blob_bytes) def _blob_stats_lock(self) -> Any: return self._redis.lock( - _blob_stats_lock_key(), + blob_stats_lock_key(), timeout=self._lock_timeout, blocking_timeout=10, ) async def _blob_stats_ready(self) -> bool: - return _stats_ready(await self._redis.hget(_blob_stats_key(), "ready")) + return stats_ready(await self._redis.hget(blob_stats_key(), "ready")) async def _incr_blob_stats(self, objects_delta: int, bytes_delta: int) -> None: pipe = self._redis.pipeline() - pipe.hincrby(_blob_stats_key(), "objects", objects_delta) - pipe.hincrby(_blob_stats_key(), "bytes", bytes_delta) + pipe.hincrby(blob_stats_key(), "objects", objects_delta) + pipe.hincrby(blob_stats_key(), "bytes", bytes_delta) await pipe.execute() async def _scan_blob_storage_totals(self) -> tuple[int, int]: @@ -506,24 +281,24 @@ async def _ensure_blob_stats(self) -> None: return blob_count, blob_bytes = await self._scan_blob_storage_totals() await self._redis.hset( - _blob_stats_key(), + blob_stats_key(), mapping={"objects": blob_count, "bytes": blob_bytes, "ready": 1}, ) async def _blob_storage_totals(self) -> tuple[int, int]: await self._ensure_blob_stats() blob_count, blob_bytes = await self._redis.hmget( - _blob_stats_key(), "objects", "bytes" + blob_stats_key(), "objects", "bytes" ) return int(blob_count or 0), int(blob_bytes or 0) async def _blob_size(self, blob_hash: str) -> int: - return await self._redis.strlen(_blob_key(blob_hash)) + return await self._redis.strlen(blob_key(blob_hash)) async def _bytes_freed_by_delete(self, commit: Commit) -> int: freed = 0 - for h in _blob_hashes(commit): - ref_count = int(await self._redis.get(_blob_ref_key(h)) or 0) + for h in blob_hashes(commit): + ref_count = int(await self._redis.get(blob_ref_key(h)) or 0) if ref_count <= 1: freed += await self._blob_size(h) return freed @@ -531,40 +306,51 @@ async def _bytes_freed_by_delete(self, commit: Commit) -> int: async def _backfill_access_index(self) -> None: all_hashes = await self._redis.zrange("cashet:index:all", 0, -1) for h in all_hashes: - h_str = _decode_hash(h) - score = await self._redis.zscore(_access_key(), h_str) + h_str = decode_hash(h) + score = await self._redis.zscore(access_key(), h_str) if score is None: - data = await self._redis.get(_commit_key(h_str)) + data = await self._redis.get(commit_key(h_str)) if data is None: continue - await self._redis.zadd(_access_key(), {h_str: _commit_access_timestamp(data)}) + await self._redis.zadd(access_key(), {h_str: commit_access_timestamp(data)}) async def evict(self, older_than: datetime, max_size_bytes: int | None = None) -> int: deleted = 0 total = await self._redis.zcard("cashet:index:all") - indexed = await self._redis.zcard(_access_key()) + indexed = await self._redis.zcard(access_key()) if total > 0 and indexed < total: await self._backfill_access_index() cutoff_ts = older_than.timestamp() - old_hashes_raw = await self._redis.zrangebyscore(_access_key(), "-inf", cutoff_ts) - old_hashes = [_decode_hash(h) for h in old_hashes_raw] + old_hashes_raw = await self._redis.zrangebyscore(access_key(), "-inf", cutoff_ts) + old_hashes = [decode_hash(h) for h in old_hashes_raw] for h_str in old_hashes: commit = await self.get_commit(h_str) if commit is None: - await self._redis.zrem(_access_key(), h_str) + await self._redis.zrem(access_key(), h_str) continue if await self.delete_commit(h_str): deleted += 1 + now_ts = datetime.now(UTC).timestamp() + expired_raw = await self._redis.zrangebyscore(expires_key(), "-inf", now_ts) + stale: list[str] = [] + for h in expired_raw: + h_str = decode_hash(h) + if await self.delete_commit(h_str): + deleted += 1 + else: + stale.append(h_str) + if stale: + await self._redis.zrem(expires_key(), *stale) if max_size_bytes is not None: current_bytes = (await self._blob_storage_totals())[1] while current_bytes > max_size_bytes: - candidates = await self._redis.zrange(_access_key(), 0, 0) + candidates = await self._redis.zrange(access_key(), 0, 0) if not candidates: break - oldest_hash = _decode_hash(candidates[0]) + oldest_hash = decode_hash(candidates[0]) commit = await self.get_commit(oldest_hash) if commit is None: - await self._redis.zrem(_access_key(), oldest_hash) + await self._redis.zrem(access_key(), oldest_hash) continue freed = await self._bytes_freed_by_delete(commit) if await self.delete_commit(oldest_hash): @@ -583,7 +369,7 @@ async def evict(self, older_than: datetime, max_size_bytes: int | None = None) - return deleted async def delete_commit(self, hash: str) -> bool: - normalized = _normalize_hash_prefix(hash) + normalized = normalize_hash_prefix(hash) if normalized is None: return False hash = normalized @@ -595,7 +381,7 @@ async def delete_commit(self, hash: str) -> bool: async def delete_by_tags(self, tags: dict[str, str | None]) -> int: set_keys: list[str] = [] for key, val in tags.items(): - set_keys.append(_tag_key(key) if val is None else _tag_value_key(key, val)) + set_keys.append(tag_key(key) if val is None else tag_value_key(key, val)) if len(set_keys) == 1: hashes = await self._redis.smembers(set_keys[0]) else: @@ -622,15 +408,15 @@ async def _delete_commit_obj_by_hash(self, hash: str) -> bool: async def _delete_commit_obj(self, commit: Commit) -> bool: resolved_hash = commit.hash pipe = self._redis.pipeline() - _remove_commit_index_commands(pipe, commit, resolved_hash) + remove_commit_index_commands(pipe, commit, resolved_hash) await pipe.execute() - for h in _blob_hashes(commit): + for h in blob_hashes(commit): deleted = await self._redis.eval( - _DECR_DELETE_SCRIPT, + DECR_DELETE_SCRIPT, 3, - _blob_ref_key(h), - _blob_key(h), - _blob_stats_key(), + blob_ref_key(h), + blob_key(h), + blob_stats_key(), ) if deleted: logger.debug("orphan blob cleaned hash=%s", h[:12]) @@ -639,7 +425,7 @@ async def _delete_commit_obj(self, commit: Commit) -> bool: async def _touch_commit(self, hash: str) -> None: now = datetime.now(UTC).timestamp() - await self._redis.zadd(_access_key(), {hash: now}) + await self._redis.zadd(access_key(), {hash: now}) async def close(self) -> None: logger.debug("closing async redis store") diff --git a/src/cashet/serializers.py b/src/cashet/serializers.py new file mode 100644 index 0000000..ccce013 --- /dev/null +++ b/src/cashet/serializers.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from typing import Any, Protocol, runtime_checkable + +_pickle_warning_issued = False + + +def warn_default_pickle() -> None: + global _pickle_warning_issued + if _pickle_warning_issued: + return + _pickle_warning_issued = True + import warnings + + warnings.warn( + "Using PickleSerializer by default — arbitrary code execution risk on " + "untrusted cached results. Pass serializer=SafePickleSerializer() " + "for safer deserialization.", + stacklevel=3, + ) + + +@runtime_checkable +class Serializer(Protocol): + def dumps(self, obj: Any) -> bytes: ... + def loads(self, data: bytes) -> Any: ... + + +class PickleSerializer: + def __init__(self, protocol: int | None = None) -> None: + import pickle + + self._pickle = pickle + self._protocol = protocol or pickle.HIGHEST_PROTOCOL + + def dumps(self, obj: Any) -> bytes: + return self._pickle.dumps(obj, protocol=self._protocol) + + def loads(self, data: bytes) -> Any: + return self._pickle.loads(data) + + +class JsonSerializer: + def dumps(self, obj: Any) -> bytes: + import json + + return json.dumps(obj, default=str, sort_keys=True).encode() + + def loads(self, data: bytes) -> Any: + import json + + return json.loads(data) + + +class SafePickleSerializer: + _cached_allowlist: list[type] | None = None + + def __init__(self, extra_classes: list[type] | None = None) -> None: + import pickle + + self._pickle = pickle + self._allowed: dict[str, type] = {} + for cls in self._default_allowlist(): + key = f"{cls.__module__}.{cls.__qualname__}" + self._allowed[key] = cls + if extra_classes: + for cls in extra_classes: + key = f"{cls.__module__}.{cls.__qualname__}" + self._allowed[key] = cls + + def dumps(self, obj: Any) -> bytes: + return self._pickle.dumps(obj, protocol=self._pickle.HIGHEST_PROTOCOL) + + def loads(self, data: bytes) -> Any: + import io + import pickle + + allowed = self._allowed + blocked_msg = " — not in allowlist. Pass it via SafePickleSerializer(extra_classes=[...])." + + class _RestrictedUnpickler(pickle.Unpickler): + def find_class(self, module: str, name: str) -> Any: # type: ignore[override] + key = f"{module}.{name}" + if key in allowed: + return allowed[key] + raise pickle.UnpicklingError(f"Blocked class {key}{blocked_msg}") + + return _RestrictedUnpickler(io.BytesIO(data)).load() + + @classmethod + def _default_allowlist(cls) -> list[type]: + if cls._cached_allowlist is not None: + return cls._cached_allowlist + import collections + import datetime + + types_list: list[type] = [ + type(None), + bool, + int, + float, + str, + bytes, + bytearray, + list, + dict, + tuple, + set, + frozenset, + slice, + range, + complex, + object, + type, + datetime.datetime, + datetime.date, + datetime.timedelta, + datetime.time, + datetime.timezone, + collections.OrderedDict, + collections.defaultdict, + collections.Counter, + collections.deque, + ] + try: + import numpy # pyright: ignore[reportMissingImports] + + types_list.append(numpy.ndarray) # type: ignore[attr-defined] + except ImportError: + pass + cls._cached_allowlist = types_list + return types_list diff --git a/src/cashet/server.py b/src/cashet/server.py index dec1f9c..8b5fae1 100644 --- a/src/cashet/server.py +++ b/src/cashet/server.py @@ -9,6 +9,7 @@ import time import types from collections.abc import Callable, Mapping +from datetime import timedelta from typing import Any from starlette.applications import Starlette @@ -19,6 +20,7 @@ from cashet.async_client import AsyncClient from cashet.client import Client +from cashet.models import Commit, TaskError logger = logging.getLogger("cashet") @@ -54,6 +56,16 @@ def _too_large_response(max_size: int) -> _CustomJSONResponse: ) +def _task_failed_response(exc: TaskError) -> _CustomJSONResponse: + # A failing task is the caller's error, not a server bug: surface the final + # traceback line ("ExceptionType: message") but never server file paths. + lines = str(exc).strip().splitlines() + detail = lines[-1] if lines else "task failed" + return _CustomJSONResponse( + {"error": "task failed", "detail": detail}, status_code=422 + ) + + class _BadRequestError(Exception): pass @@ -222,8 +234,121 @@ async def wrapper(request: Request) -> JSONResponse: return wrapper -async def _async_submit(request: Request) -> JSONResponse: - client: AsyncClient = request.app.state.client +def _submit_options(data: dict[str, Any]) -> dict[str, Any]: + # Keys match the underscore-prefixed keyword parameters of Client.submit + # and AsyncClient.submit, so the ops adapters can splat them straight in. + return { + "_cache": data.get("cache", True), + "_tags": data.get("tags", {}), + "_retries": data.get("retries", 0), + "_force": data.get("force", False), + "_timeout": data.get("timeout"), + "_ttl": data.get("ttl"), + } + + +class _AsyncOps: + def __init__(self, client: AsyncClient) -> None: + self.client = client + self.serializer: Any = client.serializer + + async def submit_and_load( + self, + func: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + options: dict[str, Any], + ) -> tuple[str, str, bytes]: + ref = await self.client.submit(func, *args, **options, **kwargs) + result = await ref.load() + return ref.commit_hash, ref.hash, self.serializer.dumps(result) + + async def get_result(self, commit_hash: str) -> bytes: + return self.serializer.dumps(await self.client.get(commit_hash)) + + async def show(self, commit_hash: str) -> Commit | None: + return await self.client.show(commit_hash) + + async def log( + self, func_name: str | None, limit: int, status: str | None + ) -> list[Commit]: + return await self.client.log(func_name=func_name, limit=limit, status=status) + + async def stats(self) -> dict[str, int]: + return await self.client.stats() + + async def gc(self, older_than: timedelta, max_size: int | None) -> int: + return await self.client.gc(older_than, max_size_bytes=max_size) + + +class _SyncOps: + # Sync Client calls run in threads so they never block the event loop. + def __init__(self, client: Client) -> None: + self.client = client + self.serializer: Any = client.serializer + + async def submit_and_load( + self, + func: Callable[..., Any], + args: tuple[Any, ...], + kwargs: dict[str, Any], + options: dict[str, Any], + ) -> tuple[str, str, bytes]: + def run() -> tuple[str, str, bytes]: + ref = self.client.submit(func, *args, **options, **kwargs) + result = ref.load() + return ref.commit_hash, ref.hash, self.serializer.dumps(result) + + return await asyncio.to_thread(run) + + async def get_result(self, commit_hash: str) -> bytes: + return await asyncio.to_thread( + lambda: self.serializer.dumps(self.client.get(commit_hash)) + ) + + async def show(self, commit_hash: str) -> Commit | None: + return await asyncio.to_thread(self.client.show, commit_hash) + + async def log( + self, func_name: str | None, limit: int, status: str | None + ) -> list[Commit]: + return await asyncio.to_thread( + lambda: self.client.log(func_name=func_name, limit=limit, status=status) + ) + + async def stats(self) -> dict[str, int]: + return await asyncio.to_thread(self.client.stats) + + async def gc(self, older_than: timedelta, max_size: int | None) -> int: + return await asyncio.to_thread( + lambda: self.client.gc(older_than, max_size_bytes=max_size) + ) + + +_ServerOps = _AsyncOps | _SyncOps + + +def _log_request(request: Request, status_code: int, start: float) -> None: + logger.info( + "request method=%s path=%s status=%d duration_ms=%d", + request.method, + request.url.path, + status_code, + int((time.perf_counter() - start) * 1000), + ) + + +def _log_request_failed(request: Request, start: float) -> None: + logger.exception( + "request failed method=%s path=%s status=500 duration_ms=%d", + request.method, + request.url.path, + int((time.perf_counter() - start) * 1000), + ) + + +async def _submit(request: Request) -> JSONResponse: + ops: _ServerOps = request.app.state.ops data = await request.json() func, error = _resolve_func( data, @@ -235,195 +360,89 @@ async def _async_submit(request: Request) -> JSONResponse: if func is None: return _CustomJSONResponse({"error": "task required"}, status_code=400) - serializer = client.serializer args, kwargs, error = _decode_call( - data, serializer, request.app.state.allow_remote_code + data, ops.serializer, request.app.state.allow_remote_code ) if error is not None: return error - cache = data.get("cache", True) - tags = data.get("tags", {}) - retries = data.get("retries", 0) - force = data.get("force", False) - timeout = data.get("timeout") - ttl = data.get("ttl") - + options = _submit_options(data) start = time.perf_counter() try: - ref = await client.submit( - func, - *args, - _cache=cache, - _tags=tags, - _retries=retries, - _force=force, - _timeout=timeout, - _ttl=ttl, - **kwargs, - ) - result = await ref.load() - result_b64 = base64.b64encode(serializer.dumps(result)).decode() - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, - request.url.path, - 200, - duration, - ) - return _CustomJSONResponse( - { - "commit_hash": ref.commit_hash, - "blob_hash": ref.hash, - "result_b64": result_b64, - } + commit_hash, blob_hash, payload = await ops.submit_and_load( + func, args, kwargs, options ) + except TaskError as exc: + _log_request(request, 422, start) + return _task_failed_response(exc) except Exception: - duration = int((time.perf_counter() - start) * 1000) - logger.exception( - "request failed method=%s path=%s status=500 duration_ms=%d", - request.method, - request.url.path, - duration, - ) + _log_request_failed(request, start) return _CustomJSONResponse({"error": "Internal server error"}, status_code=500) + _log_request(request, 200, start) + return _CustomJSONResponse( + { + "commit_hash": commit_hash, + "blob_hash": blob_hash, + "result_b64": base64.b64encode(payload).decode(), + } + ) -async def _async_result(request: Request) -> JSONResponse: - client: AsyncClient = request.app.state.client +async def _result(request: Request) -> JSONResponse: + ops: _ServerOps = request.app.state.ops commit_hash = request.path_params["commit_hash"] start = time.perf_counter() try: - result = await client.get(commit_hash) - serializer = client.serializer - result_b64 = base64.b64encode(serializer.dumps(result)).decode() - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, 200, duration, - ) - return _CustomJSONResponse({"result_b64": result_b64}) + payload = await ops.get_result(commit_hash) except (KeyError, ValueError): - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, 404, duration, - ) + _log_request(request, 404, start) return _CustomJSONResponse({"error": "not found"}, status_code=404) except Exception: - duration = int((time.perf_counter() - start) * 1000) - logger.exception( - "request failed method=%s path=%s status=500 duration_ms=%d", - request.method, request.url.path, duration, - ) + _log_request_failed(request, start) return _CustomJSONResponse({"error": "Internal server error"}, status_code=500) + _log_request(request, 200, start) + return _CustomJSONResponse({"result_b64": base64.b64encode(payload).decode()}) -async def _async_commit(request: Request) -> JSONResponse: - client: AsyncClient = request.app.state.client +async def _commit(request: Request) -> JSONResponse: + ops: _ServerOps = request.app.state.ops commit_hash = request.path_params["commit_hash"] start = time.perf_counter() - c = await client.show(commit_hash) - status_code = 200 if c is not None else 404 - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, status_code, duration, - ) + c = await ops.show(commit_hash) + _log_request(request, 200 if c is not None else 404, start) if c is None: return _CustomJSONResponse({"error": "not found"}, status_code=404) return _CustomJSONResponse(c.summary()) -async def _async_log(request: Request) -> JSONResponse: - client: AsyncClient = request.app.state.client +async def _log(request: Request) -> JSONResponse: + ops: _ServerOps = request.app.state.ops func_name = request.query_params.get("func") limit = _query_int(request, "limit", 50) status = request.query_params.get("status") start = time.perf_counter() - commits = await client.log(func_name=func_name, limit=limit, status=status) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, 200, duration, - ) + commits = await ops.log(func_name, limit, status) + _log_request(request, 200, start) return _CustomJSONResponse([c.summary() for c in commits]) -async def _async_stats(request: Request) -> JSONResponse: - client: AsyncClient = request.app.state.client +async def _stats(request: Request) -> JSONResponse: + ops: _ServerOps = request.app.state.ops start = time.perf_counter() - result = await client.stats() - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, 200, duration, - ) + result = await ops.stats() + _log_request(request, 200, start) return _CustomJSONResponse(result) -async def _async_gc(request: Request) -> JSONResponse: - from datetime import timedelta - - client: AsyncClient = request.app.state.client +async def _gc(request: Request) -> JSONResponse: + ops: _ServerOps = request.app.state.ops older_than_days, max_size = _gc_params(await _json_body(request)) start = time.perf_counter() - deleted = await client.gc(timedelta(days=older_than_days), max_size_bytes=max_size) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, 200, duration, - ) + deleted = await ops.gc(timedelta(days=older_than_days), max_size) + _log_request(request, 200, start) return _CustomJSONResponse({"deleted": deleted}) -def create_async_app( - client: AsyncClient, - require_token: str | None = None, - *, - tasks: TaskRegistry | None = None, - allow_remote_code: bool = False, - max_content_length: int = _DEFAULT_MAX_CONTENT_LENGTH, -) -> Starlette: - _validate_remote_code_options(allow_remote_code, require_token) - routes = [ - Route( - "/submit", - _require_token(_safe_handler(_async_submit), require_token), - methods=["POST"], - ), - Route( - "/result/{commit_hash}", - _require_token(_safe_handler(_async_result), require_token), - methods=["GET"], - ), - Route( - "/commit/{commit_hash}", - _require_token(_safe_handler(_async_commit), require_token), - methods=["GET"], - ), - Route("/log", _require_token(_safe_handler(_async_log), require_token), methods=["GET"]), - Route( - "/stats", - _require_token(_safe_handler(_async_stats), require_token), - methods=["GET"], - ), - Route("/gc", _require_token(_safe_handler(_async_gc), require_token), methods=["POST"]), - ] - - app = Starlette( - routes=routes, - middleware=[Middleware(_SizeLimitMiddleware)], - ) - app.state.client = client - app.state.tasks = _server_tasks(client, tasks) - app.state.allow_remote_code = allow_remote_code - app.state.max_content_length = max_content_length - app.state.require_token = require_token - return app - - class _SizeLimitMiddleware: # Pure-ASGI so the body cap is enforced on bytes actually received, not on a # client-supplied Content-Length that is absent for chunked requests. @@ -492,185 +511,12 @@ async def replay_receive() -> Any: await self.app(scope, replay_receive, send) -# Sync handlers run sync Client operations in threads so they don't block the event loop - - -async def _submit(request: Request) -> JSONResponse: - client: Client = request.app.state.client - data = await request.json() - func, error = _resolve_func( - data, - request.app.state.tasks, - request.app.state.allow_remote_code, - ) - if error is not None: - return error - if func is None: - return _CustomJSONResponse({"error": "task required"}, status_code=400) - - serializer = client.serializer - args, kwargs, error = _decode_call( - data, serializer, request.app.state.allow_remote_code - ) - if error is not None: - return error - - cache = data.get("cache", True) - tags = data.get("tags", {}) - retries = data.get("retries", 0) - force = data.get("force", False) - timeout = data.get("timeout") - ttl = data.get("ttl") - - start = time.perf_counter() - - def _run() -> JSONResponse: - try: - ref = client.submit( - func, - *args, - _cache=cache, - _tags=tags, - _retries=retries, - _force=force, - _timeout=timeout, - _ttl=ttl, - **kwargs, - ) - result = ref.load() - result_b64 = base64.b64encode(serializer.dumps(result)).decode() - return _CustomJSONResponse( - { - "commit_hash": ref.commit_hash, - "blob_hash": ref.hash, - "result_b64": result_b64, - } - ) - except Exception: - logger.exception( - "request failed method=POST path=/submit status=500", - ) - return _CustomJSONResponse({"error": "Internal server error"}, status_code=500) - - response = await asyncio.to_thread(_run) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, response.status_code, duration, - ) - return response - - -async def _result(request: Request) -> JSONResponse: - client: Client = request.app.state.client - commit_hash = request.path_params["commit_hash"] - start = time.perf_counter() - - def _run() -> JSONResponse: - try: - result = client.get(commit_hash) - serializer = client.serializer - result_b64 = base64.b64encode(serializer.dumps(result)).decode() - return _CustomJSONResponse({"result_b64": result_b64}) - except (KeyError, ValueError): - return _CustomJSONResponse({"error": "not found"}, status_code=404) - except Exception: - logger.exception( - "request failed method=%s path=%s status=500", - request.method, request.url.path, - ) - return _CustomJSONResponse({"error": "Internal server error"}, status_code=500) - - response = await asyncio.to_thread(_run) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, response.status_code, duration, - ) - return response - - -async def _commit(request: Request) -> JSONResponse: - client: Client = request.app.state.client - commit_hash = request.path_params["commit_hash"] - start = time.perf_counter() - - def _run() -> JSONResponse: - c = client.show(commit_hash) - if c is None: - return _CustomJSONResponse({"error": "not found"}, status_code=404) - return _CustomJSONResponse(c.summary()) - - response = await asyncio.to_thread(_run) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, response.status_code, duration, - ) - return response - - -async def _log(request: Request) -> JSONResponse: - client: Client = request.app.state.client - func_name = request.query_params.get("func") - limit = _query_int(request, "limit", 50) - status = request.query_params.get("status") - start = time.perf_counter() - - def _run() -> JSONResponse: - commits = client.log(func_name=func_name, limit=limit, status=status) - return _CustomJSONResponse([c.summary() for c in commits]) - - response = await asyncio.to_thread(_run) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, response.status_code, duration, - ) - return response - - -async def _stats(request: Request) -> JSONResponse: - client: Client = request.app.state.client - start = time.perf_counter() - response = await asyncio.to_thread( - lambda: _CustomJSONResponse(client.stats()) - ) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, response.status_code, duration, - ) - return response - - -async def _gc(request: Request) -> JSONResponse: - from datetime import timedelta - - client: Client = request.app.state.client - older_than_days, max_size = _gc_params(await _json_body(request)) - start = time.perf_counter() - - def _run() -> JSONResponse: - deleted = client.gc(timedelta(days=older_than_days), max_size_bytes=max_size) - return _CustomJSONResponse({"deleted": deleted}) - - response = await asyncio.to_thread(_run) - duration = int((time.perf_counter() - start) * 1000) - logger.info( - "request method=%s path=%s status=%d duration_ms=%d", - request.method, request.url.path, response.status_code, duration, - ) - return response - - -def create_app( - client: Client, - require_token: str | None = None, - *, - tasks: TaskRegistry | None = None, - allow_remote_code: bool = False, - max_content_length: int = _DEFAULT_MAX_CONTENT_LENGTH, +def _build_app( + ops: _ServerOps, + require_token: str | None, + tasks: TaskRegistry | None, + allow_remote_code: bool, + max_content_length: int, ) -> Starlette: _validate_remote_code_options(allow_remote_code, require_token) routes = [ @@ -701,9 +547,36 @@ def create_app( routes=routes, middleware=[Middleware(_SizeLimitMiddleware)], ) - app.state.client = client - app.state.tasks = _server_tasks(client, tasks) + app.state.client = ops.client + app.state.ops = ops + app.state.tasks = _server_tasks(ops.client, tasks) app.state.allow_remote_code = allow_remote_code app.state.max_content_length = max_content_length app.state.require_token = require_token return app + + +def create_app( + client: Client, + require_token: str | None = None, + *, + tasks: TaskRegistry | None = None, + allow_remote_code: bool = False, + max_content_length: int = _DEFAULT_MAX_CONTENT_LENGTH, +) -> Starlette: + return _build_app( + _SyncOps(client), require_token, tasks, allow_remote_code, max_content_length + ) + + +def create_async_app( + client: AsyncClient, + require_token: str | None = None, + *, + tasks: TaskRegistry | None = None, + allow_remote_code: bool = False, + max_content_length: int = _DEFAULT_MAX_CONTENT_LENGTH, +) -> Starlette: + return _build_app( + _AsyncOps(client), require_token, tasks, allow_remote_code, max_content_length + ) diff --git a/src/cashet/store.py b/src/cashet/store.py index 34f63c6..19c0999 100644 --- a/src/cashet/store.py +++ b/src/cashet/store.py @@ -2,871 +2,28 @@ import asyncio import hashlib -import json -import logging import sqlite3 -import threading -import zlib -from dataclasses import dataclass -from datetime import UTC, datetime, timedelta +from datetime import datetime from pathlib import Path from typing import Any +from cashet._locks import SQLiteFingerprintLock from cashet._runner import BlockingAsyncRunner -from cashet.models import Commit, ObjectRef, StorageTier, TaskDef, TaskStatus - -_BLOB_COMPRESS_THRESHOLD = 256 -_INLINE_THRESHOLD = 1024 - -logger = logging.getLogger("cashet") - - -def _normalize_hash_prefix(hash: str) -> str | None: - if 0 < len(hash) <= 64 and all(c in "0123456789abcdefABCDEF" for c in hash): - return hash.lower() - return None - - -@dataclass -class _SQLiteLockState: - thread_lock: threading.Lock - file_lock: Any - - -_SQLITE_LOCKS: dict[str, _SQLiteLockState] = {} -_SQLITE_LOCKS_GUARD = threading.Lock() - - -def _sqlite_lock_state(lock_path: str) -> _SQLiteLockState: - with _SQLITE_LOCKS_GUARD: - state = _SQLITE_LOCKS.get(lock_path) - if state is None: - from filelock import FileLock - - state = _SQLiteLockState( - thread_lock=threading.Lock(), - file_lock=FileLock(lock_path, timeout=30, thread_local=False), - ) - _SQLITE_LOCKS[lock_path] = state - return state - - -class _SQLiteStoreCore: - def __init__(self, root: Path) -> None: - self.root = root - self.objects_dir = root / "objects" - self.db_path = root / "meta.db" - self.objects_dir.mkdir(parents=True, exist_ok=True) - self._tls = threading.local() - self._lock = threading.Lock() - self._init_db() - - def _connect(self, *, immediate: bool = False) -> sqlite3.Connection: - conn: sqlite3.Connection | None = getattr(self._tls, "conn", None) - try: - if conn is not None: - conn.execute("SELECT 1") - else: - conn = None - except sqlite3.Error: - conn = None - if conn is None: - conn = sqlite3.connect(str(self.db_path), isolation_level=None) - conn.row_factory = sqlite3.Row - conn.execute("PRAGMA auto_vacuum=INCREMENTAL") - conn.execute("PRAGMA journal_mode=WAL") - conn.execute("PRAGMA busy_timeout=5000") - self._tls.conn = conn - if immediate: - conn.execute("BEGIN IMMEDIATE") - return conn - - def close(self) -> None: - logger.debug("closing sqlite store path=%s", str(self.db_path)) - conn: sqlite3.Connection | None = getattr(self._tls, "conn", None) - if conn is not None: - conn.close() - self._tls.conn = None - - def _init_db(self) -> None: - logger.debug("initializing sqlite store path=%s", str(self.db_path)) - conn = self._connect() - conn.execute(""" - CREATE TABLE IF NOT EXISTS commits ( - hash TEXT PRIMARY KEY, - fingerprint TEXT NOT NULL, - func_name TEXT NOT NULL, - func_hash TEXT NOT NULL, - args_hash TEXT NOT NULL, - args_snapshot BLOB, - func_source TEXT, - dep_versions TEXT, - cache INTEGER NOT NULL DEFAULT 1, - retries INTEGER NOT NULL DEFAULT 0, - force INTEGER NOT NULL DEFAULT 0, - timeout_seconds REAL, - ttl_seconds REAL, - input_refs TEXT, - output_hash TEXT, - output_size INTEGER, - output_tier TEXT, - parent_hash TEXT, - status TEXT NOT NULL DEFAULT 'pending', - error TEXT, - tags TEXT, - created_at TEXT NOT NULL, - claimed_at TEXT NOT NULL, - last_accessed_at TEXT, - expires_at TEXT, - FOREIGN KEY (parent_hash) REFERENCES commits(hash) - ) - """) - conn.execute("CREATE INDEX IF NOT EXISTS idx_fingerprint ON commits(fingerprint)") - conn.execute("CREATE INDEX IF NOT EXISTS idx_func_name ON commits(func_name)") - conn.execute("CREATE INDEX IF NOT EXISTS idx_created_at ON commits(created_at)") - self._migrate_last_accessed_at(conn) - self._migrate_retries(conn) - self._migrate_task_options(conn) - self._migrate_expires_at(conn) - self._migrate_ttl(conn) - self._migrate_claimed_at(conn) - conn.execute(""" - CREATE TABLE IF NOT EXISTS inline_objects ( - hash TEXT PRIMARY KEY, - data BLOB NOT NULL - ) - """) - - def _migrate_last_accessed_at(self, conn: sqlite3.Connection) -> None: - col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] - if "last_accessed_at" not in col_names: - conn.execute( - "ALTER TABLE commits ADD COLUMN last_accessed_at TEXT DEFAULT NULL" - ) - conn.execute( - "UPDATE commits SET last_accessed_at = created_at WHERE last_accessed_at IS NULL" - ) - conn.execute( - "CREATE INDEX IF NOT EXISTS idx_last_accessed_at ON commits(last_accessed_at)" - ) - - def _migrate_retries(self, conn: sqlite3.Connection) -> None: - col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] - if "retries" not in col_names: - conn.execute( - "ALTER TABLE commits ADD COLUMN retries INTEGER NOT NULL DEFAULT 0" - ) - - def _migrate_task_options(self, conn: sqlite3.Connection) -> None: - col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] - if "force" not in col_names: - conn.execute( - "ALTER TABLE commits ADD COLUMN force INTEGER NOT NULL DEFAULT 0" - ) - if "timeout_seconds" not in col_names: - conn.execute("ALTER TABLE commits ADD COLUMN timeout_seconds REAL") - - def _migrate_expires_at(self, conn: sqlite3.Connection) -> None: - col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] - if "expires_at" not in col_names: - conn.execute("ALTER TABLE commits ADD COLUMN expires_at TEXT") - - def _migrate_ttl(self, conn: sqlite3.Connection) -> None: - col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] - if "ttl_seconds" not in col_names: - conn.execute("ALTER TABLE commits ADD COLUMN ttl_seconds REAL") - - def _migrate_claimed_at(self, conn: sqlite3.Connection) -> None: - col_names = [r[1] for r in conn.execute("PRAGMA table_info(commits)").fetchall()] - if "claimed_at" not in col_names: - conn.execute( - "ALTER TABLE commits ADD COLUMN claimed_at TEXT NOT NULL DEFAULT ''" - ) - conn.execute( - "UPDATE commits SET claimed_at = created_at WHERE claimed_at = ''" - ) - - def put_blob(self, data: bytes) -> ObjectRef: - content_hash = hashlib.sha256(data).hexdigest() - if len(data) < _INLINE_THRESHOLD: - conn = self._connect() - conn.execute( - "INSERT OR IGNORE INTO inline_objects (hash, data) VALUES (?, ?)", - (content_hash, data), - ) - logger.info( - "blob stored hash=%s size=%d tier=inline", - content_hash[:12], - len(data), - ) - return ObjectRef(hash=content_hash, size=len(data), tier=StorageTier.INLINE) - prefix = content_hash[:2] - suffix = content_hash[2:] - obj_path = self.objects_dir / prefix / suffix - if obj_path.exists(): - logger.debug( - "blob deduplicated hash=%s size=%d", - content_hash[:12], - len(data), - ) - return ObjectRef(hash=content_hash, size=len(data), tier=StorageTier.BLOB) - obj_path.parent.mkdir(parents=True, exist_ok=True) - stored = data - if len(data) >= _BLOB_COMPRESS_THRESHOLD: - compressed = zlib.compress(data, level=6) - if len(compressed) < len(data): - stored = compressed - obj_path.write_bytes(stored) - logger.info( - "blob stored hash=%s size=%d tier=blob compressed=%s", - content_hash[:12], - len(data), - "true" if stored is not data else "false", - ) - return ObjectRef(hash=content_hash, size=len(data), tier=StorageTier.BLOB) - - def get_blob(self, ref: ObjectRef) -> bytes: - if ref.tier == StorageTier.INLINE: - conn = self._connect() - row = conn.execute( - "SELECT data FROM inline_objects WHERE hash = ?", (ref.hash,) - ).fetchone() - if row is None: - logger.warning("inline blob not found hash=%s", ref.hash[:12]) - raise ValueError(f"Inline blob {ref.hash} not found") - data = row["data"] - if hashlib.sha256(data).hexdigest() != ref.hash: - logger.error( - "inline blob integrity check failed hash=%s", - ref.hash[:12], - ) - raise ValueError(f"Inline blob {ref.hash} integrity check failed") - return data - prefix = ref.hash[:2] - suffix = ref.hash[2:] - obj_path = self.objects_dir / prefix / suffix - raw = obj_path.read_bytes() - try: - decompressed = zlib.decompress(raw) - if hashlib.sha256(decompressed).hexdigest() == ref.hash: - return decompressed - except zlib.error: - pass - if hashlib.sha256(raw).hexdigest() != ref.hash: - logger.error( - "blob integrity check failed hash=%s", - ref.hash[:12], - ) - raise ValueError(f"Blob {ref.hash} integrity check failed") - return raw - - def blob_exists(self, hash: str) -> bool: - if (self.objects_dir / hash[:2] / hash[2:]).exists(): - return True - conn = self._connect() - row = conn.execute( - "SELECT 1 FROM inline_objects WHERE hash = ?", (hash,) - ).fetchone() - return row is not None - - def put_commit(self, commit: Commit) -> None: - conn = self._connect(immediate=True) - now = datetime.now(UTC).isoformat() - try: - self._put_commit_row(conn, commit, now) - conn.execute("COMMIT") - except Exception: - conn.execute("ROLLBACK") - raise - - def _put_commit_row( - self, conn: sqlite3.Connection, commit: Commit, accessed_at: str - ) -> None: - output_hash = commit.output_ref.hash if commit.output_ref else None - output_size = commit.output_ref.size if commit.output_ref else None - output_tier = commit.output_ref.tier.value if commit.output_ref else None - timeout_seconds = ( - commit.task_def.timeout.total_seconds() if commit.task_def.timeout else None - ) - conn.execute( - """INSERT OR REPLACE INTO commits - (hash, fingerprint, func_name, func_hash, args_hash, args_snapshot, - func_source, dep_versions, cache, retries, force, timeout_seconds, - ttl_seconds, input_refs, output_hash, output_size, output_tier, parent_hash, - status, error, tags, created_at, - claimed_at, last_accessed_at, expires_at) - VALUES ( - ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? - )""", - ( - commit.hash, - commit.fingerprint, - commit.task_def.func_name, - commit.task_def.func_hash, - commit.task_def.args_hash, - commit.task_def.args_snapshot, - commit.task_def.func_source, - json.dumps(commit.task_def.dep_versions), - int(commit.task_def.cache), - commit.task_def.retries, - int(commit.task_def.force), - timeout_seconds, - commit.task_def.ttl.total_seconds() if commit.task_def.ttl else None, - json.dumps([r.hash for r in commit.input_refs]), - output_hash, - output_size, - output_tier, - commit.parent_hash, - commit.status.value, - commit.error, - json.dumps(commit.tags), - commit.created_at.isoformat(), - commit.claimed_at.isoformat(), - accessed_at, - commit.expires_at.isoformat() if commit.expires_at else None, - ), - ) - - def find_by_fingerprint(self, fingerprint: str) -> Commit | None: - conn = self._connect() - now = datetime.now(UTC) - now_iso = now.isoformat() - row = conn.execute( - """SELECT * FROM commits - WHERE fingerprint = ? AND status IN ('completed', 'cached') - AND (expires_at IS NULL OR expires_at > ?) - ORDER BY created_at DESC - LIMIT 1""", - (fingerprint, now_iso), - ).fetchone() - if row is None: - return None - try: - conn.execute( - "UPDATE commits SET last_accessed_at = ? WHERE hash = ?", - (now_iso, row["hash"]), - ) - except sqlite3.OperationalError: - # Access-time bump only feeds LRU ordering; never fail a cache hit - # because a concurrent writer holds the lock past busy_timeout. - logger.debug("last_accessed_at bump skipped (db locked) hash=%s", row["hash"][:12]) - return self._row_to_commit(row) - - def find_running_by_fingerprint(self, fingerprint: str) -> Commit | None: - conn = self._connect() - row = conn.execute( - """SELECT * FROM commits - WHERE fingerprint = ? AND status = 'running' - ORDER BY claimed_at DESC LIMIT 1""", - (fingerprint,), - ).fetchone() - if row is None: - return None - return self._row_to_commit(row) - - def get_commit(self, hash: str) -> Commit | None: - normalized = _normalize_hash_prefix(hash) - if normalized is None: - return None - hash = normalized - conn = self._connect() - if len(hash) < 64: - rows = conn.execute( - "SELECT * FROM commits WHERE hash LIKE ?", (hash + "%",) - ).fetchall() - if not rows: - return None - if len(rows) > 1: - matches = ", ".join(r["hash"][:12] for r in rows) - raise ValueError( - f"Ambiguous prefix {hash[:12]} matches {len(rows)} commits: {matches}" - ) - return self._row_to_commit(rows[0]) - row = conn.execute("SELECT * FROM commits WHERE hash = ?", (hash,)).fetchone() - if row is None: - return None - return self._row_to_commit(row) - - def list_commits( - self, - func_name: str | None = None, - limit: int = 50, - status: TaskStatus | None = None, - tags: dict[str, str | None] | None = None, - ) -> list[Commit]: - conn = self._connect() - query = "SELECT * FROM commits WHERE 1=1" - params: list[Any] = [] - if func_name: - query += " AND func_name = ?" - params.append(func_name) - if status: - query += " AND status = ?" - params.append(status.value) - if tags: - for key, val in tags.items(): - if val is None: - query += " AND json_extract(tags, ?) IS NOT NULL" - params.append(f"$.{key}") - else: - query += " AND json_extract(tags, ?) = ?" - params.append(f"$.{key}") - params.append(val) - query += " ORDER BY created_at DESC LIMIT ?" - params.append(limit) - rows = conn.execute(query, params).fetchall() - return [self._row_to_commit(r) for r in rows] - - def get_history(self, hash: str) -> list[Commit]: - normalized = _normalize_hash_prefix(hash) - if normalized is None: - return [] - hash = normalized - conn = self._connect() - if len(hash) < 64: - rows = conn.execute( - "SELECT * FROM commits WHERE hash LIKE ?", (hash + "%",) - ).fetchall() - if not rows: - return [] - if len(rows) > 1: - matches = ", ".join(r["hash"][:12] for r in rows) - raise ValueError( - f"Ambiguous prefix {hash[:12]} matches {len(rows)} commits: {matches}" - ) - commit = self._row_to_commit(rows[0]) - else: - row = conn.execute("SELECT * FROM commits WHERE hash = ?", (hash,)).fetchone() - if row is None: - return [] - commit = self._row_to_commit(row) - fingerprint = commit.fingerprint - rows = conn.execute( - """SELECT * FROM commits - WHERE fingerprint = ? AND status IN ('completed', 'cached') - ORDER BY created_at ASC""", - (fingerprint,), - ).fetchall() - now = datetime.now(UTC) - result: list[Commit] = [] - for r in rows: - c = self._row_to_commit(r) - if c.expires_at is not None and c.expires_at <= now: - continue - result.append(c) - return result - - def stats(self) -> dict[str, int]: - conn = self._connect() - total = conn.execute("SELECT COUNT(*) FROM commits").fetchone()[0] - completed = conn.execute( - "SELECT COUNT(*) FROM commits WHERE status IN ('completed', 'cached')" - ).fetchone()[0] - obj_count, total_bytes = self._blob_storage_totals() - inline_count, inline_bytes = self._inline_storage_totals(conn) - return { - "total_commits": total, - "completed_commits": completed, - "stored_objects": obj_count + inline_count, - "disk_bytes": total_bytes + inline_bytes, - "blob_objects": obj_count, - "blob_bytes": total_bytes, - "inline_objects": inline_count, - "inline_bytes": inline_bytes, - } - - def _blob_storage_totals(self) -> tuple[int, int]: - obj_count = 0 - total_bytes = 0 - for p in self.objects_dir.iterdir(): - if p.is_dir(): - for f in p.iterdir(): - if f.is_file(): - obj_count += 1 - total_bytes += f.stat().st_size - return obj_count, total_bytes - - def _inline_storage_totals(self, conn: sqlite3.Connection) -> tuple[int, int]: - inline_row = conn.execute( - "SELECT COUNT(*) AS cnt, COALESCE(SUM(LENGTH(data)), 0) AS bytes FROM inline_objects" - ).fetchone() - inline_count = inline_row["cnt"] if inline_row else 0 - inline_bytes = inline_row["bytes"] if inline_row else 0 - return inline_count, inline_bytes - - def _storage_bytes(self, conn: sqlite3.Connection) -> int: - return self._blob_storage_totals()[1] + self._inline_storage_totals(conn)[1] - - def evict( - self, older_than: datetime, max_size_bytes: int | None = None - ) -> int: - conn = self._connect(immediate=True) - orphans: list[str] = [] - try: - cutoff = older_than.isoformat() - candidates: set[str] = set() - evicted_hashes = [ - row[0] - for row in conn.execute( - "SELECT hash FROM commits WHERE last_accessed_at < ?", (cutoff,) - ) - ] - for row in conn.execute( - "SELECT output_hash FROM commits " - "WHERE last_accessed_at < ? AND output_hash IS NOT NULL", - (cutoff,), - ): - candidates.add(row[0]) - for row in conn.execute( - "SELECT input_refs FROM commits " - "WHERE last_accessed_at < ? AND input_refs IS NOT NULL", - (cutoff,), - ): - for h in json.loads(row[0]): - candidates.add(h) - if evicted_hashes: - placeholders = ", ".join("?" for _ in evicted_hashes) - conn.execute( - f"UPDATE commits SET parent_hash = NULL WHERE parent_hash IN ({placeholders})", - evicted_hashes, - ) - cursor = conn.execute( - "DELETE FROM commits WHERE last_accessed_at < ?", (cutoff,) - ) - deleted = cursor.rowcount - if candidates: - orphans = self._find_orphan_objects(conn, candidates) - conn.execute("COMMIT") - except Exception: - conn.execute("ROLLBACK") - raise - if orphans: - logger.info( - "orphan objects cleaned count=%d", - len(orphans), - ) - self._delete_orphan_objects(conn, orphans) - - if max_size_bytes is not None: - current_bytes = self._storage_bytes(conn) - if current_bytes > max_size_bytes: - deleted += self._evict_to_size(current_bytes, max_size_bytes) - - if deleted: - logger.info( - "eviction complete deleted=%d reason=%s", - deleted, - "size_limit" if max_size_bytes is not None else "ttl", - ) - try: - self._vacuum() - except sqlite3.OperationalError: - # Reclaiming space is best-effort; a concurrent writer holding the - # lock must not undo an eviction whose deletes already committed. - logger.debug("vacuum skipped (db busy) after eviction") - else: - logger.debug("eviction found no candidates") - - return deleted - - def _vacuum(self) -> None: - conn = self._connect() - mode = conn.execute("PRAGMA auto_vacuum").fetchone()[0] - if mode == 2: - conn.execute("PRAGMA incremental_vacuum") - else: - conn.execute("VACUUM") - - def _evict_to_size(self, current_bytes: int, max_size_bytes: int) -> int: - pending_orphans: list[str] = [] - deleted = 0 - conn = self._connect(immediate=True) - try: - ref_counts = self._object_ref_counts(conn) - rows = conn.execute( - "SELECT * FROM commits ORDER BY last_accessed_at ASC" - ).fetchall() - for target in rows: - if current_bytes <= max_size_bytes: - break - refs = self._row_object_refs(target) - conn.execute( - "UPDATE commits SET parent_hash = NULL WHERE parent_hash = ?", - (target["hash"],), - ) - conn.execute("DELETE FROM commits WHERE hash = ?", (target["hash"],)) - deleted += 1 - for obj_hash in refs: - count = ref_counts.get(obj_hash, 0) - if count <= 0: - continue - if count == 1: - ref_counts.pop(obj_hash, None) - if obj_hash not in pending_orphans: - current_bytes -= self._object_storage_size(conn, obj_hash) - pending_orphans.append(obj_hash) - else: - ref_counts[obj_hash] = count - 1 - conn.execute("COMMIT") - except Exception: - conn.execute("ROLLBACK") - raise - if pending_orphans: - logger.info( - "orphan objects cleaned count=%d", - len(pending_orphans), - ) - self._delete_orphan_objects(conn, pending_orphans) - return deleted - - def _delete_commit_body(self, conn: sqlite3.Connection, hash: str) -> tuple[bool, list[str]]: - if len(hash) < 64: - rows = conn.execute( - "SELECT * FROM commits WHERE hash LIKE ?", (hash + "%",) - ).fetchall() - if not rows: - return False, [] - if len(rows) > 1: - matches = ", ".join(r["hash"][:12] for r in rows) - raise ValueError( - f"Ambiguous prefix {hash[:12]} matches {len(rows)} commits: {matches}" - ) - target = rows[0] - else: - target = conn.execute( - "SELECT * FROM commits WHERE hash = ?", (hash,) - ).fetchone() - if target is None: - return False, [] - - candidates = set(self._row_object_refs(target)) - conn.execute( - "UPDATE commits SET parent_hash = NULL WHERE parent_hash = ?", - (target["hash"],), - ) - conn.execute("DELETE FROM commits WHERE hash = ?", (target["hash"],)) - - orphans: list[str] = [] - if candidates: - orphans = self._find_orphan_objects(conn, candidates) - return True, orphans - - def delete_commit(self, hash: str) -> bool: - normalized = _normalize_hash_prefix(hash) - if normalized is None: - return False - hash = normalized - conn = self._connect(immediate=True) - try: - success, orphans = self._delete_commit_body(conn, hash) - if not success: - # Avoid leaving a dangling transaction that poisons the next write. - conn.execute("ROLLBACK") - return False - conn.execute("COMMIT") - except Exception: - conn.execute("ROLLBACK") - raise - if orphans: - logger.info( - "orphan objects cleaned count=%d", - len(orphans), - ) - self._delete_orphan_objects(conn, orphans) - logger.debug("commit deleted hash=%s", hash[:12]) - return True - - def delete_by_tags(self, tags: dict[str, str | None]) -> int: - conn = self._connect(immediate=True) - query = "SELECT hash, output_hash, input_refs FROM commits WHERE 1=1" - params: list[Any] = [] - for key, val in tags.items(): - if val is None: - query += " AND json_extract(tags, ?) IS NOT NULL" - params.append(f"$.{key}") - else: - query += " AND json_extract(tags, ?) = ?" - params.append(f"$.{key}") - params.append(val) - rows = conn.execute(query, params).fetchall() - if not rows: - conn.execute("ROLLBACK") - return 0 - hashes = [r[0] for r in rows] - candidates: set[str] = set() - for r in rows: - if r[1]: - candidates.add(r[1]) - if r[2]: - for h in json.loads(r[2]): - candidates.add(h) - try: - placeholders = ", ".join("?" for _ in hashes) - conn.execute( - f"UPDATE commits SET parent_hash = NULL WHERE parent_hash IN ({placeholders})", - hashes, - ) - conn.execute( - f"DELETE FROM commits WHERE hash IN ({placeholders})", hashes - ) - deleted = len(hashes) - all_orphans = self._find_orphan_objects(conn, candidates) if candidates else [] - conn.execute("COMMIT") - except Exception: - conn.execute("ROLLBACK") - raise - if all_orphans: - logger.info("orphan objects cleaned count=%d", len(all_orphans)) - self._delete_orphan_objects(conn, all_orphans) - return deleted - - def _row_object_refs(self, row: sqlite3.Row) -> list[str]: - refs: list[str] = [] - if row["output_hash"]: - refs.append(row["output_hash"]) - if row["input_refs"]: - refs.extend(set(json.loads(row["input_refs"]))) - return refs - - def _object_ref_counts(self, conn: sqlite3.Connection) -> dict[str, int]: - counts: dict[str, int] = {} - for row in conn.execute("SELECT output_hash FROM commits WHERE output_hash IS NOT NULL"): - counts[row[0]] = counts.get(row[0], 0) + 1 - for row in conn.execute("SELECT input_refs FROM commits WHERE input_refs IS NOT NULL"): - for h in set(json.loads(row[0])): - counts[h] = counts.get(h, 0) + 1 - return counts - - def _object_storage_size(self, conn: sqlite3.Connection, obj_hash: str) -> int: - size = 0 - blob_path = self.objects_dir / obj_hash[:2] / obj_hash[2:] - if blob_path.exists(): - size += blob_path.stat().st_size - row = conn.execute( - "SELECT LENGTH(data) AS size FROM inline_objects WHERE hash = ?", (obj_hash,) - ).fetchone() - if row is not None: - size += row["size"] or 0 - return size - - def _find_orphan_objects(self, conn: sqlite3.Connection, candidates: set[str]) -> list[str]: - still_output: set[str] = set() - for row in conn.execute("SELECT output_hash FROM commits WHERE output_hash IS NOT NULL"): - still_output.add(row[0]) - still_input: set[str] = set() - for row in conn.execute("SELECT input_refs FROM commits WHERE input_refs IS NOT NULL"): - still_input.update(json.loads(row[0])) - return [h for h in candidates if h not in still_output and h not in still_input] - - def _delete_orphan_objects(self, conn: sqlite3.Connection | None, orphans: list[str]) -> int: - if conn is None: - conn = self._connect() - freed = 0 - for obj_hash in orphans: - blob_path = self.objects_dir / obj_hash[:2] / obj_hash[2:] - if blob_path.exists(): - freed += blob_path.stat().st_size - blob_path.unlink() - row = conn.execute( - "SELECT LENGTH(data) AS size FROM inline_objects WHERE hash = ?", (obj_hash,) - ).fetchone() - if row is not None: - freed += row["size"] or 0 - conn.execute("DELETE FROM inline_objects WHERE hash = ?", (obj_hash,)) - for prefix_dir in list(self.objects_dir.iterdir()): - if prefix_dir.is_dir() and not any(prefix_dir.iterdir()): - prefix_dir.rmdir() - return freed - - def _row_to_commit(self, row: sqlite3.Row) -> Commit: - output_ref = None - if row["output_hash"]: - tier = StorageTier(row["output_tier"]) if row["output_tier"] else StorageTier.BLOB - output_ref = ObjectRef( - hash=row["output_hash"], - size=row["output_size"] or 0, - tier=tier, - ) - input_refs: list[ObjectRef] = [] - if row["input_refs"]: - for h in json.loads(row["input_refs"]): - input_refs.append(ObjectRef(hash=h)) - dep_versions: dict[str, str] = {} - if row["dep_versions"]: - dep_versions = json.loads(row["dep_versions"]) - tags: dict[str, str] = {} - if row["tags"]: - tags = json.loads(row["tags"]) - row_keys = set(row.keys()) - timeout_seconds = row["timeout_seconds"] if "timeout_seconds" in row_keys else None - ttl_seconds = row["ttl_seconds"] if "ttl_seconds" in row_keys else None - task_def = TaskDef( - func_hash=row["func_hash"], - func_name=row["func_name"], - func_source=row["func_source"] or "", - args_hash=row["args_hash"], - args_snapshot=row["args_snapshot"] or b"", - dep_versions=dep_versions, - cache=bool(row["cache"]), - tags=tags, - retries=row["retries"] if "retries" in row_keys else 0, - force=bool(row["force"]) if "force" in row_keys else False, - timeout=timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None, - ttl=timedelta(seconds=ttl_seconds) if ttl_seconds is not None else None, - ) - created_at = row["created_at"] - if isinstance(created_at, str): - created_at = datetime.fromisoformat(created_at) - claimed_at = row["claimed_at"] - if isinstance(claimed_at, str): - claimed_at = datetime.fromisoformat(claimed_at) - expires_at = row["expires_at"] - if isinstance(expires_at, str): - expires_at = datetime.fromisoformat(expires_at) - return Commit( - hash=row["hash"], - task_def=task_def, - input_refs=input_refs, - output_ref=output_ref, - parent_hash=row["parent_hash"], - status=TaskStatus(row["status"]), - created_at=created_at, - claimed_at=claimed_at, - error=row["error"], - tags=tags, - expires_at=expires_at, - ) - - -class _SQLiteFingerprintLock: - def __init__(self, lock_path: str) -> None: - self._state = _sqlite_lock_state(lock_path) - - async def __aenter__(self) -> None: - await asyncio.to_thread(self._state.thread_lock.acquire) - try: - await asyncio.to_thread(self._state.file_lock.acquire) - except Exception: - self._state.thread_lock.release() - raise - - async def __aexit__(self, *args: Any) -> None: - try: - await asyncio.to_thread(self._state.file_lock.release) - finally: - self._state.thread_lock.release() +from cashet._sqlite_core import SQLiteStoreCore +from cashet.models import Commit, ObjectRef, TaskStatus class AsyncSQLiteStore: def __init__(self, root: Path) -> None: - self._core = _SQLiteStoreCore(root) + self._core = SQLiteStoreCore(root) self._write_lock = asyncio.Lock() - def _fingerprint_lock(self, fingerprint: str) -> _SQLiteFingerprintLock: - import hashlib - - fp_hash = hashlib.sha256(fingerprint.encode()).hexdigest()[:16] - return _SQLiteFingerprintLock(str(self._core.root / f".lock-{fp_hash}")) + def _fingerprint_lock(self, fingerprint: str) -> SQLiteFingerprintLock: + # Striping bounds the process-global lock registry at 256 entries per + # store instead of one per fingerprint forever. Distinct fingerprints + # sharing a stripe only serialize their claim sections. + stripe = hashlib.sha256(fingerprint.encode()).hexdigest()[:2] + return SQLiteFingerprintLock(str(self._core.root / f".lock-{stripe}")) @property def root(self) -> Path: diff --git a/tests/conftest.py b/tests/conftest.py index 082ae42..ce655b8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,7 @@ import pytest from cashet import Client +from tests.helpers import redis_test_url def pytest_configure(config: pytest.Config) -> None: @@ -15,7 +16,7 @@ def pytest_configure(config: pytest.Config) -> None: return try: import redis - r = redis.from_url("redis://localhost:6379/0") + r = redis.from_url(redis_test_url()) r.ping() r.close() except Exception: diff --git a/tests/helpers.py b/tests/helpers.py index 8ec2e3e..8f1f00b 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,5 +1,43 @@ from __future__ import annotations +import os +from datetime import UTC, datetime, timedelta + +from cashet.models import Commit, TaskDef, TaskStatus + +DEFAULT_TEST_REDIS_URL = "redis://localhost:6379/15" + + +def redis_test_url() -> str: + return os.environ.get("CASHET_TEST_REDIS_URL", DEFAULT_TEST_REDIS_URL) + + +def make_task_def(args_hash: str = "b" * 64) -> TaskDef: + return TaskDef( + func_hash="a" * 64, + func_name="f", + func_source="def f(): pass", + args_hash=args_hash, + args_snapshot=b"", + ) + + +def make_commit( + hash: str, + task_def: TaskDef, + *, + hours_ago: float = 0, + expires_at: datetime | None = None, + status: TaskStatus = TaskStatus.COMPLETED, +) -> Commit: + return Commit( + hash=hash, + task_def=task_def, + status=status, + created_at=datetime.now(UTC) - timedelta(hours=hours_ago), + expires_at=expires_at, + ) + class PickleableCustom: def __init__(self, val: int) -> None: diff --git a/tests/test_async_client.py b/tests/test_async_client.py index 5354838..8342fe9 100644 --- a/tests/test_async_client.py +++ b/tests/test_async_client.py @@ -404,3 +404,89 @@ def make_val(x: int) -> int: async def test_diff_missing_commit_raises(self, async_client: AsyncClient) -> None: with pytest.raises(KeyError, match="not found"): await async_client.diff("deadbeef", "cafebabe") + + +class TestHeartbeatResilience: + async def test_result_survives_failing_heartbeat_renewals(self, tmp_path: Path) -> None: + import asyncio + from datetime import timedelta + + from cashet.async_executor import AsyncLocalExecutor + from cashet.models import Commit, TaskStatus + from cashet.store import AsyncSQLiteStore + + class FlakyRenewalStore: + def __init__(self, inner: AsyncSQLiteStore) -> None: + self._inner = inner + self.running_puts = 0 + + async def put_commit(self, commit: Commit) -> None: + if commit.status is TaskStatus.RUNNING: + self.running_puts += 1 + if self.running_puts > 1: + raise RuntimeError("transient store outage") + await self._inner.put_commit(commit) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + store = FlakyRenewalStore(AsyncSQLiteStore(tmp_path / ".cashet")) + client = AsyncClient( + store_dir=tmp_path / ".cashet", + store=store, # type: ignore[arg-type] + executor=AsyncLocalExecutor(running_ttl=timedelta(milliseconds=200)), + ) + + async def slow_task() -> str: + await asyncio.sleep(0.35) + return "survived" + + ref = await client.submit(slow_task) + assert await ref.load() == "survived" + assert store.running_puts > 1 + commits = await client.log() + assert commits[0].status is TaskStatus.COMPLETED + + async def test_failed_renewal_retries_before_lease_thins(self, tmp_path: Path) -> None: + import asyncio + import time + from datetime import timedelta + + from cashet.async_executor import AsyncLocalExecutor + from cashet.models import Commit, TaskStatus + from cashet.store import AsyncSQLiteStore + + class FlakyOnceStore: + def __init__(self, inner: AsyncSQLiteStore) -> None: + self._inner = inner + self.running_puts = 0 + self.renewal_times: list[float] = [] + + async def put_commit(self, commit: Commit) -> None: + if commit.status is TaskStatus.RUNNING: + self.running_puts += 1 + if self.running_puts > 1: + self.renewal_times.append(time.monotonic()) + if len(self.renewal_times) == 1: + raise RuntimeError("transient store outage") + await self._inner.put_commit(commit) + + def __getattr__(self, name: str) -> Any: + return getattr(self._inner, name) + + store = FlakyOnceStore(AsyncSQLiteStore(tmp_path / ".cashet")) + client = AsyncClient( + store_dir=tmp_path / ".cashet", + store=store, # type: ignore[arg-type] + executor=AsyncLocalExecutor(running_ttl=timedelta(seconds=0.4)), + ) + + async def slow_task() -> str: + await asyncio.sleep(0.5) + return "done" + + ref = await client.submit(slow_task) + assert await ref.load() == "done" + assert len(store.renewal_times) >= 2 + retry_gap = store.renewal_times[1] - store.renewal_times[0] + assert retry_gap < 0.15 diff --git a/tests/test_async_redis_store.py b/tests/test_async_redis_store.py index d332443..3f2e57a 100644 --- a/tests/test_async_redis_store.py +++ b/tests/test_async_redis_store.py @@ -8,13 +8,14 @@ from cashet.models import Commit, TaskDef, TaskStatus from cashet.redis_store import AsyncRedisStore +from tests.helpers import redis_test_url pytestmark = pytest.mark.redis @pytest_asyncio.fixture async def async_redis_store() -> AsyncRedisStore: - store = AsyncRedisStore() + store = AsyncRedisStore(redis_test_url()) await store._redis.flushdb() return store diff --git a/tests/test_cli.py b/tests/test_cli.py index c60b329..5d0e1bd 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -351,6 +351,49 @@ def add(x: int, y: int) -> int: assert client2.get(ref.commit_hash) == 3 +class TestStoreDirResolution: + def test_read_command_does_not_create_store( + self, cli_runner: CliRunner, tmp_path: Path + ) -> None: + missing = tmp_path / ".cashet" + result = cli_runner.invoke(main, ["log"], env={"CASHET_DIR": str(missing)}) + assert result.exit_code == 1 + assert "No cashet store found" in result.output + assert not missing.exists() + + def test_store_dir_flag_overrides_default( + self, cli_runner: CliRunner, store_dir: Path + ) -> None: + client = Client(store_dir=store_dir) + + def val() -> int: + return 7 + + client.submit(val) + result = cli_runner.invoke(main, ["--store-dir", str(store_dir), "log"]) + assert result.exit_code == 0 + assert "completed" in result.output + + def test_import_creates_missing_store( + self, cli_runner: CliRunner, store_dir: Path, tmp_path: Path + ) -> None: + client = Client(store_dir=store_dir) + + def add(x: int, y: int) -> int: + return x + y + + ref = client.submit(add, 1, 2) + archive = tmp_path / "export.tar.gz" + client.export(archive) + + fresh = tmp_path / "fresh-store" + result = cli_runner.invoke( + main, ["--store-dir", str(fresh), "import", str(archive)] + ) + assert result.exit_code == 0 + assert Client(store_dir=fresh).get(ref.commit_hash) == 3 + + class TestInvalidate: def test_invalidate_deletes_matching(self, cli_runner: CliRunner, store_dir: Path) -> None: client = Client(store_dir=store_dir) diff --git a/tests/test_core.py b/tests/test_core.py index 39b1078..45dbbf6 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -97,6 +97,17 @@ def expensive(x: int) -> int: assert ref1.hash == ref2.hash + def test_cache_hit_leaves_stored_status_completed(self, client: Client) -> None: + def expensive(x: int) -> int: + return x * x + + client.submit(expensive, 5) + client.submit(expensive, 5) + client.submit(expensive, 5) + commits = client.log() + assert len(commits) == 1 + assert commits[0].status is TaskStatus.COMPLETED + def test_different_args_different_results(self, client: Client) -> None: def double(x: int) -> int: return x * 2 diff --git a/tests/test_hashing.py b/tests/test_hashing.py index 0ad23e0..eff9b73 100644 --- a/tests/test_hashing.py +++ b/tests/test_hashing.py @@ -343,6 +343,25 @@ def identity(data: Any) -> Any: ref2 = client.submit(identity, {1: 2}) assert ref1.hash != ref2.hash + def test_set_hash_ignores_item_repr_order(self) -> None: + reprs: dict[int, str] = {} + + class SlottedVal: + __slots__ = ("val",) + + def __init__(self, val: int) -> None: + self.val = val + + def __repr__(self) -> str: + return reprs[id(self)] + + a1, b1 = SlottedVal(1), SlottedVal(2) + a2, b2 = SlottedVal(1), SlottedVal(2) + reprs[id(a1)], reprs[id(b1)] = "0", "1" + reprs[id(a2)], reprs[id(b2)] = "1", "0" + assert hash_args({a1, b1}) == hash_args({a2, b2}) + assert hash_args(frozenset({a1, b1})) == hash_args(frozenset({a2, b2})) + class TestRecursiveStructures: def test_recursive_list_does_not_crash(self, client: Client) -> None: diff --git a/tests/test_redis_store.py b/tests/test_redis_store.py index f35683a..503f014 100644 --- a/tests/test_redis_store.py +++ b/tests/test_redis_store.py @@ -8,15 +8,17 @@ import pytest from cashet import Client +from cashet._redis_codec import access_key, commit_key, fp_key from cashet.models import Commit, TaskDef, TaskStatus -from cashet.redis_store import RedisStore, _access_key, _commit_key +from cashet.redis_store import RedisStore +from tests.helpers import make_commit, make_task_def, redis_test_url pytestmark = pytest.mark.redis @pytest.fixture def redis_store() -> RedisStore: - store = RedisStore() + store = RedisStore(redis_test_url()) store._flushdb() return store @@ -63,6 +65,56 @@ def test_find_by_fingerprint(self, redis_store: RedisStore) -> None: assert found is not None assert found.hash == commit.hash + def test_find_by_fingerprint_returns_newest_commit(self, redis_store: RedisStore) -> None: + task_def = make_task_def() + older = make_commit("f" * 64, task_def, hours_ago=2) + newer = make_commit("0" * 64, task_def, hours_ago=1) + redis_store.put_commit(older) + redis_store.put_commit(newer) + found = redis_store.find_by_fingerprint(task_def.fingerprint) + assert found is not None + assert found.hash == newer.hash + + def test_find_by_fingerprint_skips_expired_newer_commit( + self, redis_store: RedisStore + ) -> None: + task_def = make_task_def() + older = make_commit("f" * 64, task_def, hours_ago=2) + newer = make_commit( + "0" * 64, + task_def, + hours_ago=1, + expires_at=datetime.now(UTC) - timedelta(minutes=5), + ) + redis_store.put_commit(older) + redis_store.put_commit(newer) + found = redis_store.find_by_fingerprint(task_def.fingerprint) + assert found is not None + assert found.hash == older.hash + + def test_find_by_fingerprint_heals_legacy_expiry_scores( + self, redis_store: RedisStore + ) -> None: + task_def = make_task_def() + older = make_commit("f" * 64, task_def, hours_ago=2) + newer = make_commit("0" * 64, task_def, hours_ago=1) + redis_store.put_commit(older) + redis_store.put_commit(newer) + redis_client = redis_store._async_store._redis # pyright: ignore[reportPrivateUsage] + redis_store._runner.call( # pyright: ignore[reportPrivateUsage] + redis_client.zadd( + fp_key(task_def.fingerprint), + {older.hash: float("inf"), newer.hash: float("inf")}, + ) + ) + found = redis_store.find_by_fingerprint(task_def.fingerprint) + assert found is not None + assert found.hash == newer.hash + healed_score = redis_store._runner.call( # pyright: ignore[reportPrivateUsage] + redis_client.zscore(fp_key(task_def.fingerprint), newer.hash) + ) + assert healed_score == pytest.approx(newer.created_at.timestamp()) + def test_cached_put_commit_does_not_move_access_score_backwards( self, redis_store: RedisStore ) -> None: @@ -81,17 +133,17 @@ def test_cached_put_commit_does_not_move_access_score_backwards( ) redis_store.put_commit(commit) old_score = redis_store._runner.call( # pyright: ignore[reportPrivateUsage] - redis_store._async_store._redis.zscore(_access_key(), commit.hash) # pyright: ignore[reportPrivateUsage] + redis_store._async_store._redis.zscore(access_key(), commit.hash) # pyright: ignore[reportPrivateUsage] ) found = redis_store.find_by_fingerprint(task_def.fingerprint) assert found is not None touched_score = redis_store._runner.call( # pyright: ignore[reportPrivateUsage] - redis_store._async_store._redis.zscore(_access_key(), commit.hash) # pyright: ignore[reportPrivateUsage] + redis_store._async_store._redis.zscore(access_key(), commit.hash) # pyright: ignore[reportPrivateUsage] ) found.status = TaskStatus.CACHED redis_store.put_commit(found) final_score = redis_store._runner.call( # pyright: ignore[reportPrivateUsage] - redis_store._async_store._redis.zscore(_access_key(), commit.hash) # pyright: ignore[reportPrivateUsage] + redis_store._async_store._redis.zscore(access_key(), commit.hash) # pyright: ignore[reportPrivateUsage] ) assert old_score is not None assert touched_score is not None @@ -310,6 +362,37 @@ def test_evict_old_commits(self, redis_store: RedisStore) -> None: assert deleted == 1 assert redis_store.get_commit("c" * 64) is None + def test_evict_removes_ttl_expired_commits(self, redis_store: RedisStore) -> None: + expired = make_commit( + "c" * 64, + make_task_def(), + hours_ago=2, + expires_at=datetime.now(UTC) - timedelta(hours=1), + ) + fresh = make_commit("d" * 64, make_task_def(args_hash="e" * 64)) + redis_store.put_commit(expired) + redis_store.put_commit(fresh) + deleted = redis_store.evict(datetime.now(UTC) - timedelta(days=30)) + assert deleted == 1 + assert redis_store.get_commit("c" * 64) is None + assert redis_store.get_commit("d" * 64) is not None + + def test_evict_spares_running_commit_with_expired_ttl( + self, redis_store: RedisStore + ) -> None: + running = make_commit( + "c" * 64, + make_task_def(), + hours_ago=2, + expires_at=datetime.now(UTC) - timedelta(hours=1), + status=TaskStatus.RUNNING, + ) + redis_store.put_commit(running) + assert redis_store.evict(datetime.now(UTC) - timedelta(days=30)) == 0 + fetched = redis_store.get_commit("c" * 64) + assert fetched is not None + assert fetched.status is TaskStatus.RUNNING + def test_evict_backfills_partial_last_access_index(self, redis_store: RedisStore) -> None: old = datetime.now(UTC) - timedelta(days=10) hashes = ["c" * 64, "d" * 64] @@ -331,21 +414,21 @@ def test_evict_backfills_partial_last_access_index(self, redis_store: RedisStore ) redis_store.put_commit(commit) raw = redis_store._runner.call( # pyright: ignore[reportPrivateUsage] - redis_store._async_store._redis.get(_commit_key(h)) # pyright: ignore[reportPrivateUsage] + redis_store._async_store._redis.get(commit_key(h)) # pyright: ignore[reportPrivateUsage] ) data = json.loads(raw) data["last_accessed_at"] = old.isoformat() redis_store._runner.call( # pyright: ignore[reportPrivateUsage] redis_store._async_store._redis.set( # pyright: ignore[reportPrivateUsage] - _commit_key(h), json.dumps(data, separators=(",", ":")).encode() + commit_key(h), json.dumps(data, separators=(",", ":")).encode() ) ) redis_store._runner.call( # pyright: ignore[reportPrivateUsage] - redis_store._async_store._redis.zrem(_access_key(), hashes[0]) # pyright: ignore[reportPrivateUsage] + redis_store._async_store._redis.zrem(access_key(), hashes[0]) # pyright: ignore[reportPrivateUsage] ) redis_store._runner.call( # pyright: ignore[reportPrivateUsage] redis_store._async_store._redis.zadd( # pyright: ignore[reportPrivateUsage] - _access_key(), {hashes[1]: old.timestamp()} + access_key(), {hashes[1]: old.timestamp()} ) ) @@ -598,9 +681,7 @@ def test_find_running_tracks_running_index(self, redis_store: RedisStore) -> Non def test_delete_does_not_drop_blob_when_ref_counter_missing( self, redis_store: RedisStore ) -> None: - import redis as _redis - - from cashet.redis_store import _blob_ref_key + from cashet._redis_codec import blob_ref_key shared = redis_store.put_blob(b"shared-payload") for h, fh in (("a" * 64, "1" * 64), ("b" * 64, "2" * 64)): @@ -615,8 +696,12 @@ def test_delete_does_not_drop_blob_when_ref_counter_missing( status=TaskStatus.COMPLETED, ) ) - # Simulate a lost/inconsistent ref counter for the shared blob. - _redis.Redis().delete(_blob_ref_key(shared.hash)) + # Simulate a lost/inconsistent ref counter for the shared blob, on the + # same database the store under test uses. + redis_client = redis_store._async_store._redis # pyright: ignore[reportPrivateUsage] + redis_store._runner.call( # pyright: ignore[reportPrivateUsage] + redis_client.delete(blob_ref_key(shared.hash)) + ) redis_store.delete_commit("a" * 64) diff --git a/tests/test_serializer.py b/tests/test_serializer.py index 51fda45..157a6a5 100644 --- a/tests/test_serializer.py +++ b/tests/test_serializer.py @@ -11,7 +11,7 @@ class TestSerializer: def test_custom_serializer(self, store_dir: Path) -> None: - from cashet.hashing import JsonSerializer + from cashet.serializers import JsonSerializer client = Client(store_dir=store_dir, serializer=JsonSerializer()) @@ -67,11 +67,11 @@ class TestPickleWarning: def test_warns_once_per_process(self, store_dir: Path) -> None: import warnings - import cashet.hashing as hashing_mod + import cashet.serializers as serializers_mod - original = hashing_mod._pickle_warning_issued + original = serializers_mod._pickle_warning_issued try: - hashing_mod._pickle_warning_issued = False + serializers_mod._pickle_warning_issued = False with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") Client(store_dir=store_dir) @@ -79,6 +79,6 @@ def test_warns_once_per_process(self, store_dir: Path) -> None: pickle_warnings = [x for x in w if "PickleSerializer" in str(x.message)] assert len(pickle_warnings) == 1 finally: - hashing_mod._pickle_warning_issued = original + serializers_mod._pickle_warning_issued = original diff --git a/tests/test_server.py b/tests/test_server.py index 61d825f..9e02ec0 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -10,7 +10,7 @@ from cashet.async_client import AsyncClient from cashet.client import Client -from cashet.hashing import PickleSerializer +from cashet.serializers import PickleSerializer from cashet.server import create_app, create_async_app @@ -252,17 +252,33 @@ def test_submit_missing_task_raises_400(self, server_client: TestClient) -> None response = server_client.post("/submit", json={"args": [1]}) assert response.status_code == 400 - def test_submit_internal_error_is_generic(self, tmp_path: Path) -> None: + def test_submit_task_failure_returns_422_without_paths(self, tmp_path: Path) -> None: client = Client(store_dir=tmp_path / ".cashet") app = create_app(client, tasks={"boom": _boom}) tc = TestClient(app) response = tc.post("/submit", json={"task": "boom"}) - assert response.status_code == 500 - assert response.json()["error"] == "Internal server error" - assert "secret internal path" not in response.text + assert response.status_code == 422 + data = response.json() + assert data["error"] == "task failed" + assert data["detail"] == "RuntimeError: secret internal path" + assert 'File "' not in response.text class TestAsyncServerSubmit: + async def test_async_submit_task_failure_returns_422(self, tmp_path: Path) -> None: + client = AsyncClient(store_dir=tmp_path / ".cashet") + app = create_async_app(client, tasks={"boom": _boom}) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://test" + ) as ac: + response = await ac.post("/submit", json={"task": "boom"}) + assert response.status_code == 422 + data = response.json() + assert data["error"] == "task failed" + assert data["detail"] == "RuntimeError: secret internal path" + assert 'File "' not in response.text + await client.close() + async def test_async_submit_registered_task_with_sqlite(self, tmp_path: Path) -> None: client = AsyncClient(store_dir=tmp_path / ".cashet") app = create_async_app(client, tasks={"add": _add}) diff --git a/tests/test_store.py b/tests/test_store.py index 576ef7a..a935a26 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -9,9 +9,11 @@ from cashet import Client, TaskError from cashet.executor import LocalExecutor -from cashet.hashing import PickleSerializer, build_task_def +from cashet.hashing import build_task_def from cashet.models import Commit, ObjectRef, StorageTier, TaskStatus +from cashet.serializers import PickleSerializer from cashet.store import SQLiteStore +from tests.helpers import make_commit, make_task_def class TestProcessSafety: @@ -336,6 +338,46 @@ def other(v: int = i) -> int: class TestStoreOperations: + def test_access_bump_throttled_to_granularity(self, store_dir: Path) -> None: + from freezegun import freeze_time + + store = SQLiteStore(store_dir) + task_def = make_task_def() + + def last_accessed() -> str: + row = store._connect().execute( # pyright: ignore[reportPrivateUsage] + "SELECT last_accessed_at FROM commits WHERE hash = ?", ("c" * 64,) + ).fetchone() + return row[0] + + with freeze_time("2026-01-01 00:00:00") as frozen: + store.put_commit(make_commit("c" * 64, task_def)) + initial = last_accessed() + + frozen.tick(60) + assert store.find_by_fingerprint(task_def.fingerprint) is not None + assert last_accessed() == initial + + frozen.tick(7200) + assert store.find_by_fingerprint(task_def.fingerprint) is not None + assert last_accessed() > initial + + def test_connection_uses_normal_synchronous(self, store_dir: Path) -> None: + store = SQLiteStore(store_dir) + mode = store._connect().execute("PRAGMA synchronous").fetchone()[0] # pyright: ignore[reportPrivateUsage] + assert mode == 1 + + def test_find_by_fingerprint_returns_newest_commit(self, store_dir: Path) -> None: + store = SQLiteStore(store_dir) + task_def = make_task_def() + older = make_commit("f" * 64, task_def, hours_ago=2) + newer = make_commit("0" * 64, task_def, hours_ago=1) + store.put_commit(older) + store.put_commit(newer) + found = store.find_by_fingerprint(task_def.fingerprint) + assert found is not None + assert found.hash == newer.hash + def test_log(self, client: Client) -> None: def f(x: int) -> int: return x @@ -346,6 +388,19 @@ def f(x: int) -> int: log = client.log() assert len(log) == 3 + def test_fingerprint_lock_registry_is_striped(self, client: Client) -> None: + from cashet._locks import SQLITE_LOCKS + + def f(x: int) -> int: + return x + + for i in range(5): + client.submit(f, i) + root = str(client.store.root) + names = [Path(p).name for p in SQLITE_LOCKS if p.startswith(root)] + assert names + assert all(len(name) == len(".lock-xx") for name in names) + def test_retries_roundtrip(self, client: Client) -> None: def stable() -> int: return 42 @@ -570,6 +625,68 @@ def test_wrong_hash_ref_raises_valueerror(self, store_dir: Path) -> None: with pytest.raises(ValueError, match="integrity check failed"): store.get_blob(fake_ref) + def test_corrupt_blob_self_heals_on_rewrite(self, store_dir: Path) -> None: + store = SQLiteStore(store_dir) + data = b"partially written payload" * 100 + ref = store.put_blob(data) + obj_path = store.objects_dir / ref.hash[:2] / ref.hash[2:] + obj_path.write_bytes(b"truncated junk") + with pytest.raises(ValueError, match="integrity check failed"): + store.get_blob(ref) + assert not obj_path.exists() + healed = store.put_blob(data) + assert healed.hash == ref.hash + assert store.get_blob(healed) == data + + def test_put_blob_leaves_no_temp_files(self, store_dir: Path) -> None: + store = SQLiteStore(store_dir) + store.put_blob(b"z" * 5000) + leftovers = list(store.objects_dir.rglob("*.tmp")) + assert leftovers == [] + + def test_put_blob_cleans_temp_file_on_failed_write( + self, store_dir: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + import cashet._sqlite_core as core_mod + + store = SQLiteStore(store_dir) + + def boom(src: object, dst: object) -> None: + raise OSError("disk full") + + monkeypatch.setattr(core_mod.os, "replace", boom) + with pytest.raises(OSError, match="disk full"): + store.put_blob(b"z" * 5000) + assert list(store.objects_dir.rglob("*.tmp")) == [] + + def test_stats_ignores_leftover_temp_files(self, store_dir: Path) -> None: + store = SQLiteStore(store_dir) + ref = store.put_blob(b"z" * 5000) + before = store.stats() + orphan = store.objects_dir / ref.hash[:2] / "deadbeef.1.2.tmp" + orphan.write_bytes(b"crash debris") + after = store.stats() + assert after["blob_objects"] == before["blob_objects"] + assert after["blob_bytes"] == before["blob_bytes"] + + def test_evict_sweeps_stale_temp_files(self, store_dir: Path) -> None: + import os + import time + + store = SQLiteStore(store_dir) + store.put_blob(b"z" * 5000) + prefix_dir = next(p for p in store.objects_dir.iterdir() if p.is_dir()) + stale = prefix_dir / "deadbeef.1.2.tmp" + stale.write_bytes(b"crash debris") + old = time.time() - 7200 + os.utime(stale, (old, old)) + fresh = prefix_dir / "cafebabe.3.4.tmp" + fresh.write_bytes(b"in-flight write") + + store.evict(datetime.now(UTC) - timedelta(days=30)) + assert not stale.exists() + assert fresh.exists() + class TestInlineStorage: def test_small_blob_uses_inline_tier(self, store_dir: Path) -> None: @@ -624,11 +741,11 @@ def test_inline_delete_removes_orphans(self, store_dir: Path) -> None: assert store.blob_exists(ref.hash) is False def test_find_by_fingerprint_survives_locked_access_bump(self, tmp_path: Path) -> None: + from cashet._sqlite_core import SQLiteStoreCore from cashet.models import Commit, TaskDef, TaskStatus - from cashet.store import _SQLiteStoreCore root = tmp_path / ".cashet" - core = _SQLiteStoreCore(root) + core = SQLiteStoreCore(root) task_def = TaskDef( func_hash="a", func_name="f", func_source="", args_hash="b", args_snapshot=b"" ) @@ -656,11 +773,11 @@ def test_find_by_fingerprint_survives_locked_access_bump(self, tmp_path: Path) - def test_evict_survives_vacuum_failure( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: + from cashet._sqlite_core import SQLiteStoreCore from cashet.models import Commit, TaskDef, TaskStatus - from cashet.store import _SQLiteStoreCore root = tmp_path / ".cashet" - core = _SQLiteStoreCore(root) + core = SQLiteStoreCore(root) task_def = TaskDef( func_hash="a", func_name="f", func_source="", args_hash="b", args_snapshot=b"" ) @@ -723,6 +840,38 @@ def make_val(x: int) -> int: assert deleted == 0 assert client.stats()["total_commits"] == 1 + def test_gc_spares_running_commit_with_expired_ttl(self, store_dir: Path) -> None: + store = SQLiteStore(store_dir) + running = make_commit( + "c" * 64, + make_task_def(), + hours_ago=2, + expires_at=datetime.now(UTC) - timedelta(hours=1), + status=TaskStatus.RUNNING, + ) + store.put_commit(running) + assert store.evict(datetime.now(UTC) - timedelta(days=30)) == 0 + fetched = store.get_commit("c" * 64) + assert fetched is not None + assert fetched.status is TaskStatus.RUNNING + + def test_gc_removes_ttl_expired_commits(self, client: Client) -> None: + from freezegun import freeze_time + + def make_val(x: int) -> int: + return x + + with freeze_time("2026-01-01 00:00:00") as frozen: + client.submit(make_val, 1, _ttl=60) + client.submit(make_val, 2) + assert client.stats()["total_commits"] == 2 + + frozen.tick(3600) + deleted = client.gc(timedelta(days=30)) + assert deleted == 1 + assert client.stats()["total_commits"] == 1 + assert client.log()[0].expires_at is None + def test_gc_default_30_days(self, client: Client) -> None: def make_val(x: int) -> int: return x @@ -1383,7 +1532,7 @@ def _build_base_db(self, db_path: str) -> None: conn.close() def test_migration_from_base_schema(self, tmp_path: Path) -> None: - from cashet.store import _SQLiteStoreCore + from cashet._sqlite_core import SQLiteStoreCore root = tmp_path / ".cashet" root.mkdir() @@ -1392,7 +1541,7 @@ def test_migration_from_base_schema(self, tmp_path: Path) -> None: self._build_base_db(str(db_path)) - core = _SQLiteStoreCore(root) + core = SQLiteStoreCore(root) try: col_names = [ r[1] @@ -1413,7 +1562,7 @@ def test_migration_from_base_schema(self, tmp_path: Path) -> None: core.close() def test_migration_idempotent(self, tmp_path: Path) -> None: - from cashet.store import _SQLiteStoreCore + from cashet._sqlite_core import SQLiteStoreCore root = tmp_path / ".cashet" root.mkdir() @@ -1422,10 +1571,10 @@ def test_migration_idempotent(self, tmp_path: Path) -> None: self._build_base_db(str(db_path)) - core1 = _SQLiteStoreCore(root) + core1 = SQLiteStoreCore(root) core1.close() - core2 = _SQLiteStoreCore(root) + core2 = SQLiteStoreCore(root) try: col_names = [ r[1] @@ -1439,7 +1588,7 @@ def test_migration_idempotent(self, tmp_path: Path) -> None: core2.close() def test_operations_after_migration(self, tmp_path: Path) -> None: - from cashet.store import _SQLiteStoreCore + from cashet._sqlite_core import SQLiteStoreCore root = tmp_path / ".cashet" root.mkdir() @@ -1448,7 +1597,7 @@ def test_operations_after_migration(self, tmp_path: Path) -> None: self._build_base_db(str(db_path)) - core = _SQLiteStoreCore(root) + core = SQLiteStoreCore(root) try: from cashet.dag import build_commit, resolve_input_refs from cashet.hashing import build_task_def @@ -1478,7 +1627,7 @@ def add(x: int, y: int) -> int: core.close() def test_partial_migration_missing_retries_only(self, tmp_path: Path) -> None: - from cashet.store import _SQLiteStoreCore + from cashet._sqlite_core import SQLiteStoreCore root = tmp_path / ".cashet" root.mkdir() @@ -1496,7 +1645,7 @@ def test_partial_migration_missing_retries_only(self, tmp_path: Path) -> None: ) pre_conn.close() - core = _SQLiteStoreCore(root) + core = SQLiteStoreCore(root) try: col_names = [ r[1] diff --git a/uv.lock b/uv.lock index 8f80400..55b7cc6 100644 --- a/uv.lock +++ b/uv.lock @@ -26,7 +26,7 @@ wheels = [ [[package]] name = "cashet" -version = "0.4.5" +version = "0.5.0" source = { editable = "." } dependencies = [ { name = "click" },