#!/usr/bin/python3
###############################################################################
#                                                                             #
# fireperf - A network benchmarking tool                                      #
# Copyright (C) 2026 IPFire Development Team                                  #
#                                                                             #
# This program is free software: you can redistribute it and/or modify        #
# it under the terms of the GNU General Public License as published by        #
# the Free Software Foundation, either version 3 of the License, or           #
# (at your option) any later version.                                         #
#                                                                             #
# This program is distributed in the hope that it will be useful,             #
# but WITHOUT ANY WARRANTY; without even the implied warranty of              #
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
# GNU General Public License for more details.                                #
#                                                                             #
# You should have received a copy of the GNU General Public License           #
# along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
#                                                                             #
###############################################################################

"""
	Measures how fast fireperf is and how much CPU time it needs for that, compares
	it to iperf3 and iperf (if they are installed), and appends the results to a
	file, so that we can see what a change has done.

	Everything runs over loopback on this host, unless --remote is given, in which
	case another host sends to this one. See benchmarks/README.
"""

import argparse
import csv
import datetime
import os
import platform
import re
import shlex
import shutil
import signal
import socket
import statistics
import subprocess
import sys
import time

FIELDS = (
	"date", "commit", "host", "kernel", "cpu", "cores", "note",
	"test", "tool", "variant", "connections", "duration", "runs",
	"gbit_per_s", "sender_cpu_s_per_gib", "receiver_cpu_s_per_gib", "cpu_s_per_gib",
	"connections_per_s", "client_cpu_s", "server_cpu_s",
	"client_cpu_us_per_connection", "server_cpu_us_per_connection",
	"total_connections", "isolated",
	"remote", "remote_kernel", "remote_cpu", "remote_cores",
)

# Tests start at this port and every one of them uses the next
port = 15000 + (os.getpid() % 10000)

def next_port():
	global port

	port += 1
	return port

def run_pair(server, client, stop_server=True):
	"""
		Runs a server and a client and returns the CPU time that each of
		them has used (user and system) and what the client has printed
	"""
	# We read what is being printed, which must not be translated
	environ = dict(os.environ, LC_ALL="C")

	s = subprocess.Popen(server, stdin=subprocess.DEVNULL, env=environ,
		stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)

	# Give the server a moment to start
	time.sleep(1)

	c = subprocess.Popen(client, stdin=subprocess.DEVNULL, env=environ,
		stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)

	output = c.stdout.read()
	_, status, client_usage = os.wait4(c.pid, 0)

	if stop_server:
		time.sleep(0.2)
		s.send_signal(signal.SIGTERM)

	s.stdout.read()
	_, _, server_usage = os.wait4(s.pid, 0)

	if os.waitstatus_to_exitcode(status) != 0:
		raise RuntimeError("%s has failed:\n%s" % (" ".join(client), output))

	return (
		server_usage.ru_utime + server_usage.ru_stime,
		client_usage.ru_utime + client_usage.ru_stime,
		output,
	)

