# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure leverage and the cost of a cold restart.

A Zipf-distributed read stream against SQLite, first bare, then behind an LRU
dictionary cache; then the cache is cleared mid-run and the database load is
followed second by second at a steady request rate.
"""
import argparse, json, os, sqlite3, tempfile
from collections import OrderedDict
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--rows", type=int, default=100_000, help="Rows in the table (and distinct keys).")
p.add_argument("--zipf", type=float, default=1.2, help="Zipf exponent of the key stream (numpy.random.zipf, modulo rows).")
p.add_argument("--capacity", type=int, default=5_000, help="Cache capacity, in entries (LRU).")
p.add_argument("--rate", type=int, default=2_000, help="Requests per second offered to the cache (the clock is the request count).")
p.add_argument("--warm-seconds", type=int, default=30, help="Seconds of traffic before the cache is cleared.")
p.add_argument("--after-seconds", type=int, default=60, help="Seconds of traffic followed after the clear.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
rng = np.random.default_rng(args.seed)

# 1. The database: SQLite in a temporary file, one row per key
path = os.path.join(tempfile.gettempdir(), "leverage-lab.db")
import atexit; atexit.register(lambda: os.path.exists(path) and os.remove(path))   # leave nothing behind
if os.path.exists(path): os.remove(path)
db = sqlite3.connect(path)
db.execute("CREATE TABLE kv (k INTEGER PRIMARY KEY, v TEXT)")
db.executemany("INSERT INTO kv VALUES (?, ?)", ((k, f"value-{k}") for k in range(args.rows)))
db.commit()
db_reads = 0
def read_db(k):
    global db_reads
    db_reads += 1
    return db.execute("SELECT v FROM kv WHERE k = ?", (k,)).fetchone()[0]

def keys(n):
    return (rng.zipf(args.zipf, n) % args.rows).tolist()

# 2. An LRU dictionary cache in front
cache = OrderedDict()
def get(k):
    if k in cache:
        cache.move_to_end(k); return cache[k]
    v = read_db(k)
    cache[k] = v
    if len(cache) > args.capacity: cache.popitem(last=False)
    return v

# 3. Steady traffic in 100 ms ticks: database reads per second, the cache cleared after warm-seconds
TICK = 10                                               # ticks per second
per_tick = args.rate // TICK
timeline = []
for tick in range((args.warm_seconds + args.after_seconds) * TICK):
    if tick == args.warm_seconds * TICK:
        cache.clear()
    before = db_reads
    for k in keys(per_tick): get(k)
    timeline.append({"second": (tick - args.warm_seconds * TICK) / TICK, "db_reads_per_s": (db_reads - before) * TICK})

warm = [t["db_reads_per_s"] for t in timeline if -10 <= t["second"] < 0]   # the last 10 s before the clear
steady_reads = sum(warm) / len(warm)
h = 1 - steady_reads / args.rate
leverage = 1 / (1 - h)
peak = max(t["db_reads_per_s"] for t in timeline if t["second"] >= 0)
recovered = next(t["second"] for t in timeline if t["second"] >= 0 and t["db_reads_per_s"] <= 1.5 * steady_reads)
print(f"hit rate h = {h:.3f}  ->  leverage 1/(1-h) = {leverage:.1f}")
print(f"database reads per 10,000 requests, warm: {steady_reads / args.rate * 10_000:,.0f} (without the cache: 10,000)")
print(f"after the clear: peak {peak:,.0f} reads/s against {steady_reads:,.0f} reads/s warm = {peak / steady_reads:.1f}x "
      f"(1/(1-h) = {leverage:.1f}); back under 1.5x warm after {recovered:.1f} s")
for t in [t for t in timeline if 0 <= t["second"] < 1]:
    print(f"  {t['second']:.1f} s after clear: {t['db_reads_per_s']:>6,.0f} reads/s")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"timeline": timeline, "summary": [{
            "h": h, "leverage": leverage, "steady_reads": steady_reads, "peak_reads": peak,
            "spike": peak / steady_reads, "rate": args.rate, "recovered_s": recovered}]}}, f)