xref: /linux/tools/testing/selftests/drivers/net/lib/py/load.py (revision fcee7d82f27d6a8b1ddc5bbefda59b4e441e9bc0) 
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        kwargs = {"background": background, "host": self.env.remote}
52        if not background:
53            kwargs["timeout"] = duration + 5
54        return cmd(cmdline, **kwargs)
55
56    def measure_bandwidth(self, reverse=False):
57        """
58        Runs an iperf3 measurement and returns the average bandwidth (Gbps).
59        Discards the first and last few reporting intervals and uses only the
60        middle part of the run where throughput is typically stable.
61        """
62        self.start_server()
63        result = self.start_client(duration=10, reverse=reverse)
64
65        if result.ret != 0:
66            raise RuntimeError("iperf3 failed to run successfully")
67        try:
68            out = json.loads(result.stdout)
69        except json.JSONDecodeError as exc:
70            raise ValueError("Failed to parse iperf3 JSON output") from exc
71
72        intervals = out.get("intervals", [])
73        samples = [i["sum"]["bits_per_second"] / 1e9 for i in intervals]
74        if len(samples) < 10:
75            raise ValueError(f"iperf3 returned too few intervals: {len(samples)}")
76        # Discard potentially unstable first and last 3 seconds.
77        stable = samples[3:-3]
78
79        avg = sum(stable) / len(stable)
80
81        return avg
82
83
84class GenerateTraffic:
85    def __init__(self, env, port=None):
86        self.env = env
87        self.runner = Iperf3Runner(env, port)
88
89        self._iperf_server = self.runner.start_server()
90        self._iperf_client = self.runner.start_client(background=True, streams=16, duration=86400)
91
92        # Wait for traffic to ramp up
93        if not self._wait_pkts(pps=1000):
94            self.stop(verbose=True)
95            raise Exception("iperf3 traffic did not ramp up")
96
97    def _wait_pkts(self, pkt_cnt=None, pps=None):
98        """
99        Wait until we've seen pkt_cnt or until traffic ramps up to pps.
100        Only one of pkt_cnt or pss can be specified.
101        """
102        pkt_start = ip("-s link show dev " + self.env.ifname, json=True)[0]["stats64"]["rx"]["packets"]
103        for _ in range(50):
104            time.sleep(0.1)
105            pkt_now = ip("-s link show dev " + self.env.ifname, json=True)[0]["stats64"]["rx"]["packets"]
106            if pps:
107                if pkt_now - pkt_start > pps / 10:
108                    return True
109                pkt_start = pkt_now
110            elif pkt_cnt:
111                if pkt_now - pkt_start > pkt_cnt:
112                    return True
113        return False
114
115    def wait_pkts_and_stop(self, pkt_cnt):
116        failed = not self._wait_pkts(pkt_cnt=pkt_cnt)
117        self.stop(verbose=failed)
118
119    def stop(self, verbose=None):
120        self._iperf_client.process(terminate=True)
121        if verbose:
122            ksft_pr(">> Client:")
123            ksft_pr(self._iperf_client.stdout)
124            ksft_pr(self._iperf_client.stderr)
125        self._iperf_server.process(terminate=True)
126        if verbose:
127            ksft_pr(">> Server:")
128            ksft_pr(self._iperf_server.stdout)
129            ksft_pr(self._iperf_server.stderr)
130        self._wait_client_stopped()
131
132    def _wait_client_stopped(self, sleep=0.005, timeout=5):
133        end = time.monotonic() + timeout
134
135        live_port_pattern = re.compile(fr":{self.runner.port:04X} 0[^6] ")
136
137        while time.monotonic() < end:
138            data = cmd("cat /proc/net/tcp*", host=self.env.remote).stdout
139            if not live_port_pattern.search(data):
140                return
141            time.sleep(sleep)
142        raise Exception(f"Waiting for client to stop timed out after {timeout}s")
143