plot.py

#!/usr/bin/env python3
"""Render plots from a measure.sh result directory.

Builds one panel per measurement module whose files are present in the
run directory (paths, udp, dscp, series, flows, failover, session,
health) and writes summary.png. When called on a top-level smoketest run
it also folds in the panels from the health/ and fo_* sub-runs, so the
single summary.png is a complete overview of every module collected.
"""
import json
import re
import sys
from math import ceil
from pathlib import Path

import matplotlib.pyplot as plt


def parse_ping(path: Path) -> list[float]:
    pattern = re.compile(r"time=([\d.]+) ms")
    return [float(m.group(1)) for m in pattern.finditer(path.read_text())]


def parse_iperf(path: Path) -> float:
    data = json.loads(path.read_text())
    return data["end"]["sum_received"]["bits_per_second"] / 1e6  # Mbit/s


def panel_paths_rtt(ax, run: Path) -> None:
    for path in ("A", "B"):
        rtts = parse_ping(run / f"ping_{path}.txt")
        ax.plot(rtts, label=f"WAN-{path}")
    ax.set_xlabel("probe")
    ax.set_ylabel("RTT [ms]")
    ax.set_title("Per-WAN ping RTT")
    ax.legend()
    ax.grid(True, alpha=0.3)


def panel_paths_tp(ax, run: Path) -> None:
    tps = [parse_iperf(run / f"iperf_{p}.json") for p in ("A", "B")]
    ax.bar(["WAN-A", "WAN-B"], tps)
    ax.set_ylabel("Throughput [Mbit/s]")
    ax.set_title("Per-WAN iperf3 TCP")
    ax.grid(True, alpha=0.3, axis="y")


def panel_dscp(ax, run: Path) -> None:
    classes = [
        ("ping_svc_dscp10.txt", "DSCP 10"),
        ("ping_svc_dscp20.txt", "DSCP 20"),
        ("ping_svc_unmarked.txt", "unmarked"),
    ]
    series = [parse_ping(run / fname) for fname, _ in classes]
    ax.boxplot(series)
    ax.set_xticks(range(1, len(classes) + 1))
    ax.set_xticklabels([label for _, label in classes])
    ax.set_ylabel("RTT [ms]")
    ax.set_title("PBR: RTT by traffic class\n(identical destination IP)")
    ax.grid(True, alpha=0.3, axis="y")


def panel_series(ax, run: Path) -> None:
    rtts = parse_ping(run / "ping_svc_series.txt")
    ax.plot(range(1, len(rtts) + 1), rtts, "o-", markersize=3, linewidth=0.8)
    ax.set_xlabel("ICMP sequence")
    ax.set_ylabel("RTT [ms]")
    ax.set_title("Service-IP ping series\n(per-packet path alternation)")
    ax.grid(True, alpha=0.3)


def panel_flows(ax, run: Path) -> None:
    data = json.loads((run / "iperf_svc_parallel.json").read_text())
    streams = data["end"]["streams"]
    rtt_ms = [s["sender"]["mean_rtt"] / 1000 for s in streams]
    mbps = [s["receiver"]["bits_per_second"] / 1e6 for s in streams]
    total = data["end"]["sum_received"]["bits_per_second"] / 1e6
    ax.scatter(rtt_ms, mbps)
    ax.set_xlabel("stream mean RTT [ms]")
    ax.set_ylabel("stream goodput [Mbit/s]")
    ax.set_title(
        f"Parallel flows to service IP\n"
        f"{len(streams)} streams, aggregate {total:.1f} Mbit/s"
    )
    ax.set_xlim(left=0)
    ax.set_ylim(bottom=0)
    ax.grid(True, alpha=0.3)


def read_events(run: Path) -> dict:
    path = run / "failover_events.txt"
    if not path.exists():
        return {}
    return dict(
        line.split() for line in path.read_text().splitlines() if line.strip()
    )


def draw_event_lines(ax, events: dict) -> None:
    for key, label, color in (
        ("down_at_s", "WAN-A link down", "tab:red"),
        ("up_at_s", "WAN-A link up", "tab:green"),
    ):
        if key in events:
            ax.axvline(float(events[key]), color=color, ls="--", lw=1, label=label)


def panel_failover(ax, run: Path) -> None:
    events = read_events(run)
    interval = float(events.get("interval_s", 0.1))
    pts = [
        (int(m.group(1)) * interval, float(m.group(2)))
        for m in re.finditer(
            r"icmp_seq=(\d+).*time=([\d.]+) ms",
            (run / "ping_svc_failover.txt").read_text(),
        )
    ]
    ax.plot([t for t, _ in pts], [r for _, r in pts], ".", markersize=3)
    draw_event_lines(ax, events)
    ax.set_xlabel("time [s]")
    ax.set_ylabel("RTT [ms]")
    ax.set_title("Failover: service-IP RTT\nduring WAN-A outage")
    ax.legend()
    ax.grid(True, alpha=0.3)


