#!/usr/bin/env python3
"""SQLite WAL checkpoint starvation lab.

What it shows
    In WAL mode a checkpoint may only copy frames that no open read
    transaction still needs (walCheckpoint() in src/wal.c caps mxSafeFrame at
    the oldest in-use read mark), and the writer only rewinds the WAL to
    frame 1 when no reader is using it (walRestartLog()). If read
    transactions always overlap, both conditions are never true at the same
    time, so the -wal file grows for as long as the overlap lasts. It also
    never shrinks by itself afterwards, because a reset rewinds the write
    position without truncating the file.

Scenarios (all fixed, see SCENARIOS below)
    a  writer only, default wal_autocheckpoint (1000 frames)
    b  writer + 2 readers with gaps: each holds a snapshot 0.4 s per 1.0 s,
       staggered by 0.5 s, so about 20% of the time no reader is open
    c  writer + 3 readers that always overlap: each holds 1.0 s, idles 0.2 s,
       staggered by 0.4 s, so at least 2 are open at every instant
    d  c + a maintenance thread running PRAGMA wal_checkpoint(TRUNCATE)
       every 5 s with a 5 s busy timeout
    e  c + PRAGMA journal_size_limit = 4 MiB on the writer
    f  b + PRAGMA journal_size_limit = 4 MiB (control for e: the same limit
       when resets can happen)
    g  c with wal_autocheckpoint=0 on the writer (control: how much of the
       writer latency under starvation comes from the autocheckpoint that
       now runs after every commit)

Common setup
    Fresh database per run, journal_mode=WAL, synchronous=NORMAL,
    page_size=4096, table t(id INTEGER PRIMARY KEY, ts REAL, b BLOB).
    One writer process commits one 1000-byte row per transaction, open loop
    at 2000 commits/s. Readers are separate processes (multiprocessing, so
    real multi-process locking and no shared GIL). Each run lasts 30 s; the
    first 5 s are discarded from the latency and slope statistics. Each
    scenario runs 3 times, interleaved (a..g, a..g, a..g).

Measurements
    every 100 ms: size of the database, -wal and -shm files, number of open
        read transactions, and the WAL header salts (to count resets)
    every 1 s (at 0.5, 1.5, 2.5 ... s, so never at the same instant as the
        maintenance thread's calls at 5, 10, 15 ... s):
        PRAGMA wal_checkpoint(PASSIVE) from a separate monitor
        connection, giving (busy, log, checkpointed). This probe is itself a
        checkpoint and runs in every scenario. Where the writer's
        autocheckpoint is on, it copies no more than autocheckpoint would.
        In scenario g (autocheckpoint off) it is the only checkpointer.
        Run with --probe-every 0 to turn it off.
    writer: per-commit service time (inside execute) and response time
        (completion minus scheduled start, so stalls count in full)
    readers: open/close times and the time of a 100-row lookup
    maintenance thread: every checkpoint's result tuple and duration

Usage
    python3 run.py                     # full matrix, about 12 minutes
    python3 run.py --only c,d --runs 1 --duration 10
    python3 run.py --workdir /some/dir # where the databases live (default:
                                       # a fresh directory under the OS temp dir)

Output (in ./results next to this file)
    env.json            machine, OS, Python, SQLite version and compile options
    runs.csv            one row of metrics per run
    summary.csv         median, min and max of every metric per scenario
    timeline.csv        100 ms samples of file sizes, open readers, probes
    writer_seconds.csv  per-second writer commits and latency percentiles
    latency_hist.csv    writer service and response time histograms
    readers.csv         every read transaction
    checkpoints.csv     every maintenance-thread checkpoint call
    results.json        env + scenario definitions + summary

Only the Python standard library is used. Database files are deleted after
each run; the WAL can reach hundreds of MB during scenario c. A run is cut
short if the WAL passes --max-wal-gb (default 20).
"""

import argparse
import csv
import json
import multiprocessing as mp
import os
import platform
import shutil
import sqlite3
import subprocess
import sys
import tempfile
import threading
import time
from array import array

import workload

HERE = os.path.dirname(os.path.abspath(__file__))
RESULTS = os.path.join(HERE, "results")
MIB = 1024 * 1024
PAGE = 4096
FRAME = PAGE + 24  # WAL frame = 24-byte frame header + one page
WAL_HDR = 32

COMMON = {"rate": 2000, "duration": 30.0, "warmup": 5.0}

