"""Does listen.udp.workers help or hurt, on this platform, with this tracker?

Starts a tracker once per worker count and drives it with tests/udp_load.c,
which speaks the real protocol, so what comes out is the tracker's own
throughput rather than an echo server's.

    cc -O2 -pthread -o /tmp/udp_load tests/udp_load.c
    make opentracker
    python3 tests/udp_scaling.py

Worker count 0 means no dedicated workers: udp is served from the main event
loop, which is what a default build does.

LOAD_THREADS and WINDOW decide how much load is offered. Both matter. Too
little and the tracker idles while the figure measures round trips; too much
and the kernel drops datagrams, every lost reply costs the generator a
timeout, and the figure collapses.

Three columns say which side the number came from. "cores" is the tracker's
cpu, "load" the generator's. Whichever is near its machine's core count is
the one being measured. "lost" counts datagrams that went unanswered: once
it leaves zero the throughput figure is timeouts, not service, and the run
needs a smaller WINDOW.
"""
import os
import resource
import signal
import socket
import subprocess
import sys
import tempfile
import time

TRACKER = os.environ.get("OPENTRACKER", "./opentracker")
LOADER = os.environ.get("UDP_LOAD", "/tmp/udp_load")
WORKER_COUNTS = [int(n) for n in os.environ.get("WORKERS", "0 1 2 4 8").split()]
LOAD_THREADS = int(os.environ.get("LOAD_THREADS", "4"))
SECONDS = float(os.environ.get("SECONDS", "5"))
TORRENTS = int(os.environ.get("TORRENTS", "10000"))
WINDOW = int(os.environ.get("WINDOW", "32"))
REPEATS = int(os.environ.get("REPEATS", "3"))
PORT_BASE = int(os.environ.get("PORT_BASE", "19960"))


def start(port, workers):
    directory = tempfile.mkdtemp(prefix="ot-udpscale-")
    config = "listen.tcp 127.0.0.1:%d\nlisten.udp 127.0.0.1:%d\n" % (port, port)
    if workers:
        config += "listen.udp.workers %d\n" % workers
    with open(os.path.join(directory, "conf"), "w") as handle:
        handle.write(config)
    log = open(os.path.join(directory, "log"), "w+b")
    tracker = subprocess.Popen([TRACKER, "-f", os.path.join(directory, "conf")], stdout=log, stderr=log)
    for _ in range(300):
        if tracker.poll() is not None:
            log.seek(0)
            raise SystemExit("tracker exited at once on port %d with %d workers:\n%s"
                             % (port, workers, log.read().decode("utf-8", "replace")))
        try:
            socket.create_connection(("127.0.0.1", port), timeout=0.2).close()
            return tracker, log
        except OSError:
            time.sleep(0.02)
    tracker.kill()
    raise SystemExit("tracker never listened on %d" % port)


def cpu_seconds(pid):
    """Cumulative cpu time of a process, the one spelling ps has everywhere."""
    try:
        out = subprocess.run(["ps", "-o", "time=", "-p", str(pid)], capture_output=True, text=True).stdout.strip()
    except OSError:
        return None
    if not out:
        return None
    total, parts = 0.0, out.replace("-", ":").split(":")
    for part in parts:
        total = total * 60 + float(part)
    return total


def reaped_cpu_seconds():
    """Cpu of every child already waited for. The generator is the only one
       that starts and finishes inside a measured run."""
    spent = resource.getrusage(resource.RUSAGE_CHILDREN)
    return spent.ru_utime + spent.ru_stime


def stop(tracker):
    tracker.send_signal(signal.SIGINT)
    try:
        tracker.wait(timeout=10)
    except subprocess.TimeoutExpired:
        tracker.kill()
        tracker.wait()


def parse(line):
    """replies/s and the count of announces that were never answered"""
    words = line.split()
    if not words or not words[0].replace(".", "").isdigit():
        return None, None
    lost = 0
    if "lost," in words:
        lost = int(words[words.index("lost,") - 1])
    return float(words[0]), lost


def main():
    if not os.path.exists(LOADER):
        raise SystemExit("build the generator first: cc -O2 -pthread -o %s tests/udp_load.c" % LOADER)

    print("  tracker %s, %d load threads x %d in flight, %d torrents, %.0f s x %d runs"
          % (TRACKER, LOAD_THREADS, WINDOW, TORRENTS, SECONDS, REPEATS))
    print("  %-8s %12s %7s %7s %7s %10s   %s"
          % ("workers", "median/s", "spread", "cores", "load", "lost", "vs 0 workers"))
    baseline = None
    for index, workers in enumerate(WORKER_COUNTS):
        results = []
        for repeat in range(REPEATS):
            port = PORT_BASE + index * 16 + repeat
            tracker, log = start(port, workers)
            try:
                # Warm up until the torrent set exists and the peers are in
                # place, so the measured run is renewals rather than insertions
                subprocess.run([LOADER, "127.0.0.1", str(port), str(LOAD_THREADS), "2", str(TORRENTS), str(WINDOW)],
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                spent, offered, wall = cpu_seconds(tracker.pid), reaped_cpu_seconds(), time.time()
                result = subprocess.run(
                    [LOADER, "127.0.0.1", str(port), str(LOAD_THREADS), str(SECONDS), str(TORRENTS), str(WINDOW)],
                    capture_output=True, text=True,
                )
                used, drove, elapsed = cpu_seconds(tracker.pid), reaped_cpu_seconds(), max(time.time() - wall, 1e-9)
            finally:
                stop(tracker)
                log.close()
            rate, lost = parse(result.stdout.strip())
            if rate is None:
                continue
            cores = (used - spent) / elapsed if (used is not None and spent is not None) else None
            results.append((rate, cores, (drove - offered) / elapsed, lost))

        if not results:
            print("  %-8s no result" % workers)
            continue
        results.sort()
        rate, cores, load, lost = results[len(results) // 2]
        if baseline is None:
            baseline = rate
        # A spread wider than the differences between rows means the rows are
        # not telling you anything
        spread = 100.0 * (results[-1][0] - results[0][0]) / rate
        print("  %-8s %12.0f %6.0f%% %7s %7.2f %10d   %.2fx"
              % (workers, rate, spread, ("%.2f" % cores) if cores else "?", load, lost, rate / baseline))
        sys.stdout.flush()


if __name__ == "__main__":
    main()