def panel_session(ax, run: Path) -> None:
    events = read_events(run)
    outcome = ""
    result = run / "session_result.txt"
    if result.exists():
        outcome = (
            "session survived"
            if "survived=yes" in result.read_text()
            else "session died"
        )
    try:
        data = json.loads((run / "iperf_svc_session.json").read_text())
    except (json.JSONDecodeError, OSError):
        data = {}
    intervals = data.get("intervals", [])
    xs = [iv["sum"]["end"] for iv in intervals]
    ys = [iv["sum"]["bits_per_second"] / 1e6 for iv in intervals]
    if xs:
        ax.plot(xs, ys, drawstyle="steps-post")
    else:
        ax.text(0.5, 0.5, "no interval data", ha="center", va="center",
                transform=ax.transAxes)
    draw_event_lines(ax, events)
    ax.set_xlabel("time [s]")
    ax.set_ylabel("goodput [Mbit/s]")
    ax.set_title(f"Long-lived TCP across failover\n{outcome}")
    ax.set_ylim(bottom=0)
    ax.legend()
    ax.grid(True, alpha=0.3)


def panel_udp(ax, run: Path) -> None:
    labels, jitter, loss = [], [], []
    for path in ("A", "B"):
        data = json.loads((run / f"iperf_udp_{path}.json").read_text())
        labels.append(f"WAN-{path}")
        jitter.append(data["end"]["sum"]["jitter_ms"])
        loss.append(data["end"]["sum"]["lost_percent"])
    x = range(len(labels))
    ax.bar([i - 0.2 for i in x], jitter, width=0.4, label="jitter [ms]")
    ax.bar([i + 0.2 for i in x], loss, width=0.4, label="loss [%]")
    ax.set_xticks(list(x))
    ax.set_xticklabels(labels)
    ax.set_title("UDP stream: jitter and loss")
    ax.legend()
    ax.grid(True, alpha=0.3, axis="y")


def parse_health(path: Path) -> dict:
    """Group health_monitor.jsonl score samples by 'wan:protocol'."""
    series: dict[str, list[tuple[float, float]]] = {}
    for line in path.read_text().splitlines():
        if not line.strip():
            continue
        d = json.loads(line)
        # Actuation events share the file with the scores; skip them
        # here (pre-actuation logs carry no 'type' and are all scores).
        if d.get("type", "score") != "score":
            continue
        key = f"{d['wan']}:{d['protocol']}"
        series.setdefault(key, []).append((d["timestamp"], d["score"]))
    return series


def panel_health(ax, run: Path) -> None:
    series = parse_health(run / "health_monitor.jsonl")
    # One line per WAN using the tcp (connectivity) score; tcp and udp are
    # currently the same ICMP stand-in, so plotting both would overlap.
    t0 = min((pts[0][0] for pts in series.values() if pts), default=0.0)
    for key in sorted(series):
        wan, proto = key.split(":")
        if proto != "tcp":
            continue
        pts = series[key]
        ax.plot([t - t0 for t, _ in pts], [s for _, s in pts],
                label=f"WAN-{wan}")
    ax.set_xlabel("time [s]")
    ax.set_ylabel("health score")
    ax.set_title("Health monitor: per-WAN score\n(WAN-A blackholed mid-run)")
    ax.set_ylim(-0.05, 1.05)
    ax.legend()
    ax.grid(True, alpha=0.3)


PANELS = [
    ("ping_A.txt", panel_paths_rtt),
    ("iperf_A.json", panel_paths_tp),
    ("iperf_udp_A.json", panel_udp),
    ("ping_svc_dscp10.txt", panel_dscp),
    ("ping_svc_series.txt", panel_series),
    ("iperf_svc_parallel.json", panel_flows),
    ("ping_svc_failover.txt", panel_failover),
    ("iperf_svc_session.json", panel_session),
    ("health_monitor.jsonl", panel_health),
]


def panels_in(d: Path, prefix: str = "") -> list[tuple]:
    """(panel_fn, dir, prefix) for every module whose files exist in d."""
    return [(fn, d, prefix) for probe, fn in PANELS if (d / probe).exists()]


def main(rundir: str) -> None:
    run = Path(rundir)
    panels = panels_in(run)
    # Fold sub-runs into the top-level summary so one summary.png covers
    # every module: health/ keeps its own title, fo_* panels are tagged
    # with the variant (plain/nat/wg). Sub-runs still get their own
    # summary.png when plot.py is called on them directly.
    for sub in sorted(p for p in run.iterdir() if p.is_dir()):
        if sub.name == "health":
            panels += panels_in(sub)
        elif sub.name.startswith("fo_"):
            panels += panels_in(sub, prefix=sub.name.removeprefix("fo_"))
    if not panels:
        sys.exit(f"no known measurement files in {rundir}")

    ncols = min(3, len(panels))
    nrows = ceil(len(panels) / ncols)
    fig, axes = plt.subplots(
        nrows, ncols, figsize=(5 * ncols, 4.2 * nrows), squeeze=False
    )
    flat = list(axes.flat)
    for (fn, d, prefix), ax in zip(panels, flat):
        fn(ax, d)
        if prefix:
            ax.set_title(f"[{prefix}] {ax.get_title()}")
    for ax in flat[len(panels):]:
        ax.set_visible(False)

    fig.suptitle(run.name)
    fig.tight_layout(rect=(0, 0, 1, 0.98))  # leave room for the suptitle
    out = run / "summary.png"
    fig.savefig(out, dpi=120)
    print(f"wrote {out}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("usage: plot.py <results/run-dir>", file=sys.stderr)
        sys.exit(1)
    main(sys.argv[1])