class Remote(object):
	"""
		Another host that we reach over SSH and that sends to us
	"""
	def __init__(self, host, directory):
		# This can be USER@HOST
		self.host = host
		self.name = host.rsplit("@", 1)[-1]
		self.directory = directory

		# Find the address of the other host, and ours as it sees it
		self.address, self.local_address = self.addresses(self.name)

	@staticmethod
	def addresses(host):
		for family, _, _, _, address in socket.getaddrinfo(host, None, type=socket.SOCK_STREAM):
			# Connecting a UDP socket sends nothing, but tells which address we would use
			s = socket.socket(family, socket.SOCK_DGRAM)

			try:
				s.connect((address[0], 9))
			except OSError:
				continue
			finally:
				local = s.getsockname()[0]
				s.close()

			return address[0], local

		raise RuntimeError("Could not find a route to %s" % host)

	def ssh(self, command, **kwargs):
		return subprocess.run(["ssh", "-o", "BatchMode=yes", self.host, command],
			stdin=subprocess.DEVNULL, text=True, **kwargs)

	def check_output(self, command):
		return self.ssh(command, stdout=subprocess.PIPE, check=True).stdout

	def has(self, command):
		return self.ssh("command -v %s" % shlex.quote(command),
			stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0

	def copy(self, path):
		"""
			Copies a file into our directory on the other host, which must allow
			to execute what is in it (home directories often don't)
		"""
		self.check_output("mkdir -p %s" % shlex.quote(self.directory))

		destination = "%s/%s" % (self.directory, os.path.basename(path))

		subprocess.check_call(["scp", "-q", "-o", "BatchMode=yes", path,
			"%s:%s" % (self.host, destination)])

		return destination

	def environment(self):
		output = self.check_output("uname -r; nproc; "
			"grep -m1 '^model name' /proc/cpuinfo | cut -d: -f2-")

		kernel, cores, cpu = (output.splitlines() + ["", "", ""])[:3]

		return {
			"remote"        : self.name,
			"remote_kernel" : kernel.strip(),
			"remote_cpu"    : cpu.strip(),
			"remote_cores"  : cores.strip(),
		}

	def start(self, command):
		"""
			Runs a command on the other host in the background
		"""
		# bash tells us the CPU time (user and system) of what it has launched
		command = "bash -c %s" % shlex.quote("%s </dev/null 2>&1; times" % shlex.join(command))

		return subprocess.Popen(["ssh", "-o", "BatchMode=yes", self.host, command],
			stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)

	@staticmethod
	def finish(process):
		"""
			Waits for something that has been started and returns its
			CPU time and what it has printed
		"""
		output = process.stdout.read()
		process.wait()

		# The last line is the CPU time of the children: "0m1.234s 0m0.567s"
		lines = output.splitlines()

		m = re.match(r"(\d+)m([0-9.]+)s (\d+)m([0-9.]+)s", lines[-1])
		if not m:
			raise RuntimeError("Could not find the CPU time in:\n%s" % output)

		cpu = int(m.group(1)) * 60 + float(m.group(2)) + int(m.group(3)) * 60 + float(m.group(4))

		return cpu, "\n".join(lines[:-2])

def run_remote(remote, server, client, server_is_remote):
	"""
		Runs a server and a client where one of them is on the other host
		and returns the CPU time of the remote and of the local one, and what
		the client has printed
	"""
	environ = dict(os.environ, LC_ALL="C")

	if server_is_remote:
		r = remote.start(server)
		time.sleep(2)

		l = subprocess.Popen(client, stdin=subprocess.DEVNULL, env=environ,
			stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)

		output = l.stdout.read()
		_, status, usage = os.wait4(l.pid, 0)

		if os.waitstatus_to_exitcode(status) != 0:
			raise RuntimeError("%s has failed:\n%s" % (" ".join(client), output))

		# The server terminates on its own (see --timeout)
		remote_cpu, _ = Remote.finish(r)

	else:
		l = subprocess.Popen(server, stdin=subprocess.DEVNULL, env=environ,
			stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

		time.sleep(1)

		remote_cpu, output = Remote.finish(remote.start(client))

		# Stop the server if it does not terminate on its own
		time.sleep(0.2)

		pid, _, usage = os.wait4(l.pid, os.WNOHANG)
		if not pid:
			l.send_signal(signal.SIGTERM)

			_, _, usage = os.wait4(l.pid, 0)

	return remote_cpu, usage.ru_utime + usage.ru_stime, output

def fireperf_bytes(output):
	m = re.search(r"^Transferred: (\d+) bytes sent, (\d+) bytes received", output, re.M)

	return int(m.group(1)) + int(m.group(2))

def fireperf_connections_per_second(output):
	# The summary is the line that follows the separator and its columns are
	# those of the header (as long as we stay left of the ones with a unit)
	lines = output.splitlines()

	header = [l for l in lines if l.split()[:2] == ["Time", "Open"]][0].split()
	summary = lines[[i for i, l in enumerate(lines) if l.startswith("- - ")][-1] + 1].split()

	return float(summary[header.index("New/s")])

def iperf_bytes(output, pattern):
	units = { "K" : 2**10, "M" : 2**20, "G" : 2**30, "T" : 2**40 }

	value, unit = re.findall(pattern, output, re.M)[-1]

	return float(value) * units[unit]

class Benchmark(object):
	def __init__(self, args):
		self.args = args
		self.rows = []

		self.fireperf = os.path.abspath(args.fireperf)
		self.iperf3 = shutil.which(os.environ.get("IPERF3", "iperf3"))
		self.iperf = shutil.which(os.environ.get("IPERF", "iperf"))

		self.environment = {
			"date"   : datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%d %H:%M"),
			"commit" : self.commit(),
			"host"   : platform.node(),
			"kernel" : platform.release(),
			"cpu"    : self.cpu(),
			"cores"  : os.cpu_count(),
			"note"   : args.note,
		}

		self.remote = None

		if args.remote:
			self.remote = Remote(args.remote, args.remote_directory)
			self.environment.update(self.remote.environment())

			# Copy fireperf over and run it from there
			self.remote_fireperf = self.remote.copy(self.fireperf)

			# The other host needs the same tools
			self.remote_iperf3 = self.iperf3 and self.remote.has("iperf3") and "iperf3"
			self.remote_iperf  = self.iperf  and self.remote.has("iperf") and "iperf"

	@staticmethod
	def commit():
		try:
			return subprocess.check_output(["git", "describe", "--always", "--dirty"],
				stderr=subprocess.DEVNULL, text=True).strip()
		except (OSError, subprocess.CalledProcessError):
			return "unknown"

	@staticmethod
	def cpu():
		try:
			with open("/proc/cpuinfo") as f:
				for line in f:
					if line.startswith("model name"):
						return line.split(":", 1)[1].strip()
		except OSError:
			pass

		return platform.machine()

	def version(self, command):
		try:
			output = subprocess.run([command, "--version"],
				stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True).stdout
		except OSError:
			return "?"

		m = re.search(r"\d+(\.\d+)+", output)
		if m:
			return m.group(0)

		return "?"

	def record(self, test, tool, variant, connections, **values):
		row = dict(self.environment, test=test, tool=tool, variant=variant,
			connections=connections, duration=self.args.duration, runs=self.args.runs)
		row.update(values)

		self.rows.append(row)

	# Throughput

	def throughput(self, connections):
		t = "%s" % self.args.duration
		P = "%s" % connections

		# Who sends is what costs the most. For fireperf and iperf3 (in reverse mode)
		# it is the server, for iperf it is the client.
		candidates = []

		for backend in ("epoll", "io_uring"):
			for zerocopy in ("auto", "never"):
				candidates.append((
					"fireperf", self.version(self.fireperf), "%s, zerocopy %s" % (backend, zerocopy),
					lambda p, b=backend, z=zerocopy: [self.fireperf, "-s", "-p", p, "-b", b, "--zerocopy", z],
					lambda p, b=backend, z=zerocopy: [self.fireperf, "-c", "127.0.0.1", "-p", p, "-b", b, "--zerocopy", z, "-t", t, "-P", P],
					fireperf_bytes, True,
				))

		if self.iperf3:
			pattern = r"\[ *\d+\].*?([0-9.]+) ([KMGT])Bytes.*receiver"
			if connections > 1:
				pattern = r"\[SUM\].*?([0-9.]+) ([KMGT])Bytes.*receiver"

			for variant, extra in (("default", []), ("zerocopy", ["-Z"])):
				candidates.append((
					"iperf3", self.version(self.iperf3), variant,
					lambda p: [self.iperf3, "-s", "-p", p, "-1"],
					lambda p, e=extra: [self.iperf3, "-c", "127.0.0.1", "-p", p, "-t", t, "-P", P, "-R"] + e,
					lambda output, pt=pattern: iperf_bytes(output, pt), True,
				))

		if self.iperf:
			pattern = r"\[ *\d+\].*?([0-9.]+) ([KMGT])Bytes"
			if connections > 1:
				pattern = r"\[SUM\].*?([0-9.]+) ([KMGT])Bytes"

			candidates.append((
				"iperf", self.version(self.iperf), "default",
				lambda p: [self.iperf, "-s", "-p", p],
				lambda p: [self.iperf, "-c", "127.0.0.1", "-p", p, "-t", t, "-P", P],
				lambda output, pt=pattern: iperf_bytes(output, pt), False,
			))

		print("\nThroughput: %d connection(s), %d seconds, median of %d run(s)" % (
			connections, self.args.duration, self.args.runs))
		print("%-10s %-8s %-26s %8s %14s %16s %14s" % ("", "", "", "Gbit/s",
			"sender s/GiB", "receiver s/GiB", "total s/GiB"))

		for tool, version, variant, server, client, parse, server_sends in candidates:
			results = []

			for _ in range(self.args.runs):
				p = "%s" % next_port()

				# iperf3 terminates on its own after one test
				server_cpu, client_cpu, output = run_pair(
					server(p), client(p), stop_server=(tool != "iperf3"))

				gib = parse(output) / 2**30

				sender_cpu, receiver_cpu = client_cpu, server_cpu
				if server_sends:
					sender_cpu, receiver_cpu = server_cpu, client_cpu

				results.append((gib, sender_cpu / gib, receiver_cpu / gib))

			gib      = statistics.median(r[0] for r in results)
			sender   = statistics.median(r[1] for r in results)
			receiver = statistics.median(r[2] for r in results)
			gbit     = gib * 2**30 * 8 / 1e9 / self.args.duration

			print("%-10s %-8s %-26s %8.1f %14.3f %16.3f %14.3f" % (
				tool, version, variant, gbit, sender, receiver, sender + receiver))

			self.record("throughput", "%s %s" % (tool, version), variant, connections,
				gbit_per_s="%.2f" % gbit, sender_cpu_s_per_gib="%.4f" % sender,
				receiver_cpu_s_per_gib="%.4f" % receiver, cpu_s_per_gib="%.4f" % (sender + receiver))

	def remote_throughput(self, connections):
		"""
			The other host sends to us. We only care about what it costs to send,
			so the receiving side is always the same.
		"""
		t = "%s" % self.args.duration
		P = "%s" % connections

		# The server terminates on its own a little after the client is done
		timeout = "%s" % (self.args.duration + 6)

		address = self.remote.address
		local   = self.remote.local_address

		candidates = []

		# The server sends, and it is on the other host
		for backend, zerocopy in (("epoll", "auto"), ("epoll", "never"),
				("io_uring", "auto"), ("io_uring", "always"), ("io_uring", "never")):
			candidates.append((
				"fireperf", self.version(self.fireperf), "%s, zerocopy %s" % (backend, zerocopy),
				lambda p, b=backend, z=zerocopy: [self.remote_fireperf, "-s", "-p", p, "-b", b,
					"--zerocopy", z, "-t", timeout, "-n"],
				lambda p: [self.fireperf, "-c", address, "-p", p, "-b", "epoll", "-t", t, "-P", P, "-n"],
				fireperf_bytes, True,
			))

		# The client sends, and it is on the other host
		if self.remote_iperf3:
			pattern = r"\[ *\d+\].*?([0-9.]+) ([KMGT])Bytes.*receiver"
			if connections > 1:
				pattern = r"\[SUM\].*?([0-9.]+) ([KMGT])Bytes.*receiver"

			for variant, extra in (("default", []), ("zerocopy", ["-Z"])):
				candidates.append((
					"iperf3", self.version(self.iperf3), variant,
					lambda p: [self.iperf3, "-s", "-p", p, "-1"],
					lambda p, e=extra: [self.remote_iperf3, "-c", local, "-p", p, "-t", t, "-P", P] + e,
					lambda output, pt=pattern: iperf_bytes(output, pt), False,
				))

		if self.remote_iperf:
			pattern = r"\[ *\d+\].*?([0-9.]+) ([KMGT])Bytes"
			if connections > 1:
				pattern = r"\[SUM\].*?([0-9.]+) ([KMGT])Bytes"

			candidates.append((
				"iperf", self.version(self.iperf), "default",
				lambda p: [self.iperf, "-s", "-p", p],
				lambda p: [self.remote_iperf, "-c", local, "-p", p, "-t", t, "-P", P],
				lambda output, pt=pattern: iperf_bytes(output, pt), False,
			))

		print("\nThroughput from %s: %d connection(s), %d seconds, median of %d run(s)" % (
			self.remote.name, connections, self.args.duration, self.args.runs))
		print("%-10s %-8s %-26s %8s %14s %16s %14s" % ("", "", "", "Gbit/s",
			"sender s/GiB", "receiver s/GiB", "total s/GiB"))

		for tool, version, variant, server, client, parse, server_is_remote in candidates:
			results = []

			for _ in range(self.args.runs):
				p = "%s" % next_port()

				sender_cpu, receiver_cpu, output = run_remote(self.remote,
					server(p), client(p), server_is_remote)

				gib = parse(output) / 2**30

				results.append((gib, sender_cpu / gib, receiver_cpu / gib))

			gib      = statistics.median(r[0] for r in results)
			sender   = statistics.median(r[1] for r in results)
			receiver = statistics.median(r[2] for r in results)
			gbit     = gib * 2**30 * 8 / 1e9 / self.args.duration

			print("%-10s %-8s %-26s %8.1f %14.3f %16.3f %14.3f" % (
				tool, version, variant, gbit, sender, receiver, sender + receiver))

			self.record("throughput", "%s %s" % (tool, version), variant, connections,
				gbit_per_s="%.2f" % gbit, sender_cpu_s_per_gib="%.4f" % sender,
				receiver_cpu_s_per_gib="%.4f" % receiver, cpu_s_per_gib="%.4f" % (sender + receiver))

	# Connections

	def isolate(self):
		"""
			Moves us (and everything that we launch from now on) into a network
			namespace of our own.

			Every connection stays in the connection tracking table of a host for
			a couple of minutes after it has been closed. Opening as many as it
			takes to measure anything would fill up that table and break the network
			for everything else on this host. A new namespace has a table of its own,
			and does not track connections at all as long as nobody asks it to.
		"""
		try:
			os.unshare(os.CLONE_NEWNET)
		except (AttributeError, OSError) as e:
			print("Could not create a network namespace: %s" % e)
			return False

		# The loopback interface of a new namespace is down
		try:
			subprocess.check_call(["ip", "link", "set", "lo", "up"])
		except (OSError, subprocess.CalledProcessError) as e:
			print("Could not bring up the loopback interface: %s" % e)
			sys.exit(1)

		return True

	def connections(self):
		"""
			How fast can we open connections? Only fireperf can do this.
		"""
		connections = self.args.connections
		isolated = self.isolate()

		# Without a namespace of our own, we have to be careful
		if not isolated:
			connections = min(connections, 10000)

			print("Only opening %d connections per test to not fill up the connection "
				"tracking table of this host. Run this as root for better results." % connections)

		n = "%s" % connections

		tests = (
			("handshake",              ["--close"]),
			("exchange 100:1000",      ["--exchange", "100:1000"]),
		)

		print("\nConnections: %s connections with 16 at a time, median of %d run(s)" % (n, self.args.runs))
		print("%-10s %-8s %-26s %14s %16s %16s" % ("", "", "", "connections/s",
			"client CPU us/c", "server CPU us/c"))

		version = self.version(self.fireperf)

		for test, extra in tests:
			for backend in ("epoll", "io_uring"):
				results = []

				for _ in range(self.args.runs):
					# A client can only open so many connections to one port before it
					# runs out of ports of its own, so the server listens on a couple
					first = next_port()

					for _ in range(15):
						last = next_port()

					p = "%s:%s" % (first, last)

					server_cpu, client_cpu, output = run_pair(
						[self.fireperf, "-s", "-p", p, "-b", backend],
						[self.fireperf, "-c", "127.0.0.1", "-p", p, "-b", backend,
							"-P", "16", "-C", n, "-t", "300"] + extra,
					)

					results.append((fireperf_connections_per_second(output),
						client_cpu / connections * 1e6, server_cpu / connections * 1e6))

				rate   = statistics.median(r[0] for r in results)
				client = statistics.median(r[1] for r in results)
				server = statistics.median(r[2] for r in results)

				print("%-10s %-8s %-26s %14.0f %16.1f %16.1f" % (
					"fireperf", version, "%s, %s" % (backend, test), rate, client, server))

				self.record("connections (%s)" % test, "fireperf %s" % version, backend, 16,
					duration="", connections_per_s="%.0f" % rate,
					client_cpu_us_per_connection="%.2f" % client,
					server_cpu_us_per_connection="%.2f" % server,
					total_connections=n, isolated=int(isolated))

	def write(self):
		if not self.args.output:
			return

		rows = []

		# Read what is there, so that it still lines up if we have added columns
		try:
			with open(self.args.output, newline="") as f:
				rows = list(csv.DictReader(f))
		except FileNotFoundError:
			pass

		rows += self.rows

		with open(self.args.output, "w", newline="") as f:
			writer = csv.DictWriter(f, fieldnames=FIELDS, restval="", lineterminator="\n")

			writer.writeheader()
			writer.writerows(rows)

		print("\n%d result(s) have been added to %s" % (len(self.rows), self.args.output))

def main():
	parser = argparse.ArgumentParser(description="Benchmarks fireperf")

	parser.add_argument("--fireperf", default="./fireperf",
		help="The fireperf binary to benchmark (default: %(default)s)")
	parser.add_argument("--duration", type=int, default=5,
		help="How long every throughput test runs in seconds (default: %(default)s)")
	parser.add_argument("--runs", type=int, default=3,
		help="How often every throughput test is being run (default: %(default)s)")
	parser.add_argument("--connections", type=int, default=200000,
		help="How many connections are being opened in every connection test. "
			"This is limited to 10000 if the tests cannot be run in a network "
			"namespace of their own, which needs root (default: %(default)s)")
	parser.add_argument("--no-connections", action="store_true",
		help="Don't run the connection tests")
	parser.add_argument("--output", default=None,
		help="Append the results to this CSV file")
	parser.add_argument("--note", default="",
		help="Something to remember this run by")
	parser.add_argument("--remote", default=None, metavar="[USER@]HOST",
		help="Let this host (which we reach over SSH) send to us instead of "
			"running over loopback. This only runs the throughput tests.")
	parser.add_argument("--remote-directory", default="/tmp/fireperf-benchmark",
		help="Where to put fireperf on the other host, which must allow to "
			"execute what is in it (default: %(default)s)")

	args = parser.parse_args()

	benchmark = Benchmark(args)

	print("fireperf %s on %s (%s, %s cores, Linux %s)" % (
		benchmark.environment["commit"], benchmark.environment["host"],
		benchmark.environment["cpu"], benchmark.environment["cores"],
		benchmark.environment["kernel"]))

	if benchmark.remote:
		print("Sent from %s (%s, %s cores, Linux %s) at %s to %s" % (
			benchmark.remote.name, benchmark.environment["remote_cpu"],
			benchmark.environment["remote_cores"], benchmark.environment["remote_kernel"],
			benchmark.remote.address, benchmark.remote.local_address))

		for connections in (1, 8):
			benchmark.remote_throughput(connections)

	else:
		for connections in (1, 8):
			benchmark.throughput(connections)

		if not args.no_connections:
			benchmark.connections()

	benchmark.write()

main()
