russell@unturf.com
7bda7450c8
Some rows in shards-genesis-v2 came back from the chunks.content BLOB column as instead of bytes — SQLite's type-affinity rule lets a BLOB-affinity column hold any storage class. The strip-and-upload path crashed on bytes(some_str) without an encoding. latin-1 preserves arbitrary byte values 1:1, so the leaf_hash check still matches.
419 lines
15 KiB
Python
419 lines
15 KiB
Python
"""Live-snapshot tier — clone raw SQLite shards from the bucket.
|
||
|
||
Two channels live in the same bucket alongside the pack tier
|
||
(``packs/<hash>.<kind>.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/<snap-id>/<NN>.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/<hash>``. 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": [...]}
|
||
<snap-id>/
|
||
000.db
|
||
001.db
|
||
002.db
|
||
003.db
|
||
...
|
||
blobs/
|
||
<leaf_hash[:2]>/<leaf_hash[2:]> # 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/<snap-id>/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/<leaf_hash>`` 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()
|
||
# SQLite type affinity: a `BLOB` column can come back as
|
||
# bytes / bytearray / memoryview, OR as `str` if the value
|
||
# was originally INSERTed as text. latin-1 preserves
|
||
# arbitrary byte values 1:1 (codepoints 0..255) so the
|
||
# hash check below still matches.
|
||
def _as_bytes(v):
|
||
if isinstance(v, (bytes, bytearray)):
|
||
return bytes(v)
|
||
if isinstance(v, memoryview):
|
||
return v.tobytes()
|
||
return v.encode("latin-1")
|
||
triples = [
|
||
(cid, lh, _as_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 _snapshot_one_shard(
|
||
src: Path,
|
||
backend: ObjectStoreBackend,
|
||
*,
|
||
snap_id: str,
|
||
just_enough: bool,
|
||
) -> ShardEntry:
|
||
"""ThreadPool worker: SQLite-Backup a shard + (optionally strip blobs)
|
||
+ multipart-upload. Returns the manifest entry."""
|
||
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
|
||
stripped_s = 0.0
|
||
if just_enough:
|
||
n_blobs = _strip_chunk_content_and_dump_blobs(tmp_path, backend)
|
||
stripped_s = time.time() - t0 - backed_up_s
|
||
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
|
||
msg = (
|
||
f"[stream-snapshot] {src.name}: backup {backed_up_s:.1f}s, "
|
||
f"{'strip %d blobs in %.1fs, ' % (n_blobs, stripped_s) if just_enough else ''}"
|
||
f"upload {up_s:.1f}s, {size / 1e9:.2f} GB"
|
||
)
|
||
print(msg, file=sys.stderr, flush=True)
|
||
return ShardEntry(name=src.name, key=key, bytes=size)
|
||
finally:
|
||
try:
|
||
tmp_path.unlink()
|
||
tmp_path.parent.rmdir()
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def stream_snapshot(
|
||
shards_dir: Path,
|
||
backend: ObjectStoreBackend,
|
||
*,
|
||
just_enough: bool = False,
|
||
snapshot_id: Optional[str] = None,
|
||
workers: int = 0,
|
||
) -> CloneManifest:
|
||
"""Snapshot every shard in ``shards_dir`` via SQLite Backup API and
|
||
upload to the bucket. Shards process in parallel (default: one
|
||
worker per shard) — each worker has its own multipart-upload
|
||
concurrency on top, so 4 shards × 10 parts in flight saturates the
|
||
link in a way the previous serial loop could not. 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())}"
|
||
n_workers = max(1, min(workers if workers > 0 else len(db_files), len(db_files)))
|
||
print(f"[stream-snapshot] snapshot_id={snap_id} just_enough={just_enough} "
|
||
f"shards={len(db_files)} workers={n_workers}",
|
||
file=sys.stderr, flush=True)
|
||
|
||
t_all = time.time()
|
||
with ThreadPoolExecutor(max_workers=n_workers) as ex:
|
||
futs = [
|
||
ex.submit(
|
||
_snapshot_one_shard, src, backend,
|
||
snap_id=snap_id, just_enough=just_enough,
|
||
)
|
||
for src in db_files
|
||
]
|
||
entries = [f.result() for f in futs]
|
||
# Sort by shard name so the manifest is stable across runs.
|
||
entries.sort(key=lambda e: e.name)
|
||
|
||
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} ({time.time() - t_all:.1f}s total)",
|
||
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
|