"""Bulk throughput through a TCP port. server: bulk.py serve PORT; client: bulk.py PORT MiB RUNS"""
import socket, statistics, sys, time
CH = 1 << 20
def sink(s):
    buf = bytearray(CH); n = 0
    while (r := s.recv_into(buf)): n += r
    return n
def source(s, total):
    mv = memoryview(bytes(CH)); sent = 0
    while sent < total: s.sendall(mv); sent += CH
    s.shutdown(socket.SHUT_WR)
if sys.argv[1] == "serve":
    l = socket.socket(); l.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    l.bind(("0.0.0.0", int(sys.argv[2]))); l.listen()
    while True:
        s, _ = l.accept(); f = s.makefile("rb"); cmd = f.readline().split()
        if not cmd:              # readiness probe
            f.close(); s.close(); continue
        if cmd[0] == b"up":      # client sends, we count and reply with byte count
            buf = bytearray(CH); n = 0
            while (r := f.readinto(buf)): n += r
            s.sendall(str(n).encode()); f.close(); s.close()
        else:                    # down: we send N bytes
            source(s, int(cmd[1])); sink(s); f.close(); s.close()
else:
    port, mib, runs = int(sys.argv[1]), int(sys.argv[2]), int(sys.argv[3]); total = mib * CH
    for d in (sys.argv[4:] or ["up", "down"]):
        rates = []
        for _ in range(runs):
            s = socket.create_connection(("127.0.0.1", port)); t0 = time.perf_counter()
            s.sendall(f"{d} {total}\n".encode())
            if d == "up":
                source(s, total)
                got = b""
                while (c := s.recv(64)): got += c
                assert int(got) == total, got
            else:
                assert sink(s) == total
            rates.append(total / (time.perf_counter() - t0) / 1e9 * 8)
            s.close()
        print(f"{d:4}: median {statistics.median(rates):6.2f} Gbit/s  (" + " ".join(f"{r:.2f}" for r in rates) + ")", flush=True)