# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Closed loop against open loop, and coordinated omission.

A small server (its own process) handles each request in 40 ms with a fixed
number of workers, and slows to 400 ms halfway through. A closed-loop client
(a fixed number of users, each waiting for its reply) and an open-loop client
(a fixed arrival rate) test it. Then the server freezes once for one second
while a rate-limited client measures latency two ways: from the actual send
(no correction, like wrk) and from the intended send time (like wrk2).
"""
import argparse, asyncio, json, os, subprocess, sys, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--users", type=int, default=20, help="Closed-loop users (and server workers).")
p.add_argument("--fast-ms", type=float, default=40, help="Service time before the slowdown.")
p.add_argument("--slow-ms", type=float, default=400, help="Service time after the slowdown.")
p.add_argument("--rate", type=float, default=250, help="Open-loop arrival rate (requests/s).")
p.add_argument("--phase", type=float, default=6, help="Seconds before and after the slowdown.")
p.add_argument("--freeze-rate", type=float, default=1000, help="Request rate during the freeze test.")
p.add_argument("--port", type=int, default=58123, help="Local port for the server.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

SERVER = r"""
import asyncio, sys, time
port, workers, fast, slow, switch_after, freeze_at = int(sys.argv[1]), int(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4]), float(sys.argv[5]), float(sys.argv[6])
t0 = None; sem = asyncio.Semaphore(workers); froze = False
async def handle(rid, w):
    global froze
    async with sem:                                   # a fixed pool of workers
        el = time.monotonic() - t0
        if freeze_at >= 0 and el >= freeze_at and not froze:
            froze = True; time.sleep(1.0)             # one blocking one-second freeze
        await asyncio.sleep((fast if switch_after < 0 or el < switch_after else slow) / 1000)
    w.write(rid + b"\n")
async def conn(r, w):
    global t0
    while line := await r.readline():
        if t0 is None: t0 = time.monotonic()
        asyncio.ensure_future(handle(line.strip(), w))
async def main():
    s = await asyncio.start_server(conn, "127.0.0.1", port); print("ready", flush=True)
    async with s: await s.serve_forever()
asyncio.run(main())
"""

def server(workers, fast, slow, switch_after, freeze_at=-1):
    pr = subprocess.Popen([sys.executable, "-c", SERVER, str(args.port), str(workers), str(fast), str(slow), str(switch_after), str(freeze_at)],
                          stdout=subprocess.PIPE, text=True)
    pr.stdout.readline(); return pr

class Client:
    async def open(self):
        self.r, self.w = await asyncio.open_connection("127.0.0.1", args.port); self.waiting, self.n = {}, 0
        self.task = asyncio.ensure_future(self.read())
    async def read(self):
        while line := await self.r.readline():
            self.waiting.pop(line.strip()).set_result(time.monotonic())
    def send(self):
        self.n += 1; rid = str(self.n).encode(); fut = asyncio.get_running_loop().create_future()
        self.waiting[rid] = fut; self.w.write(rid + b"\n"); return fut

def pct(xs, q):
    xs = sorted(xs); return xs[min(len(xs) - 1, int(q * len(xs)))] if xs else None

async def closed():
    c = Client(); await c.open(); t0 = time.monotonic(); end = t0 + 2 * args.phase; done = []
    async def user():
        while time.monotonic() < end:
            s = time.monotonic(); r = await c.send(); done.append((s - t0, r - s))
    await asyncio.gather(*(user() for _ in range(args.users))); return done

async def opened():
    c = Client(); await c.open(); t0 = time.monotonic(); done = []; futs = []
    n = int(2 * args.phase * args.rate)
    async def one(intended, fut): r = await fut; done.append((intended - t0, r - intended))
    for i in range(n):
        intended = t0 + i / args.rate
        await asyncio.sleep(max(0, intended - time.monotonic()))
        futs.append(asyncio.ensure_future(one(intended, c.send())))
    deadline = time.monotonic() + 30
    while len(done) < n and time.monotonic() < deadline: await asyncio.sleep(0.05)   # let the queue drain (bounded)
    return done

async def freeze():
    c = Client(); await c.open(); t0 = time.monotonic(); raw, corrected = [], []
    for i in range(int(2 * args.phase * args.freeze_rate)):
        intended = t0 + i / args.freeze_rate
        await asyncio.sleep(max(0, intended - time.monotonic()))
        s = time.monotonic(); r = await c.send()            # one request in flight, like a single wrk connection
        raw.append(r - s); corrected.append(r - intended)
    return raw, corrected

def per_second(done, value):
    out = {}
    for t, v in done: out.setdefault(int(t), []).append(v)
    return out

rows = []
srv = server(args.users, args.fast_ms, args.slow_ms, args.phase)
try: cl = asyncio.run(closed())
finally: srv.kill(); srv.wait()
srv = server(args.users, args.fast_ms, args.slow_ms, args.phase)
try: op = asyncio.run(opened())
finally: srv.kill(); srv.wait()
cps, ops = per_second(cl, 0), per_second(op, 0)
for sec in range(int(2 * args.phase)):
    rows.append({"second": sec, "closed_rate": len(cps.get(sec, [])),
                 "closed_p99_ms": (pct([v for _, v in cl if int(_) == sec], .99) or 0) * 1000,
                 "open_p99_ms": (pct([v for _, v in op if int(_) == sec], .99) or 0) * 1000})
little = [{"phase": f"service time {ms:g} ms", "predicted_rate": args.users / (ms / 1000),
           "measured_rate": sum(len(cps.get(s, [])) for s in rng) / len(rng)}
          for ms, rng in ((args.fast_ms, range(1, int(args.phase))), (args.slow_ms, range(int(args.phase) + 1, int(2 * args.phase))))]
summary = {"closed_p99_ms": pct([v for _, v in cl], .99) * 1000, "open_p99_ms": pct([v for _, v in op], .99) * 1000,
           "open_max_ms": max(v for _, v in op) * 1000, "open_completed": len(op), "open_sent": int(2 * args.phase * args.rate)}

srv = server(1000, 0, 0, -1, freeze_at=args.phase)                 # no service time: only the freeze
try: raw, corr = asyncio.run(freeze())
finally: srv.kill(); srv.wait()
fz = [{"measurement": m, "samples": len(xs), "slow_samples": sum(x > 0.1 for x in xs),
       "p99_ms": pct(xs, .99) * 1000, "p999_ms": pct(xs, .999) * 1000, "max_ms": max(xs) * 1000}
      for m, xs in (("from actual send (uncorrected)", raw), ("from intended send (corrected)", corr))]

for r in little: print(f"closed loop, {r['phase']}: Little's Law {r['predicted_rate']:.0f}/s, measured {r['measured_rate']:.0f}/s")
print(f"p99 over the whole run: closed loop {summary['closed_p99_ms']:.0f} ms, open loop {summary['open_p99_ms']:.0f} ms (max {summary['open_max_ms']:.0f} ms)")
for r in fz: print(f"freeze, {r['measurement']:<31} {r['slow_samples']:>5} samples over 100 ms, p99 {r['p99_ms']:.1f} ms, max {r['max_ms']:.0f} ms")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"summary": summary, "tables": {"timeline": rows, "little": little, "freeze": fz}}, f)