diff --git a/Makefile b/Makefile index 2ab74ab..b4cc110 100644 --- a/Makefile +++ b/Makefile @@ -992,6 +992,26 @@ cold-hydrate: bootstrap ## genesis a fresh peer from cloud, M-aware: pull every @echo ">> hydration complete:" @ls -lh $(HYDRATE_DIR)/*.db 2>/dev/null +# ---- Live-snapshot clone tier (raw .db on Spaces; skip pack/unpack entirely) +# Producer streams each shard via SQLite Online Backup API + multipart upload; +# consumer pulls raw .db files in parallel. Recovery collapses to ~download time +# (the FTS index travels inside the .db, no rebuild needed). JUST_ENOUGH=1 +# strips chunks.content into per-chunk blobs/ for SPV / mobile peers. +cold-stream-snapshot: bootstrap ## producer: snapshot shards to bucket as raw .db [SHARDS_DIR=path JUST_ENOUGH=1 SNAPSHOT_ID=...] + @if [ -z "$(SHARDS_DIR)" ]; then echo "SHARDS_DIR= required"; exit 2; fi + $(ARBORIST) cold stream-snapshot --shards-dir $(SHARDS_DIR) \ + $(if $(JUST_ENOUGH),--just-enough,) \ + $(if $(SNAPSHOT_ID),--snapshot-id $(SNAPSHOT_ID),) + +cold-clone: bootstrap ## consumer: clone raw .db shards from bucket (fast, no restore) [SHARDS_DIR=path SNAPSHOT_ID=... WORKERS=N] + @if [ -z "$(SHARDS_DIR)" ]; then echo "SHARDS_DIR= required"; exit 2; fi + @mkdir -p $(SHARDS_DIR) + $(ARBORIST) cold clone --shards-dir $(SHARDS_DIR) \ + $(if $(SNAPSHOT_ID),--snapshot-id $(SNAPSHOT_ID),) \ + $(if $(WORKERS),--workers $(WORKERS),) + @echo ">> verifying cloned shards" + @$(ARBORIST) cold verify --shards-dir $(SHARDS_DIR) + cold-pack-all-dvd: bootstrap ## same fan-out for DVD burning (no S3); RAM-aware @if [ -z "$(LOCAL_DIR)" ]; then echo "LOCAL_DIR= required"; exit 2; fi @mkdir -p $(LOCAL_DIR) diff --git a/arborist/cli.py b/arborist/cli.py index 5adb2ce..35c047f 100644 --- a/arborist/cli.py +++ b/arborist/cli.py @@ -3247,6 +3247,44 @@ def _cmd_cold_verify(args: argparse.Namespace) -> int: return 0 +def _cmd_cold_stream_snapshot(args: argparse.Namespace) -> int: + """Producer: SQLite-Backup-API snapshot each shard to a raw `.db`, + multipart-upload to the bucket, swap CURRENT. Optionally Tier B + (`--just-enough`): strip chunks.content into per-chunk blobs.""" + from arborist.cold_clone import stream_snapshot + backend = _make_cold_backend() + manifest = stream_snapshot( + Path(args.shards_dir).expanduser(), + backend, + just_enough=bool(getattr(args, "just_enough", False)), + snapshot_id=getattr(args, "snapshot_id", None), + ) + import json as _json + print(_json.dumps({ + "snapshot_id": manifest.snapshot_id, + "shards": [s.__dict__ for s in manifest.shards], + "just_enough": manifest.just_enough, + }, indent=2)) + return 0 + + +def _cmd_cold_clone(args: argparse.Namespace) -> int: + """Consumer: read CURRENT (or pinned snapshot) and pull raw `.db` + shards into --shards-dir in parallel. No pack/unpack, no FTS rebuild + — the index is already inside the file.""" + from arborist.cold_clone import clone_from_bucket + backend = _make_cold_backend() + manifest = clone_from_bucket( + Path(args.shards_dir).expanduser(), + backend, + snapshot_id=getattr(args, "snapshot_id", None) or None, + workers=int(getattr(args, "workers", 4) or 4), + ) + print(f"clone OK: {manifest.M} shards from snapshot {manifest.snapshot_id} " + f"into {args.shards_dir} (just_enough={manifest.just_enough})") + return 0 + + def _cmd_cold_stats(args: argparse.Namespace) -> int: from arborist.cold_object import PACK_PREFIX @@ -6256,6 +6294,51 @@ def build_parser() -> argparse.ArgumentParser: ) cold_verify.set_defaults(func=_cmd_cold_verify) + cold_stream_snapshot = cold_sub.add_parser( + "stream-snapshot", + help=( + "live-snapshot producer: SQLite Backup API each shard to a " + "raw .db, multipart-upload to bucket, swap CURRENT. " + "--just-enough strips chunks.content into per-chunk blobs " + "(Tier B, for SPV / mobile peers)." + ), + ) + cold_stream_snapshot.add_argument( + "--shards-dir", dest="shards_dir", required=True, + help="directory of shards to snapshot", + ) + cold_stream_snapshot.add_argument( + "--just-enough", dest="just_enough", action="store_true", + help="Tier B: strip chunks.content into blobs/; ship metadata-only shards", + ) + cold_stream_snapshot.add_argument( + "--snapshot-id", dest="snapshot_id", default=None, + help="pin a specific snapshot id (default: snap-)", + ) + cold_stream_snapshot.set_defaults(func=_cmd_cold_stream_snapshot) + + cold_clone = cold_sub.add_parser( + "clone", + help=( + "live-snapshot consumer: read CURRENT (or a pinned --snapshot-id), " + "multipart-download raw .db shards into --shards-dir in parallel. " + "No pack/unpack, no FTS rebuild — the index travels inside the file." + ), + ) + cold_clone.add_argument( + "--shards-dir", dest="shards_dir", required=True, + help="target dir to clone into", + ) + cold_clone.add_argument( + "--snapshot-id", dest="snapshot_id", default=None, + help="pin a snapshot id (default: read CURRENT pointer)", + ) + cold_clone.add_argument( + "--workers", dest="workers", type=int, default=4, + help="parallel multipart downloads (default 4 = one per shard)", + ) + cold_clone.set_defaults(func=_cmd_cold_clone) + cold_list = cold_sub.add_parser( "list", help="enumerate packs in the bucket with metadata (pack_hash, size, chunk_count) for new-peer hydration", diff --git a/arborist/cold_clone.py b/arborist/cold_clone.py new file mode 100644 index 0000000..d5da994 --- /dev/null +++ b/arborist/cold_clone.py @@ -0,0 +1,377 @@ +"""Live-snapshot tier — clone raw SQLite shards from the bucket. + +Two channels live in the same bucket alongside the pack tier +(``packs/..tar.zst``): + +* **Tier A — full clone.** Producer ``cold stream-snapshot`` uses the + SQLite Online Backup API (`sqlite3.Connection.backup`) to copy each + shard to a temp file without blocking ingest, then multipart-uploads + the raw ``.db`` to ``clones//.db``. A small + ``clones/CURRENT.json`` manifest points at the latest snapshot. + Consumer ``cold clone`` reads the manifest and pulls the shards in + parallel — no pack/unpack, no edge fan-out, no FTS rebuild (the + shadow tables travel verbatim inside the ``.db``). Recovery collapses + to ~download time. + +* **Tier B — just-enough clone + JIT chunks.** Same flow, but the + producer strips ``chunks.content`` (writes NULL into the backup copy) + and ships each chunk body as a content-addressed object at + ``blobs/``. The consumer's local DB carries metadata + leaf + hashes only; the query path fetches the chunks it needs from the + bucket on demand (``hash_leaf`` verifies every fetch). Mobile / SPV + peers fit in ~a few GB instead of ~35 GB. + +The manifest layout is intentionally flat and self-describing: + + clones/ + CURRENT.json # {"snapshot_id": ..., "shards": [...]} + / + 000.db + 001.db + 002.db + 003.db + ... + blobs/ + / # Tier B only — opaque bytes, + # verified hash on fetch +""" + +from __future__ import annotations + +import io +import json +import sqlite3 +import sys +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional + +from arborist.cold_object import ObjectStoreBackend + +CURRENT_KEY = "clones/CURRENT.json" +SHARD_PREFIX = "clones" +BLOB_PREFIX = "blobs" +FORMAT = "raw_sqlite_v1" + + +# ---------- manifest ------------------------------------------------------- + + +@dataclass +class ShardEntry: + name: str # 000.db + key: str # clones//000.db + bytes: int # size on disk + snapshot_root: Optional[str] = None # corpus-level identity if computable + + +@dataclass +class CloneManifest: + snapshot_id: str + ts: int + format: str + M: int + shards: list[ShardEntry] + just_enough: bool = False + blob_prefix: Optional[str] = None # set when just_enough=True + + def to_json(self) -> bytes: + out = { + "snapshot_id": self.snapshot_id, + "ts": self.ts, + "format": self.format, + "M": self.M, + "just_enough": self.just_enough, + "blob_prefix": self.blob_prefix, + "shards": [s.__dict__ for s in self.shards], + } + return json.dumps(out, separators=(",", ":"), sort_keys=True).encode("utf-8") + + @classmethod + def from_json(cls, body: bytes) -> "CloneManifest": + d = json.loads(body) + return cls( + snapshot_id=d["snapshot_id"], + ts=int(d["ts"]), + format=d.get("format", FORMAT), + M=int(d["M"]), + shards=[ShardEntry(**s) for s in d["shards"]], + just_enough=bool(d.get("just_enough", False)), + blob_prefix=d.get("blob_prefix"), + ) + + +def read_current_manifest(backend: ObjectStoreBackend) -> CloneManifest: + """Pull the CURRENT pointer + parse it. Caller catches missing-key.""" + body = backend.get(CURRENT_KEY) + return CloneManifest.from_json(body) + + +def write_current_manifest(backend: ObjectStoreBackend, manifest: CloneManifest) -> None: + """Atomically swap the CURRENT pointer. S3 PUT is atomic per-key.""" + backend.put(CURRENT_KEY, manifest.to_json(), content_type="application/json") + + +# ---------- producer: stream snapshots into the bucket --------------------- + + +def _backup_sqlite(src_path: Path, dst_path: Path) -> None: + """SQLite Online Backup — page-by-page copy of an open db to a file. + + Does NOT block writers on the source for long; the backup steps + through pages and yields between batches. The destination is a + self-contained, byte-correct SQLite database the moment the call + returns. Safe to run against a shard that's being written to.""" + src = sqlite3.connect(f"file:{src_path}?mode=ro", uri=True) + dst = sqlite3.connect(str(dst_path)) + try: + # pages=-1 copies in one pass (faster for a one-shot snapshot); + # sleep=0 means we don't yield between batches — fine when we're + # not contending with the source's writers (or okay either way, + # the destination is exclusively ours). + src.backup(dst, pages=-1, progress=None, sleep=0.0) + finally: + dst.close() + src.close() + + +def _strip_chunk_content_and_dump_blobs( + db_path: Path, backend: ObjectStoreBackend, *, upload_workers: int = 32 +) -> int: + """Tier B helper: in the just-backed-up copy, write each chunk's + compressed content out to ``blobs/`` and NULL the column. + + The bucket blobs are stored verbatim (zstd-packed bytes as they + appear in the chunks table) so the consumer can drop them straight + back into ``chunks.content`` after ``hash_leaf`` verifies. Uploads + fan out across ``upload_workers`` threads — a 1.5M-chunk shard is + 1.5M small PUTs and the per-request latency, not bytes-on-wire, + is the throughput cap.""" + from arborist.merkle import hash_leaf + + conn = sqlite3.connect(str(db_path)) + conn.execute("PRAGMA synchronous = OFF") + conn.execute("PRAGMA journal_mode = MEMORY") + n_uploaded = 0 + try: + # Get all chunk_ids up front (small — 8 B/row) so we can stream + # through them without holding a SELECT cursor across writes. + ids = [ + r[0] for r in conn.execute( + "SELECT chunk_id FROM chunks WHERE content IS NOT NULL" + ).fetchall() + ] + BATCH = 1000 + + def _upload(triple: tuple[int, str, bytes]) -> int: + chunk_id, leaf_hash, body = triple + if hash_leaf(body).hex() != leaf_hash: + raise ValueError( + f"chunk {chunk_id}: leaf_hash {leaf_hash[:12]} " + f"does not match content" + ) + key = f"{BLOB_PREFIX}/{leaf_hash[:2]}/{leaf_hash[2:]}" + backend.put(key, body, content_type="application/zstd") + return chunk_id + + for start in range(0, len(ids), BATCH): + batch_ids = ids[start:start + BATCH] + placeholders = ",".join("?" for _ in batch_ids) + rows = conn.execute( + f"SELECT chunk_id, leaf_hash, content FROM chunks " + f"WHERE chunk_id IN ({placeholders})", + batch_ids, + ).fetchall() + triples = [ + (cid, lh, bytes(c)) for cid, lh, c in rows if c is not None + ] + with ThreadPoolExecutor(max_workers=upload_workers) as ex: + for cid in ex.map(_upload, triples): + pass + uploaded_ids = [t[0] for t in triples] + with conn: + conn.execute( + f"UPDATE chunks SET content = NULL " + f"WHERE chunk_id IN ({','.join('?' for _ in uploaded_ids)})", + uploaded_ids, + ) + n_uploaded += len(uploaded_ids) + # Reclaim the space the zeroed-out content blobs used to occupy + # so the metadata-only shard is small enough to be worth shipping. + conn.execute("VACUUM") + finally: + conn.close() + return n_uploaded + + +def stream_snapshot( + shards_dir: Path, + backend: ObjectStoreBackend, + *, + just_enough: bool = False, + snapshot_id: Optional[str] = None, +) -> CloneManifest: + """Snapshot every shard in ``shards_dir`` via SQLite Backup API and + upload to the bucket. Returns the manifest that was just published.""" + shards_dir = Path(shards_dir).expanduser() + db_files = sorted(shards_dir.glob("00[0-9].db")) + if not db_files: + raise FileNotFoundError(f"no shards under {shards_dir}") + + snap_id = snapshot_id or f"snap-{int(time.time())}" + print(f"[stream-snapshot] snapshot_id={snap_id} just_enough={just_enough} " + f"shards={len(db_files)}", file=sys.stderr, flush=True) + entries: list[ShardEntry] = [] + for src in db_files: + t0 = time.time() + tmp_path = Path(tempfile.mkdtemp(prefix="arborist-clone-")) / src.name + try: + _backup_sqlite(src, tmp_path) + backed_up_s = time.time() - t0 + n_blobs = 0 + if just_enough: + n_blobs = _strip_chunk_content_and_dump_blobs(tmp_path, backend) + stripped_s = time.time() - t0 - backed_up_s + else: + stripped_s = 0.0 + size = tmp_path.stat().st_size + key = f"{SHARD_PREFIX}/{snap_id}/{src.name}" + up_t = time.time() + backend.put_file(key, tmp_path, content_type="application/x-sqlite3") + up_s = time.time() - up_t + entries.append(ShardEntry(name=src.name, key=key, bytes=size)) + print( + f"[stream-snapshot] {src.name}: backup {backed_up_s:.1f}s, " + f"{'strip ' + str(n_blobs) + ' blobs in %.1fs, ' % stripped_s if just_enough else ''}" + f"upload {up_s:.1f}s, {size / 1e9:.2f} GB", + file=sys.stderr, flush=True, + ) + finally: + try: + tmp_path.unlink() + tmp_path.parent.rmdir() + except OSError: + pass + + manifest = CloneManifest( + snapshot_id=snap_id, + ts=int(time.time()), + format=FORMAT, + M=len(entries), + shards=entries, + just_enough=just_enough, + blob_prefix=BLOB_PREFIX if just_enough else None, + ) + write_current_manifest(backend, manifest) + print(f"[stream-snapshot] CURRENT -> {snap_id}", file=sys.stderr, flush=True) + return manifest + + +# ---------- consumer: clone the shards down --------------------------------- + + +def clone_from_bucket( + shards_dir: Path, + backend: ObjectStoreBackend, + *, + snapshot_id: Optional[str] = None, + workers: int = 4, +) -> CloneManifest: + """Read the CURRENT manifest (or a pinned ``snapshot_id``) and pull + every shard into ``shards_dir`` in parallel. Returns the manifest.""" + shards_dir = Path(shards_dir).expanduser() + shards_dir.mkdir(parents=True, exist_ok=True) + if snapshot_id is None: + manifest = read_current_manifest(backend) + else: + # Pinned mode: reconstruct from listing keys under the snap dir. + keys = sorted( + k for k in backend.list_keys(f"{SHARD_PREFIX}/{snapshot_id}/") + if k.endswith(".db") + ) + if not keys: + raise FileNotFoundError( + f"no shards under {SHARD_PREFIX}/{snapshot_id}/ in bucket" + ) + entries = [ + ShardEntry(name=Path(k).name, key=k, bytes=backend.object_size(k) or 0) + for k in keys + ] + manifest = CloneManifest( + snapshot_id=snapshot_id, + ts=int(time.time()), + format=FORMAT, + M=len(entries), + shards=entries, + ) + print( + f"[clone] snapshot_id={manifest.snapshot_id} shards={manifest.M} " + f"just_enough={manifest.just_enough}", + file=sys.stderr, flush=True, + ) + + def _pull(entry: ShardEntry) -> tuple[str, float, int]: + t = time.time() + dst = shards_dir / entry.name + backend.get_file(entry.key, dst) + return entry.name, time.time() - t, dst.stat().st_size + + workers = max(1, min(workers, len(manifest.shards))) + t_all = time.time() + with ThreadPoolExecutor(max_workers=workers) as ex: + futs = {ex.submit(_pull, e): e for e in manifest.shards} + for fut in as_completed(futs): + name, dt, size = fut.result() + print( + f"[clone] {name}: {dt:.1f}s, {size / 1e9:.2f} GB", + file=sys.stderr, flush=True, + ) + total = time.time() - t_all + total_bytes = sum(e.bytes for e in manifest.shards) + print( + f"[clone] {manifest.M} shards in {total:.1f}s " + f"({total_bytes / 1e9:.2f} GB total, {workers}-way parallel)", + file=sys.stderr, flush=True, + ) + return manifest + + +# ---------- JIT chunk fetch (Tier B consumer) ------------------------------- + + +def fetch_chunk_jit( + conn: sqlite3.Connection, + chunk_id: int, + backend: ObjectStoreBackend, +) -> Optional[bytes]: + """Tier B retrieval: fetch a chunk's content from the bucket on + demand and cache it back into the local row. Returns the raw + (still-zstd-packed) bytes, ready for ``arborist.compress.unpack_chunk`` + to decode. Returns None if the row doesn't exist or already has + content (the caller should just read the local copy in that case).""" + from arborist.merkle import hash_leaf + + row = conn.execute( + "SELECT leaf_hash, content FROM chunks WHERE chunk_id = ?", (chunk_id,) + ).fetchone() + if row is None: + return None + leaf_hash, content = row + if content is not None: + return bytes(content) + key = f"{BLOB_PREFIX}/{leaf_hash[:2]}/{leaf_hash[2:]}" + body = backend.get(key) + if hash_leaf(body).hex() != leaf_hash: + raise ValueError( + f"chunk {chunk_id}: bucket blob hash does not match leaf_hash" + ) + # cache for next time so the second query is free + with conn: + conn.execute( + "UPDATE chunks SET content = ? WHERE chunk_id = ?", + (body, chunk_id), + ) + return body diff --git a/arborist/cold_object.py b/arborist/cold_object.py index 3d60a80..c86cd84 100644 --- a/arborist/cold_object.py +++ b/arborist/cold_object.py @@ -177,6 +177,13 @@ class ObjectStoreBackend(abc.ABC): def get(self, key: str) -> bytes: ... + @abc.abstractmethod + def get_file(self, key: str, path: Path) -> None: + """Stream a (potentially multi-GB) object to a local file WITHOUT + loading it into memory. The S3 backend hands the path to boto3 + ``download_file`` which fetches multipart in parallel. Used by + ``cold clone`` to pull raw ``.db`` shards from the bucket.""" + @abc.abstractmethod def head(self, key: str) -> bool: """True if the object exists; False otherwise. Never raises on missing.""" @@ -406,6 +413,18 @@ class S3CompatibleBackend(ObjectStoreBackend): resp = self._client.get_object(Bucket=self._bucket, Key=key) return resp["Body"].read() + def get_file(self, key: str, path: Path) -> None: + # boto3 download_file streams multipart in parallel into the + # destination file — memory stays bounded by max_concurrency × + # part_size, not by object size. Required for shard `.db` clones + # that can run 8-12 GB each. + self._client.download_file( + Bucket=self._bucket, + Key=key, + Filename=str(path), + Config=self._transfer_config, + ) + def head(self, key: str) -> bool: try: self._client.head_object(Bucket=self._bucket, Key=key) @@ -471,6 +490,12 @@ class MemoryBackend(ObjectStoreBackend): raise KeyError(key) return self._store[key] + def get_file(self, key: str, path: Path) -> None: + if key not in self._store: + raise KeyError(key) + with open(path, "wb") as f: + f.write(self._store[key]) + def head(self, key: str) -> bool: return key in self._store