# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Reproduce incast on a laptop.

One receiver asks N senders (here: N connections to a server process) for a
32 KB response at the same instant, and waits for all of them, like a
scatter-gather query. N grows from 1 to 500; p50, p99 and p999 of the
per-response latency are recorded. Then the same with jittered sending: each
sender waits a random 0-2 ms before answering. On loopback there is no switch
buffer to overflow, so the drop-and-RTO behaviour of real incast may not
appear: the run records what loopback does.
"""
import argparse, asyncio, json, resource, subprocess, sys, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--senders", default="1,10,50,100,200,300,500", help="Values of N.")
p.add_argument("--rounds", type=int, default=40, help="Queries per N.")
p.add_argument("--kb", type=int, default=32, help="Response size (KB).")
p.add_argument("--jitter-ms", type=float, default=2.0, help="Maximum jitter in the jittered run (ms).")
p.add_argument("--port", type=int, default=58401, help="Local port.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 8192), hard))

SERVER = r"""
import asyncio, random, resource, sys
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE); resource.setrlimit(resource.RLIMIT_NOFILE, (min(hard, 8192), hard))
port, size, jitter = int(sys.argv[1]), int(sys.argv[2]) * 1024, float(sys.argv[3]) / 1000
body = b"x" * size
async def h(r, w):
    while await r.readline():
        if jitter: await asyncio.sleep(random.uniform(0, jitter))
        w.write(body); await w.drain()
async def main():
    s = await asyncio.start_server(h, "127.0.0.1", port, backlog=4096); print("ready", flush=True)
    async with s: await s.serve_forever()
asyncio.run(main())
"""

async def query_rounds(n):
    conns = [await asyncio.open_connection("127.0.0.1", args.port) for _ in range(n)]
    lat = []; size = args.kb * 1024
    for _ in range(args.rounds):
        t0 = time.perf_counter()
        for r, w in conns: w.write(b"go\n")
        async def read(r):
            await r.readexactly(size); return time.perf_counter() - t0
        lat += await asyncio.gather(*(read(r) for r, _ in conns))
    for _, w in conns: w.close()
    return lat

def pct(xs, q): xs = sorted(xs); return xs[min(len(xs) - 1, int(q * len(xs)))]
rows = []
for jitter in (0.0, args.jitter_ms):
    srv = subprocess.Popen([sys.executable, "-c", SERVER, str(args.port), str(args.kb), str(jitter)], stdout=subprocess.PIPE, text=True)
    srv.stdout.readline()
    try:
        for n in map(int, args.senders.split(",")):
            lat = asyncio.run(query_rounds(n))
            rows.append({"sending": "jittered" if jitter else "synchronized", "N": n, "p50_ms": pct(lat, .5) * 1000,
                         "p99_ms": pct(lat, .99) * 1000, "p999_ms": pct(lat, .999) * 1000, "max_ms": max(lat) * 1000})
    finally: srv.kill(); srv.wait()
for r in rows: print(f"{r['sending']:<13} N={r['N']:>4}  p50 {r['p50_ms']:8.2f} ms  p99 {r['p99_ms']:8.2f} ms  p999 {r['p999_ms']:8.2f} ms")

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