# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Measure retry amplification and a retry budget.

A backend that handles 2,000 requests per second stalls for 3 seconds;
clients time out after 100 ms and retry up to r times. Offered load per
second without a budget, and with a budget that allows a retry only while
retries are below 10% of recent successes. Then r = 1, 2, 3.
"""
import argparse, collections, json

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--capacity", type=int, default=2000, help="Requests per second the backend finishes.")
p.add_argument("--base", type=int, default=1000, help="New requests per second.")
p.add_argument("--seconds", type=int, default=30, help="Simulated seconds.")
p.add_argument("--retries", type=int, default=3, help="Retries per request.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
TICK = 0.01

def run(budget, retries):
    q, offered, log = collections.deque(), 0, []
    recent_ok = recent_retry = 0
    for t in range(int(args.seconds / TICK)):
        now = t * TICK
        sick = 5 <= now < 8
        cap = 2 if sick else args.capacity * TICK
        for _ in range(int(args.base * TICK)): q.append((now, 0)); offered += 1
        nq = collections.deque()
        for (t0, n) in q:
            if now - t0 > 0.1:
                if n < retries and (not budget or recent_retry < 0.1 * max(recent_ok, 1)):
                    nq.append((now, n + 1)); offered += 1; recent_retry += 1
            else: nq.append((t0, n))
        q = nq
        for _ in range(int(cap)):
            if q: q.popleft(); recent_ok += 1
        if t % 100 == 99:
            recent_ok *= 0.5; recent_retry *= 0.5
            log.append({"second": round(now), "offered": offered}); offered = 0
    return log

rows, peaks = [], []
for budget in (False, True):
    log = run(budget, args.retries)
    label = "retry budget (10%)" if budget else "no budget"
    for r in log: rows.append({**r, "policy": label})
    peak = max(r["offered"] for r in log) / args.base
    peaks.append({"policy": label, "retries": args.retries, "peak_ratio": peak, "formula": (1 + args.retries) if not budget else 1.1})
    print(f"{label:<20} peak offered load {peak:.2f}x base")
for r in (1, 2, 3):
    peak = max(x["offered"] for x in run(False, r)) / args.base
    peaks.append({"policy": "no budget", "retries": r, "peak_ratio": peak, "formula": 1 + r})
    print(f"no budget, r = {r}: peak {peak:.2f}x   (1 + r)^1 = {1 + r}")

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