SCENARIOS = {
    "a": {"name": "writer only", "readers": 0},
    "b": {"name": "readers with gaps", "readers": 2, "hold": 0.4, "idle": 0.6},
    "c": {"name": "overlapping readers", "readers": 3, "hold": 1.0, "idle": 0.2},
    "d": {"name": "overlapping + TRUNCATE every 5 s", "readers": 3, "hold": 1.0, "idle": 0.2,
          "maintenance": {"mode": "TRUNCATE", "period": 5.0, "busy_timeout_s": 5.0}},
    "e": {"name": "overlapping + journal_size_limit 4 MiB", "readers": 3, "hold": 1.0, "idle": 0.2,
          "journal_size_limit": 4 * MIB},
    "f": {"name": "gaps + journal_size_limit 4 MiB", "readers": 2, "hold": 0.4, "idle": 0.6,
          "journal_size_limit": 4 * MIB},
    "g": {"name": "overlapping, wal_autocheckpoint=0", "readers": 3, "hold": 1.0, "idle": 0.2,
          "wal_autocheckpoint": 0},
}


# ---------------------------------------------------------------- environment

def sh(*cmd):
    try:
        return subprocess.run(cmd, capture_output=True, text=True, timeout=20).stdout.strip()
    except Exception:
        return ""


def filesystem_of(path):
    df = sh("df", "-P", path).splitlines()
    dev = df[-1].split()[0] if len(df) > 1 else ""
    fstype = ""
    for line in sh("mount").splitlines():
        if line.startswith(dev + " "):
            fstype = line.split("(")[-1].split(",")[0].strip(")")
    info = {}
    for line in sh("diskutil", "info", dev).splitlines():
        if ":" in line:
            k, v = line.split(":", 1)
            if k.strip() in ("Solid State", "Device Location", "File System Personality", "Protocol"):
                info[k.strip()] = v.strip()
    return {"path": path, "device": dev, "type": fstype, **info}


def environment(workdir):
    con = sqlite3.connect(":memory:")
    opts = [r[0] for r in con.execute("PRAGMA compile_options")]
    con.close()
    return {
        "date_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "cpu": sh("sysctl", "-n", "machdep.cpu.brand_string"),
        "cpu_count": os.cpu_count(),
        "memory_bytes": int(sh("sysctl", "-n", "hw.memsize") or 0),
        "machine": platform.machine(),
        "os": sh("sw_vers").replace("\t", " ").splitlines(),
        "kernel": platform.release(),
        "python": sys.version,
        "sqlite_version": sqlite3.sqlite_version,
        "sqlite_compile_options": opts,
        "workdir_filesystem": filesystem_of(workdir),
    }


# ---------------------------------------------------------------- statistics

def pct(sorted_vals, p):
    if not sorted_vals:
        return None
    k = min(len(sorted_vals) - 1, max(0, int(round(p / 100 * (len(sorted_vals) - 1)))))
    return sorted_vals[k]


