# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Find the metastable band.

A server of capacity C = 2,000 requests/s where, once the queue is longer
than a threshold, every request costs k = 3 times the work (retries, cold
caches). Arrival rates from 500 to 1,900/s, from a healthy start and after a
5-second burst of 6,000/s; then the recovery procedure: shed load below C/k,
wait, ramp back up.
"""
import argparse, json

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--capacity", type=float, default=2000, help="C: requests per second of work capacity.")
p.add_argument("--k", type=float, default=3, help="Cost multiplier once the queue is long.")
p.add_argument("--threshold", type=float, default=1000, help="Queue length that switches to the expensive mode.")
p.add_argument("--rates", type=lambda s: [float(v) for v in s.split(",")], default=[500, 1000, 1500, 1900], help="Arrival rates.")
p.add_argument("--seconds", type=int, default=400, help="Simulated seconds.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
C, k, DT = args.capacity, args.k, 0.01

def run(lam, burst, shed=False):
    q, series, state = 0.0, [], {}
    for i in range(int(args.seconds / DT)):
        t = i * DT
        a = 6000 if burst and 10 <= t < 15 else lam
        if shed and t >= 20:                                                  # shed well below C/k until the queue drains, then ramp
            if state.get("ramp_at") is None:
                if q < args.threshold / 2: state["ramp_at"] = t
                else: a = min(a, C / k * 0.5)
            if state.get("ramp_at") is not None and t < state["ramp_at"] + 30:
                a = min(a, C / k * 0.5 + (lam - C / k * 0.5) * (t - state["ramp_at"]) / 30)
        cost = k if q > args.threshold else 1.0
        served = min(q + a * DT, C / cost * DT)
        q = q + a * DT - served
        if i % 100 == 0: series.append({"t": round(t), "goodput": served / DT * (1 if cost == 1 else 1 / k), "queue": q})
    end = series[-1]
    return series, end["queue"] > args.threshold

rows, states = [], []
for lam in args.rates:
    for burst in (False, True):
        s, stuck = run(lam, burst)
        states.append({"rate": lam, "start": "after a burst" if burst else "healthy", "end_state": "degraded" if stuck else "healthy"})
        if lam == 1000: rows += [dict(r, run="after a burst" if burst else "healthy start") for r in s]
    print(f"lambda = {lam:>6.0f}: healthy start -> {states[-2]['end_state']:<9} after a burst -> {states[-1]['end_state']}")
s, stuck = run(1000, True, shed=True)
rows += [dict(r, run="after a burst, then shed and ramp") for r in s]
print(f"lambda = 1000 after a burst, with shed-and-ramp: {'degraded' if stuck else 'healthy'}   (band: {C / k:.0f} < lambda < {C:.0f})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"goodput": rows, "states": states, "summary": [{"band_low": C / k, "band_high": C,
                   "shed_recovers": "no" if stuck else "yes"}]}}, f)