# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Nagle plus delayed ACK: reproduce the 40 ms stall.

A client sends a request as two small writes (header, then body) and waits
for a short reply. Round-trip time with Nagle on (the default), with
TCP_NODELAY, and with header and body in one write. Results depend on the
operating system's delayed-ACK behaviour; report what you see.
"""
import argparse, json, platform, socket, statistics, threading, time

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--calls", type=int, default=200, help="Calls per variant.")
p.add_argument("--host", default="127.0.0.1", help="Server address (another machine shows the effect more reliably).")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

def server():
    srv = socket.socket(); srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("0.0.0.0", 0)); srv.listen(4)
    def loop():
        while True:
            conn, _ = srv.accept()
            def handle(c):
                while True:
                    data = b""
                    while len(data) < 104:                       # 4-byte header + 100-byte body
                        chunk = c.recv(4096)
                        if not chunk: return
                        data += chunk
                    c.sendall(b"ok")
            threading.Thread(target=handle, args=(conn,), daemon=True).start()
    threading.Thread(target=loop, daemon=True).start()
    return srv.getsockname()[1]

port = server()
rows = []
for name, nodelay, one_write in (("Nagle on, two writes", False, False), ("TCP_NODELAY, two writes", True, False),
                                 ("Nagle on, one write", False, True)):
    c = socket.create_connection((args.host, port))
    if nodelay: c.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    rtts = []
    for _ in range(args.calls):
        t = time.perf_counter()
        if one_write: c.sendall(b"HDR!" + b"x" * 100)
        else: c.sendall(b"HDR!"); c.sendall(b"x" * 100)
        c.recv(2); rtts.append((time.perf_counter() - t) * 1000)
    c.close()
    rows.append({"variant": name, "median_ms": statistics.median(rtts), "p99_ms": sorted(rtts)[int(len(rtts) * 0.99) - 1]})
    print(f"{name:<26} median {rows[-1]['median_ms']:7.3f} ms   p99 {rows[-1]['p99_ms']:7.3f} ms")
print(f"OS: {platform.system()} {platform.release()}")

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