def median(vals):
    v = sorted(x for x in vals if x is not None)
    if not v:
        return None
    n = len(v)
    return v[n // 2] if n % 2 else (v[n // 2 - 1] + v[n // 2]) / 2


def slope(points):
    """Least-squares slope of (t, y)."""
    n = len(points)
    if n < 2:
        return None
    mt = sum(t for t, _ in points) / n
    my = sum(y for _, y in points) / n
    den = sum((t - mt) ** 2 for t, _ in points)
    return sum((t - mt) * (y - my) for t, y in points) / den if den else None


def coverage(intervals, lo, hi):
    """Fraction of [lo, hi] with zero open readers, and the minimum open count."""
    events = []
    for a, b in intervals:
        a, b = max(a, lo), min(b, hi)
        if b > a:
            events += [(a, 1), (b, -1)]
    events.sort()
    events.append((hi, 0))
    zero, t, n, min_open = 0.0, lo, 0, None
    for x, d in events:
        if x > t:  # segment [t, x) has n open readers
            if n == 0:
                zero += x - t
            min_open = n if min_open is None else min(min_open, n)
        n += d
        t = max(t, x)
    return zero / (hi - lo), min_open


HIST_EDGES = [round(10 ** (k / 10)) for k in range(0, 71)]  # 1 us .. 10 s, 10 buckets per decade


def histogram(vals):
    counts = [0] * (len(HIST_EDGES) + 1)
    for v in vals:
        lo, hi = 0, len(HIST_EDGES)
        while lo < hi:  # first edge >= v
            mid = (lo + hi) // 2
            if HIST_EDGES[mid] < v:
                lo = mid + 1
            else:
                hi = mid
        counts[lo] += 1
    return counts


# ---------------------------------------------------------------- one run

def remove_db(path):
    for suffix in ("", "-wal", "-shm", "-journal"):
        try:
            os.remove(path + suffix)
        except FileNotFoundError:
            pass


def create_db(path):
    remove_db(path)
    con = sqlite3.connect(path, isolation_level=None)
    con.execute(f"PRAGMA page_size={PAGE}")
    assert con.execute("PRAGMA journal_mode=WAL").fetchone()[0] == "wal"
    con.execute("CREATE TABLE t(id INTEGER PRIMARY KEY, ts REAL, b BLOB)")
    con.close()  # last connection closes: WAL is checkpointed and deleted


def run_once(key, run_no, workdir, args, out_rows):
    sc = {**COMMON, **SCENARIOS[key]}
    sc["duration"], sc["rate"] = args.duration, args.rate
    warmup = min(sc["warmup"], sc["duration"] / 3)
    path = os.path.join(workdir, f"lab_{key}_{run_no}.db")
    wal, shm = path + "-wal", path + "-shm"
    create_db(path)

    ctx = mp.get_context("spawn")
    q = ctx.Queue()
    stop = ctx.Event()
    open_count = ctx.Value("i", 0)
    t0 = time.time() + 2.0  # time for the spawned interpreters to start
    procs = [ctx.Process(target=workload.writer,
                         args=(path, t0, sc["duration"], sc["rate"], sc.get("journal_size_limit"),
                               sc.get("wal_autocheckpoint"), q))]
    n = sc["readers"]
    for i in range(n):
        period = sc["hold"] + sc["idle"]
        procs.append(ctx.Process(target=workload.reader,
                                 args=(path, t0, i, sc["hold"], sc["idle"], i * period / n, stop, open_count, q)))
    for p in procs:
        p.start()

    ckpt_rows = []
    maint = None
    if "maintenance" in sc:
        m = sc["maintenance"]
        maint = threading.Thread(target=workload.maintenance,
                                 args=(path, t0, sc["duration"], m["period"], m["mode"], stop, ckpt_rows))
        maint.start()

    monitor = sqlite3.connect(path, isolation_level=None, timeout=0.0)
    samples, resets, prev_salts, aborted = [], 0, None, False
    k = 0
    next_probe = 0.5
    workload.wait_until(t0)
    while True:
        t_due = k * 0.1
        if t_due > sc["duration"]:
            break
        workload.wait_until(t0 + t_due)
        t = time.time() - t0
        salts = workload.wal_salts(wal)
        if salts and prev_salts and salts != prev_salts:
            d = (salts[0] - prev_salts[0]) % 2**32
            resets += d if 1 <= d <= 1000 else 1  # random new salts (nCkpt==0 path) count as one
        if salts:
            prev_salts = salts
        wal_b = workload.file_size(wal)
        row = {"t": round(t, 3), "wal_bytes": wal_b, "shm_bytes": workload.file_size(shm),
               "db_bytes": workload.file_size(path), "readers_open": open_count.value,
               "resets_so_far": resets, "probe_busy": "", "probe_log": "", "probe_ckpt": "", "probe_ms": ""}
        if args.probe_every and t_due >= next_probe - 1e-9:
            next_probe += args.probe_every
            s = time.perf_counter()
            b, lg, ck = monitor.execute("PRAGMA wal_checkpoint(PASSIVE)").fetchone()
            row.update(probe_busy=b, probe_log=lg, probe_ckpt=ck,
                       probe_ms=round((time.perf_counter() - s) * 1e3, 3))
        samples.append(row)
        if wal_b > args.max_wal_gb * 1024**3:
            aborted = True
            break
        k += 1

    stop.set()
    got = {"writer": None, "reader": []}
    for _ in procs:
        kind, payload = q.get()
        if kind == "writer":
            got["writer"] = payload
        else:
            got["reader"].extend(payload)
    for p in procs:
        p.join()
    if maint:
        maint.join()

    # After everything stopped: one PASSIVE with no readers left.
    wal_before_final = workload.file_size(wal)
    final = monitor.execute("PRAGMA wal_checkpoint(PASSIVE)").fetchone()
    wal_after_final = workload.file_size(wal)
    monitor.close()
    remove_db(path)

    # ---- writer statistics
    w = got["writer"]
    start = array("d"); start.frombytes(w["start"])
    service = array("q"); service.frombytes(w["service"])
    response = array("q"); response.frombytes(w["response"])
    keep = [i for i in range(len(start)) if start[i] >= warmup]
    sv = sorted(service[i] for i in keep)
    rv = sorted(response[i] for i in keep)
    measured = sc["duration"] - warmup

    # per-second writer rows
    per_sec = {}
    for i in range(len(start)):
        per_sec.setdefault(int(start[i]), []).append(i)
    for sec, idxs in sorted(per_sec.items()):
        s_ = sorted(service[i] for i in idxs)
        r_ = sorted(response[i] for i in idxs)
        out_rows["writer_seconds"].append({
            "scenario": key, "run": run_no, "second": sec, "commits": len(idxs),
            "service_p50_us": pct(s_, 50), "service_p99_us": pct(s_, 99), "service_max_us": s_[-1],
            "response_p99_us": pct(r_, 99), "response_max_us": r_[-1]})
    for kind, vals in (("service", sv), ("response", rv)):
        for j, c in enumerate(histogram(vals)):
            if c:
                out_rows["latency_hist"].append({
                    "scenario": key, "run": run_no, "kind": kind,
                    "upper_us": HIST_EDGES[j] if j < len(HIST_EDGES) else "inf", "count": c})

    # ---- WAL statistics
    post = [s for s in samples if s["t"] >= warmup]
    wal_mb = [(s["t"], s["wal_bytes"] / MIB) for s in post]
    probes = [s for s in samples if s["probe_log"] != ""]
    lag = [s["probe_log"] - s["probe_ckpt"] for s in probes if s["probe_log"] >= 0]
    commits = len(start)

    # ---- readers
    rd = sorted(got["reader"], key=lambda r: r[1])
    zero_frac, min_open = coverage([(r[1], r[2]) for r in rd], warmup, sc["duration"])
    lookups = sorted(r[5] for r in rd if r[1] >= warmup)
    for r in rd:
        out_rows["readers"].append({"scenario": key, "run": run_no, "reader": r[0], "t_open": r[1],
                                    "t_close": r[2], "snapshot_max_id": r[3], "rows": r[4], "lookup_us": r[5]})

    for c in ckpt_rows:
        out_rows["checkpoints"].append({"scenario": key, "run": run_no, "t": c[0], "mode": c[1], "busy": c[2],
                                        "log": c[3], "checkpointed": c[4], "duration_ms": c[5]})
    ck_ms = sorted(c[5] for c in ckpt_rows)

    for s in samples:
        out_rows["timeline"].append({"scenario": key, "run": run_no, **s})

    frames_final = (samples[-1]["wal_bytes"] - WAL_HDR) / FRAME if samples[-1]["wal_bytes"] else 0
    metrics = {
        "scenario": key, "run": run_no, "aborted": int(aborted),
        "commits": commits, "target_commits": int(sc["rate"] * sc["duration"]),
        "achieved_rate": round(len(keep) / measured, 1), "write_errors": w["errors"],
        "service_p50_us": pct(sv, 50), "service_p99_us": pct(sv, 99), "service_p999_us": pct(sv, 99.9),
        "service_max_us": sv[-1] if sv else None,
        "response_p50_us": pct(rv, 50), "response_p99_us": pct(rv, 99), "response_max_us": rv[-1] if rv else None,
        "commits_over_100ms": sum(1 for v in rv if v > 100_000),
        "wal_final_mb": round(samples[-1]["wal_bytes"] / MIB, 2),
        "wal_high_water_mb": round(max(s["wal_bytes"] for s in samples) / MIB, 2),
        "wal_slope_mb_s": round(slope(wal_mb), 3) if slope(wal_mb) is not None else None,
        "shm_final_kb": round(samples[-1]["shm_bytes"] / 1024, 1),
        "wal_resets": resets,
        "frames_per_commit": round(frames_final / commits, 3) if resets == 0 and commits else None,
        "probe_count": len(probes), "probe_busy_count": sum(1 for s in probes if s["probe_busy"] == 1),
        "probe_last_log": probes[-1]["probe_log"] if probes else None,
        "probe_last_ckpt": probes[-1]["probe_ckpt"] if probes else None,
        "probe_max_lag_frames": max(lag) if lag else None,
        "final_passive": "/".join(map(str, final)),
        "wal_mb_before_final_passive": round(wal_before_final / MIB, 2),
        "wal_mb_after_final_passive": round(wal_after_final / MIB, 2),
        "zero_reader_fraction": round(zero_frac, 4), "min_readers_open": min_open,
        "read_txns": len(rd),
        "lookup_p50_us": pct(lookups, 50), "lookup_p99_us": pct(lookups, 99),
        "maint_calls": len(ckpt_rows), "maint_busy": sum(1 for c in ckpt_rows if c[2] == 1),
        "maint_full_reset": sum(1 for c in ckpt_rows if c[2] == 0 and c[3] == 0 and c[4] == 0),
        "maint_median_ms": median(ck_ms), "maint_max_ms": ck_ms[-1] if ck_ms else None,
        "writer_wal_autocheckpoint": w["settings"]["wal_autocheckpoint"],
        "writer_journal_size_limit": w["settings"]["journal_size_limit"],
        "writer_synchronous": w["settings"]["synchronous"],
    }
    out_rows["runs"].append(metrics)
    return metrics


# ---------------------------------------------------------------- output

def write_csv(name, rows):
    if not rows:
        return
    fields = list(rows[0].keys())
    for r in rows:
        for k in r:
            if k not in fields:
                fields.append(k)
    with open(os.path.join(RESULTS, name), "w", newline="") as f:
        wr = csv.DictWriter(f, fieldnames=fields)
        wr.writeheader()
        wr.writerows(rows)


def summarize(runs):
    out = []
    skip = {"scenario", "run", "final_passive"}
    for key in SCENARIOS:
        rs = [r for r in runs if r["scenario"] == key]
        if not rs:
            continue
        for m in rs[0]:
            if m in skip:
                continue
            vals = [r[m] for r in rs if isinstance(r[m], (int, float))]
            if not vals:
                continue
            out.append({"scenario": key, "name": SCENARIOS[key]["name"], "metric": m, "runs": len(vals),
                        "median": median(vals), "min": min(vals), "max": max(vals)})
    return out


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--only", default=",".join(SCENARIOS), help="comma-separated scenario keys")
    ap.add_argument("--runs", type=int, default=3)
    ap.add_argument("--duration", type=float, default=COMMON["duration"])
    ap.add_argument("--rate", type=float, default=COMMON["rate"])
    ap.add_argument("--probe-every", type=float, default=1.0, help="seconds between PASSIVE probes, 0 = off")
    ap.add_argument("--workdir", default=None)
    ap.add_argument("--max-wal-gb", type=float, default=20.0)
    args = ap.parse_args()

    keys = [k.strip() for k in args.only.split(",") if k.strip()]
    workdir = args.workdir or tempfile.mkdtemp(prefix="sqlite-wal-lab-")
    os.makedirs(workdir, exist_ok=True)
    os.makedirs(RESULTS, exist_ok=True)
    env = environment(workdir)
    env["parameters"] = {**COMMON, "duration": args.duration, "rate": args.rate, "runs": args.runs,
                         "probe_every_s": args.probe_every, "sample_every_s": 0.1}
    with open(os.path.join(RESULTS, "env.json"), "w") as f:
        json.dump(env, f, indent=2)
    print(f"SQLite {env['sqlite_version']}, Python {platform.python_version()}, {env['cpu']}, "
          f"{' '.join(env['os'])}, workdir {workdir} ({env['workdir_filesystem'].get('type')})", flush=True)

    rows = {n: [] for n in ("runs", "timeline", "writer_seconds", "latency_hist", "readers", "checkpoints")}
    t_start = time.time()
    for run_no in range(1, args.runs + 1):
        for key in keys:
            m = run_once(key, run_no, workdir, args, rows)
            print(f"[{time.time() - t_start:6.0f}s] {key} run {run_no}: wal {m['wal_final_mb']} MB "
                  f"(hw {m['wal_high_water_mb']}, {m['wal_slope_mb_s']} MB/s), resets {m['wal_resets']}, "
                  f"p50/p99 {m['service_p50_us']}/{m['service_p99_us']} us, resp max {m['response_max_us']} us, "
                  f"zero-reader {m['zero_reader_fraction']}, maint {m['maint_full_reset']}/{m['maint_calls']}",
                  flush=True)

    summary = summarize(rows["runs"])
    for name, r in rows.items():
        write_csv(f"{name}.csv", r)
    write_csv("summary.csv", summary)
    with open(os.path.join(RESULTS, "results.json"), "w") as f:
        json.dump({"env": env, "scenarios": SCENARIOS, "summary": summary}, f, indent=2)
    if not args.workdir:
        shutil.rmtree(workdir, ignore_errors=True)
    print(f"done in {time.time() - t_start:.0f}s, results in {RESULTS}", flush=True)


if __name__ == "__main__":
    main()