# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Fan-out tail and the power of two choices.

(1) How often a request that fans out to N shards contains a p99 event.
(2) Ten clients route 10,000 requests to 20 servers from a load snapshot
refreshed every k requests: to the least-loaded server in the snapshot, to the
lesser of two random servers, or to one random server.
"""
import argparse, json, random
import numpy as np

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--fanouts", type=lambda s: [int(v) for v in s.split(",")], default=[1, 10, 40, 100], help="Shards per request (N).")
p.add_argument("--trials", type=int, default=100_000, help="Requests per N.")
p.add_argument("--servers", type=int, default=20, help="Servers.")
p.add_argument("--requests", type=int, default=10_000, help="Requests routed.")
p.add_argument("--intervals", type=lambda s: [int(v) for v in s.split(",")], default=[10, 100, 500, 2000],
               help="Requests between snapshot refreshes, comma-separated.")
p.add_argument("--utilization", type=float, default=0.9, help="Arrival rate over total service capacity (in-flight load model).")
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); random.seed(args.seed)

fan = []
p99 = float(np.exp(2.326))                                 # p99 of a standard log-normal
for n in args.fanouts:
    worst = np.exp(rng.standard_normal((args.trials, n))).max(axis=1)
    fan.append({"N": n, "measured_pct": float(np.mean(worst > p99) * 100), "formula_pct": (1 - 0.99 ** n) * 100})
    print(f"N = {n:>3}: request contains a p99 event {fan[-1]['measured_pct']:5.1f}%  (1 - 0.99^N = {fan[-1]['formula_pct']:5.1f}%)")

def route(policy, interval):
    """In-flight load: one request arrives per step; each server finishes one every servers/utilization steps."""
    load = [0] * args.servers; snap = load[:]; peak = 0
    for i in range(args.requests):
        if i % interval == 0: snap = load[:]
        if policy == "least loaded (snapshot)": s = min(range(args.servers), key=lambda j: snap[j])
        elif policy == "two random choices": a, b = random.sample(range(args.servers), 2); s = a if snap[a] <= snap[b] else b
        else: s = random.randrange(args.servers)
        load[s] += 1
        if i % max(1, round(args.servers / args.utilization)) == 0:     # every server finishes one in-flight request
            load = [max(0, x - 1) for x in load]
        peak = max(peak, max(load))
    return peak

lb = []
for interval in args.intervals:
    for policy in ("least loaded (snapshot)", "two random choices", "one random"):
        lb.append({"interval": interval, "policy": policy, "max_load": route(policy, interval)})
    print(f"snapshot every {interval:>5}: " + "   ".join(f"{r['policy']} {r['max_load']}" for r in lb[-3:]))

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"fanout": fan, "routing": lb}}, f)