load.py

  1# SPDX-License-Identifier: GPL-2.0
  2
  3import re
  4import time
  5import json
  6
  7from lib.py import ksft_pr, cmd, ip, rand_port, wait_port_listen
  8
  9
 10class Iperf3Runner:
 11    """
 12    Sets up and runs iperf3 traffic.
 13    """
 14    def __init__(self, env, port=None, server_ip=None, client_ip=None):
 15        env.require_cmd("iperf3", local=True, remote=True)
 16        self.env = env
 17        self.port = rand_port() if port is None else port
 18        self.server_ip = server_ip
 19        self.client_ip = client_ip
 20
 21    def _build_server(self):
 22        cmdline = f"iperf3 -s -1 -p {self.port}"
 23        if self.server_ip:
 24            cmdline += f" -B {self.server_ip}"
 25        return cmdline
 26
 27    def _build_client(self, streams, duration, reverse):
 28        host = self.env.addr if self.server_ip is None else self.server_ip
 29        cmdline = f"iperf3 -c {host} -p {self.port} -P {streams} -t {duration} -J"
 30        if self.client_ip:
 31            cmdline += f" -B {self.client_ip}"
 32        if reverse:
 33            cmdline += " --reverse"
 34        return cmdline
 35
 36    def start_server(self):
 37        """
 38        Starts an iperf3 server with optional bind IP.
 39        """
 40        cmdline = self._build_server()
 41        proc = cmd(cmdline, background=True)
 42        wait_port_listen(self.port)
 43        time.sleep(0.1)
 44        return proc
 45
 46    def start_client(self, background=False, streams=1, duration=10, reverse=False):
 47        """
 48        Starts the iperf3 client with the configured options.
 49        """
 50        cmdline = self._build_client(streams, duration, reverse)
 51        return cmd(cmdline, background=background, host=self.env.remote)
 52
 53    def measure_bandwidth(self, reverse=False):
 54        """
 55        Runs an iperf3 measurement and returns the average bandwidth (Gbps).
 56        Discards the first and last few reporting intervals and uses only the
 57        middle part of the run where throughput is typically stable.
 58        """
 59        self.start_server()
 60        result = self.start_client(duration=10, reverse=reverse)
 61
 62        if result.ret != 0:
 63            raise RuntimeError("iperf3 failed to run successfully")
 64        try:
 65            out = json.loads(result.stdout)
 66        except json.JSONDecodeError as exc:
 67            raise ValueError("Failed to parse iperf3 JSON output") from exc
 68
 69        intervals = out.get("intervals", [])
 70        samples = [i["sum"]["bits_per_second"] / 1e9 for i in intervals]
 71        if len(samples) < 10:
 72            raise ValueError(f"iperf3 returned too few intervals: {len(samples)}")
 73        # Discard potentially unstable first and last 3 seconds.
 74        stable = samples[3:-3]
 75
 76        avg = sum(stable) / len(stable)
 77
 78        return avg
 79
 80
 81class GenerateTraffic:
 82    def __init__(self, env, port=None):
 83        self.env = env
 84        self.runner = Iperf3Runner(env, port)
 85
 86        self._iperf_server = self.runner.start_server()
 87        self._iperf_client = self.runner.start_client(background=True, streams=16, duration=86400)
 88
 89        # Wait for traffic to ramp up
 90        if not self._wait_pkts(pps=1000):
 91            self.stop(verbose=True)
 92            raise Exception("iperf3 traffic did not ramp up")
 93
 94    def _wait_pkts(self, pkt_cnt=None, pps=None):
 95        """
 96        Wait until we've seen pkt_cnt or until traffic ramps up to pps.
 97        Only one of pkt_cnt or pss can be specified.
 98        """
 99        pkt_start = ip("-s link show dev " + self.env.ifname, json=True)[0]["stats64"]["rx"]["packets"]
100        for _ in range(50):
101            time.sleep(0.1)
102            pkt_now = ip("-s link show dev " + self.env.ifname, json=True)[0]["stats64"]["rx"]["packets"]
103            if pps:
104                if pkt_now - pkt_start > pps / 10:
105                    return True
106                pkt_start = pkt_now
107            elif pkt_cnt:
108                if pkt_now - pkt_start > pkt_cnt:
109                    return True
110        return False
111
112    def wait_pkts_and_stop(self, pkt_cnt):
113        failed = not self._wait_pkts(pkt_cnt=pkt_cnt)
114        self.stop(verbose=failed)
115
116    def stop(self, verbose=None):
117        self._iperf_client.process(terminate=True)
118        if verbose:
119            ksft_pr(">> Client:")
120            ksft_pr(self._iperf_client.stdout)
121            ksft_pr(self._iperf_client.stderr)
122        self._iperf_server.process(terminate=True)
123        if verbose:
124            ksft_pr(">> Server:")
125            ksft_pr(self._iperf_server.stdout)
126            ksft_pr(self._iperf_server.stderr)
127        self._wait_client_stopped()
128
129    def _wait_client_stopped(self, sleep=0.005, timeout=5):
130        end = time.monotonic() + timeout
131
132        live_port_pattern = re.compile(fr":{self.runner.port:04X} 0[^6] ")
133
134        while time.monotonic() < end:
135            data = cmd("cat /proc/net/tcp*", host=self.env.remote).stdout
136            if not live_port_pattern.search(data):
137                return
138            time.sleep(sleep)
139        raise Exception(f"Waiting for client to stop timed out after {timeout}s")