# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Retry storm: no jitter, equal jitter, full jitter.

400 clients against a server that serves at most 50 requests per second and
is down for the first 2 seconds; exponential backoff (base 0.1 s, cap 10 s)
with no jitter, equal jitter and full jitter. Completion time, the busiest
second, and total attempts.
"""
import argparse, heapq, json, random

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--clients", type=int, default=400, help="Clients.")
p.add_argument("--capacity", type=int, default=50, help="Requests served per second.")
p.add_argument("--down", type=float, default=2.0, help="Seconds the server is down at the start.")
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()

def run(jitter):
    random.seed(args.seed)
    events = [(random.uniform(0, 0.01), i, 0) for i in range(args.clients)]
    heapq.heapify(events)
    served, attempts, done_at, arrivals = {}, 0, 0.0, {}
    while events:
        t, c, n = heapq.heappop(events); attempts += 1
        sec = int(t); arrivals[sec] = arrivals.get(sec, 0) + 1
        if t < args.down or served.get(sec, 0) >= args.capacity:
            wait = min(10, 0.1 * 2 ** n)
            if jitter == "full": wait = random.uniform(0, wait)
            elif jitter == "equal": wait = wait / 2 + random.uniform(0, wait / 2)
            heapq.heappush(events, (t + wait, c, n + 1))
        else:
            served[sec] = served.get(sec, 0) + 1; done_at = max(done_at, t)
    return {"jitter": jitter, "all_done_s": done_at, "busiest_second": max(arrivals.values()), "attempts": attempts}

rows = [run(m) for m in ("none", "equal", "full")]
for r in rows: print(f"jitter {r['jitter']:<6} all served after {r['all_done_s']:6.1f} s   busiest second {r['busiest_second']:>4} arrivals   attempts {r['attempts']:,}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"modes": rows}}, f)