arborist/bench/chunk_fetch_speed.py

"""Measure per-chunk fetch latency for two cloud paths:

1. apsw HttpRangeVFS on the big shard .db — current live path
   (FtsSidecarShardClient falls back here today because blobs/<hash>
   isn't published).
2. Direct HTTP GET — proxied by GETting a known small bucket object
   (the slim FTS5 sidecar's first bytes) to capture real RTT +
   transfer time per chunk.

Output: per-chunk wallclock for both paths, plus the speedup ratio
the JUST_ENOUGH=1 blob-publish move would buy us.

Reads-only; never writes to the bucket. Uses real cloud manifest +
real chunks. ~30 s on first run (cold cache), ~5 s on warm.
"""
from __future__ import annotations

import argparse
import sqlite3
import statistics
import time
import urllib.request
from pathlib import Path


DEFAULT_MANIFEST = (
    "https://nyc3.digitaloceanspaces.com/arborist/clones/manifest-fts.json"
)
DEFAULT_N = 10


def _pick_leaf_hashes(sidecar_path: Path, n: int) -> list[str]:
    """Pull n random leaf_hashes from a local slim sidecar."""
    c = sqlite3.connect(f"file:{sidecar_path}?mode=ro", uri=True)
    rows = c.execute(
        "SELECT leaf_hash FROM chunks ORDER BY RANDOM() LIMIT ?",
        (n,),
    ).fetchall()
    c.close()
    return [r[0] for r in rows]


def _time_apsw_fetches(shard_url: str, leaf_hashes: list[str]) -> list[float]:
    """For each leaf_hash, time the FtsSidecarShardClient fallback path:
    apsw HttpRangeVFS on the big shard + SELECT content FROM chunks
    WHERE leaf_hash = ?.
    """
    from arborist.wallet.bucket import BucketClient, BucketEndpoint
    # Cache_bytes=64 MB matches the default in MultiShardSidecarCorpus.
    # First chunk pays connection setup + page-cache cold misses; later
    # chunks may benefit from cached pages — report per-chunk times so
    # the asymmetry is visible.
    client = BucketClient(
        BucketEndpoint(shard_url=shard_url, blob_base=""),
        cache_bytes=64 * 1024 * 1024,
    )
    times: list[float] = []
    try:
        for h in leaf_hashes:
            t0 = time.time()
            rows = list(client.conn.execute(
                "SELECT content FROM chunks WHERE leaf_hash = ? "
                "AND content IS NOT NULL LIMIT 1",
                (h,),
            ))
            elapsed = time.time() - t0
            times.append(elapsed)
            if not rows:
                # No content — leaf_hash missing in this shard. Skip
                # but keep timing the trip.
                continue
    finally:
        client.close()
    return times


def _time_blob_fetches(probe_url: str, n: int) -> list[float]:
    """Proxy direct-blob-fetch cost by GETting the same bucket object
    n times (cold connection each time so RTT + small-payload transfer
    dominate, same shape a per-chunk blob GET would have).

    Uses a small known object (sidecar manifest JSON or first 4 KB of
    a published file). We're measuring "open-HTTP-connection-to-bucket
    + transfer ~few-KB" — not the apsw multi-roundtrip cost.
    """
    times: list[float] = []
    for _ in range(n):
        req = urllib.request.Request(probe_url, headers={"Range": "bytes=0-2048"})
        t0 = time.time()
        with urllib.request.urlopen(req, timeout=30) as resp:
            _ = resp.read()
        elapsed = time.time() - t0
        times.append(elapsed)
    return times


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--manifest", default=DEFAULT_MANIFEST)
    p.add_argument("--shard-idx", type=int, default=0,
                   help="which manifest shard to probe (0..n-1)")
    p.add_argument("--n", type=int, default=DEFAULT_N,
                   help="number of chunks to fetch on each path")
    args = p.parse_args()

    import json
    with urllib.request.urlopen(args.manifest, timeout=30) as resp:
        m = json.loads(resp.read())
    sh = m["shards"][args.shard_idx]
    shard_url = sh["url"]
    fts_url = sh["fts_sidecar_url"]
    fts_basename = fts_url.rsplit("/", 1)[-1]   # 000.idx.db
    n = fts_basename.split(".")[0]              # 000
    local_sidecar = Path.home() / ".arborist" / "sidecar-fts" / f"{n}.idx.db"
    if not local_sidecar.exists():
        print(f"need local slim sidecar at {local_sidecar} to pick "
              f"known-good leaf_hashes; build via `make sidecar-build-fts-all` "
              f"or download from {fts_url}")
        return 2

    print(f"shard: {shard_url}")
    print(f"sample-leaf-hashes drawn from local: {local_sidecar}")
    print(f"n={args.n} per path")
    print()

    leaf_hashes = _pick_leaf_hashes(local_sidecar, args.n)

    print("=== apsw HttpRangeVFS path (FtsSidecarShardClient fallback today) ===")
    apsw_times = _time_apsw_fetches(shard_url, leaf_hashes)
    for i, t in enumerate(apsw_times, 1):
        print(f"  chunk {i:2d}:  {t*1000:7.1f} ms")
    print(f"  apsw median: {statistics.median(apsw_times)*1000:.1f} ms")
    print(f"  apsw  mean : {statistics.mean(apsw_times)*1000:.1f} ms")
    print(f"  apsw  total: {sum(apsw_times):.2f} s")
    print()

    print("=== direct blob GET path (proxied by 2KB Range GET on a same-bucket object) ===")
    blob_times = _time_blob_fetches(fts_url, args.n)
    for i, t in enumerate(blob_times, 1):
        print(f"  blob  {i:2d}:  {t*1000:7.1f} ms")
    print(f"  blob median: {statistics.median(blob_times)*1000:.1f} ms")
    print(f"  blob  mean : {statistics.mean(blob_times)*1000:.1f} ms")
    print(f"  blob  total: {sum(blob_times):.2f} s")
    print()

    apsw_med = statistics.median(apsw_times)
    blob_med = statistics.median(blob_times)
    speedup = apsw_med / blob_med if blob_med > 0 else float("inf")
    print(f"=== speedup ===")
    print(f"  per-chunk median: apsw {apsw_med*1000:.1f} ms / blob {blob_med*1000:.1f} ms = {speedup:.1f}×")
    return 0


if __name__ == "__main__":
    import sys
    sys.exit(main())