# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Stampede with and without protection.

A fake database call that sleeps and counts its calls, behind a TTL cache.
200 threads ask for the same key the moment it expires: first unprotected,
then with per-key coalescing; then steady traffic with early probabilistic expiry.
"""
import argparse, json, math, random, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--threads", type=int, default=200, help="Threads that ask for the key at the moment it expires.")
p.add_argument("--recompute", type=float, default=2.0, help="Seconds one database call takes (W).")
p.add_argument("--ttl", type=float, default=120.0, help="TTL of the entry, in seconds, for the steady-traffic part.")
p.add_argument("--periods", type=int, default=3, help="TTL periods of steady traffic.")
p.add_argument("--clients", type=int, default=50, help="Steady-traffic clients, each asking every 10 ms.")
p.add_argument("--beta", type=float, default=1.0, help="beta of early probabilistic expiry.")
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()
random.seed(args.seed)

calls = 0
call_lock = threading.Lock()
def db():
    global calls
    with call_lock: calls += 1
    time.sleep(args.recompute); return "page"


def burst(get):
    """All threads ask at once for a key that has just expired; returns database calls."""
    global calls
    calls, cache = 0, {}
    start = threading.Barrier(args.threads)
    def worker():
        start.wait(); get("home", cache)
    ts = [threading.Thread(target=worker) for _ in range(args.threads)]
    for t in ts: t.start()
    for t in ts: t.join()
    return calls


# 1. No protection: every thread that misses calls the database
def get_plain(key, cache):
    v = cache.get(key)
    if v: return v
    cache[key] = db(); return cache[key]

# 2. Per-key coalescing: one thread computes, the others wait on an event
lock, events = threading.Lock(), {}
def get_coalesced(key, cache):
    v = cache.get(key)
    if v: return v
    with lock:
        ev = events.get(key)
        leader = ev is None
        if leader: events[key] = ev = threading.Event()
    if leader:
        cache[key] = db(); ev.set()
        with lock: events.pop(key)
    else:
        ev.wait()
    return cache[key]

plain = burst(get_plain)
coalesced = burst(get_coalesced)
print(f"{args.threads} threads at expiry, no protection: {plain} database calls "
      f"(L = lambda_k * W: every request in the window recomputes)")
print(f"{args.threads} threads at expiry, coalesced:     {coalesced} database call(s)")

# 3. Early probabilistic expiry under steady traffic: a request that draws
#    now - delta*beta*ln(rand) >= expiry refreshes in the background; everyone is served the current value
calls = 0
entry = {"value": db(), "expiry": time.time() + args.ttl, "refreshing": False}
delta = args.recompute                                    # the measured recompute time
refresh_log, waits, stop = [], [], threading.Event()
entry_lock = threading.Lock()
t0 = time.time()

def refresh():
    v = db()
    with entry_lock:
        entry.update(value=v, expiry=time.time() + args.ttl, refreshing=False)

def client():
    while not stop.is_set():
        t = time.time()
        with entry_lock:
            early = t - delta * args.beta * math.log(random.random()) >= entry["expiry"]
            if early and not entry["refreshing"]:
                entry["refreshing"] = True
                refresh_log.append(t - t0)
                threading.Thread(target=refresh).start()
            _ = entry["value"]
        waits.append(time.time() - t)
        time.sleep(0.01)

calls = 0
cs = [threading.Thread(target=client) for _ in range(args.clients)]
for c in cs: c.start()
time.sleep(args.ttl * args.periods)
stop.set()
for c in cs: c.join()

per_period = [sum(1 for r in refresh_log if i * args.ttl <= r < (i + 1) * args.ttl) for i in range(args.periods)]
rate = len(waits) / (args.ttl * args.periods)
lead = delta * args.beta * math.log(rate * delta * args.beta)     # where lambda * delta * e^(-x/delta) = 1
print(f"steady traffic {rate:,.0f} requests/s: early expiry should fire about delta*beta*ln(lambda*delta) = {lead:.1f} s "
      f"before expiry, so about {args.ttl / (args.ttl - lead + delta):.2f} recomputes per {args.ttl:g} s period")
print(f"early expiry, {args.periods} TTL periods of {args.ttl:g} s, {args.clients} clients: recomputes per period {per_period} "
      f"(mean {sum(per_period) / len(per_period):.2f}); longest wait {max(waits) * 1000:.1f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {
            "periods": [{"period": i + 1, "recomputes": n} for i, n in enumerate(per_period)],
            "summary": [{"plain_calls": plain, "coalesced_calls": coalesced,
                         "mean_recomputes": sum(per_period) / len(per_period), "max_wait_ms": max(waits) * 1000,
                         "rate": rate, "lead_s": lead, "predicted": args.ttl / (args.ttl - lead + delta),
                         "threads": args.threads}]}}, f)