#! /usr/bin/python3
# SPDX-License-Identifier: GPL-2.0-or-later
#
# Benchmark for passt/pasta and other things for comparison
#
# Copyright Red Hat
# Author: David Gibson <david@gibson.dropbear.id.au>
#

# How to use
#
# For best results, run on an idle system with external networks
# disabled.
#
#  1. Adjust TIME, OMIT, INTERVAL globals to select the length of the
#     test
#  2. Run
#        $ ./benchmark.py <scenarios>
#
# Scenarios:
#
# loopback - Direct transfer over lo interface
#    Run within a network namespace both to isolate and remove impact
#    of any netfilter rules.
#
# veth_{in,out} - Transfer over kernel veth device
#    'in' is from parent netns to child, 'out' is reverse
#
# pasta_{in,out} - Transfer over pasta managed netns
#
# slirp4netns_{in,out} - Transfer over slirp4netns managed netns
#
# qemu_user_{in,out} - Transfer host <-> qemu guest with -net user
#
# passt_{in,out} - Transfer host <-> qemu guest with passt
#
# qemu_tap_{in,out} - Transfer host <-> qemu guest with -net tap
#

import sys
import os
import time
import socket
import math
import contextlib
import json
import subprocess
import statistics

TIME, OMIT, INTERVAL = 6, 1, 1
# TIME, OMIT, INTERVAL = 600, 10, 5

NSTOOL = "../../test/nstool"
PASST = "../../passt"
PASTA = "../../pasta"
MBUTO = "../../test/mbuto.img"
GUEST_KEY = "../../test/guest-key"

DEBUG = True

SERVER = ["iperf3", "-s", "-1", "-I", "server.pid"]
PORT = 5201

IP1 = "192.0.2.1"
IP2 = "192.0.2.2"
IP3 = "192.0.2.3"
IP4 = "192.0.2.4"


def dbg(fmt, *args, **kwargs):
    if DEBUG:
        print("DBG: " + fmt.format(*args, **kwargs), file=sys.stderr)


# Base class for command runners
class Cmd:
    def popen(self, args, **kwargs):
        return subprocess.Popen(self.cmd_prefix + args, **kwargs)

    def run(self, args, **kwargs):
        return subprocess.run(self.cmd_prefix + args, **kwargs)


class HostCmd(Cmd):
    def __init__(self):
        self.cmd_prefix = []


class NsTool(contextlib.AbstractContextManager, Cmd):
    def __init__(self, ctl, parent, holdprefix):
        try:
            os.remove(ctl)
        except FileNotFoundError:
            pass
        self.ctl = ctl
        self.parent = parent
        self.proc = parent.popen(holdprefix + [NSTOOL, "hold", ctl],
                                 stdout=subprocess.DEVNULL, stderr=None)
        self.cmd_prefix = [NSTOOL, "exec", ctl, "--"]

    def __enter__(self):
        self.parent.run([NSTOOL, "info", "-w", self.ctl],
                        stdout=subprocess.DEVNULL).check_returncode()
        self.proc.__enter__()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return self.proc.__exit__(exc_type, exc_value, traceback)

    def relpid(self):
        out = self.parent.run([NSTOOL, "info", "-p", self.ctl],
                              stdout=subprocess.PIPE)
        return int(out.stdout)

    def stop(self):
        return self.parent.run([NSTOOL, "stop", self.ctl]).check_returncode()


# We put everything into an isolated network namespace for a few reasons:
# 1. Make sure test traffic doesn't leak out onto external interfaces
# 2. Lets us allocate our own addresses freely
# 3. Avoids firewall rules which might be set up on the host affecting
#    performance
@contextlib.contextmanager
def isolate(ctl):
    with NsTool(ctl, HostCmd(), ["unshare", "-Unrp"]) as isons:
        isons.run(["ip", "link", "set", "lo", "up"]).check_returncode()
        yield isons
        isons.stop()


@contextlib.contextmanager
def loopback(isons):
    with isons.popen(SERVER):
        yield (isons, "127.0.0.1")


@contextlib.contextmanager
def veth_ns(isons, ip_outer, ip_inner):
    isons.run(["ip", "link", "add", "type", "veth"])
    with NsTool("veth_in.ctl", isons, ["unshare", "-n"]) as guestns:
        dbg("veth_ns guest ns ready")
        isons.run(["ip", "link", "set", "veth1",
                   "netns", str(guestns.relpid())])
        isons.run(["ip", "link", "set", "veth0", "up"])
        isons.run(["ip", "addr", "add", f"{ip_outer}/24", "dev", "veth0"])
        guestns.run(["ip", "link", "set", "veth1", "up"])
        guestns.run(["ip", "addr", "add", f"{ip_inner}/24", "dev", "veth1"])
        yield guestns
        guestns.stop()


@contextlib.contextmanager
def veth_in(isons):
    outer = IP3
    inner = IP4
    with veth_ns(isons, outer, inner) as guestns:
        dbg("veth_in guestns ready")
        with guestns.popen(SERVER):
            yield (isons, inner)


@contextlib.contextmanager
def veth_out(isons):
    outer = IP3
    inner = IP4
    with veth_ns(isons, outer, inner) as guestns:
        dbg("veth_in guestns ready")
        with isons.popen(SERVER):
            yield (guestns, outer)


def passt_config_isons(isons):
    isons.run(["ip", "link", "add", "type", "dummy"]).check_returncode()
    isons.run(["ip", "link", "set", "dummy0", "up"]).check_returncode()
    isons.run(["ip", "addr", "add", f"{IP1}/24",
               "dev", "dummy0"]).check_returncode()
    isons.run(["ip", "route", "add", "default", "via", str(IP2),
               "dev", "dummy0"]).check_returncode()


@contextlib.contextmanager
def pasta_in(isons):
    passt_config_isons(isons)
    cmd = [PASTA, "--config-net", "-t", str(PORT), "--"] + SERVER
    with isons.popen(cmd):
        time.sleep(1)
        yield (isons, IP1)


@contextlib.contextmanager
def pasta_out(isons):
    passt_config_isons(isons)
    pasta = NsTool("pasta_out.ctl", isons, [PASTA, "--config-net", "--"])
    with isons.popen(SERVER), pasta as pasta:
        time.sleep(1)
        yield (pasta, IP2)
        pasta.stop()


class Slirp4NetNs(contextlib.AbstractContextManager):
    def __init__(self, ns, name, pid):
        self.apipath = f"{name}.api"
        self.ns = ns
        self.pid = pid

    def api(self, req):
        with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
            sock.connect(self.apipath)
            sock.sendall(req)

    def __enter__(self):
        exit_r, exit_w = os.pipe()
        self.exit = os.fdopen(exit_w, "wb")
        cmd = ["slirp4netns", "-c", "-a", self.apipath, "-e", str(exit_r),
               str(self.pid), "tap0"]
        self.proc = self.ns.popen(cmd, pass_fds=(exit_r,))
        self.proc.__enter__()
        time.sleep(1)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.exit.close()
        return self.proc.__exit__(exc_type, exc_value, traceback)


@contextlib.contextmanager
def slirp4netns_in(isons):
    with NsTool("slirp4netns_in.ctl", isons, ["unshare", "-n"]) as guestns:
        with Slirp4NetNs(isons, "slirp4netns_in", guestns.relpid()) as slirp:
            req = f'''
{{
    "execute": "add_hostfwd",
    "arguments": {{
        "proto": "tcp",
        "host_addr": "0.0.0.0",
        "host_port": {PORT},
        "guest_addr": "10.0.2.100",
        "guest_port": {PORT}
    }}
}}'''
            slirp.api(bytes(req, 'UTF-8'))
            with guestns.popen(SERVER):
                yield (isons, "127.0.0.1")
                guestns.stop()
            dbg("guestns closed")
        dbg("slirp exited")


@contextlib.contextmanager
def slirp4netns_out(isons):
    with isons.popen(SERVER), \
         NsTool("slirp4netns_in.ctl", isons, ["unshare", "-n"]) as guestns:
        with Slirp4NetNs(isons, "slirp4netns_out", guestns.relpid()):
            yield (guestns, "10.0.2.2")
            guestns.stop()


class MbutoQemu(Cmd, contextlib.AbstractContextManager):
    def __init__(self, parent, name, opts):
        self.ns = parent
        self.name = name
        self.ssh_config = f"{name}.ssh"
        self.pidfile = f"{name}.pid"

        cid = hash(name) & 0xffffffff
        self.hosts = os.path.realpath(f"{name}.hosts")
        try:
            os.remove(self.hosts)
        except FileNotFoundError:
            pass
        key = os.path.realpath(GUEST_KEY)

        with open(self.ssh_config, "w") as c:
            c.write(f"""
Host {name}
    User root
    UserKnownHostsFile {self.hosts}
    StrictHostKeyChecking no
    IdentityFile {key}
    IdentityAgent none
    ProxyCommand socat - VSOCK-CONNECT:{cid}:22
""")

        kernel = f"/boot/vmlinuz-{os.uname().release}"

        self.qemu_cmd = ["qemu-kvm", "-kernel", kernel, "-initrd", MBUTO,
                         "-append", '"console=ttyS0"',
                         "-nodefaults", "-nographic", "-serial", "stdio",
                         "-device", f"vhost-vsock-pci,guest-cid={cid}",
                         "-pidfile", self.pidfile] + list(opts)

        self.cmd_prefix = ["ssh", "-F", self.ssh_config, "-t", name, "--"]

    def __enter__(self):
        dbg(f"Starting {self.qemu_cmd}")
        self.qemu = self.ns.popen(self.qemu_cmd)
        self.qemu.__enter__()
        while self.run([":"], stderr=subprocess.DEVNULL).returncode != 0:
            pass

        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pid = int(open(self.pidfile, "r").read())
        self.ns.run(["kill", str(pid)])
        ret = self.qemu.__exit__(exc_type, exc_value, traceback)
        os.remove(self.ssh_config)
        os.remove(self.hosts)
        return ret

    def dhcp(self):
        self.run(["ip", "link", "set", "lo", "up"]).check_returncode()
        self.run(["ip", "link", "set", "eth0", "up"]).check_returncode()
        self.run(["dhclient", "-v", "eth0"]).check_returncode()

    def static(self, ip):
        self.run(["ip", "link", "set", "lo", "up"]).check_returncode()
        self.run(["ip", "link", "set", "eth0", "up"]).check_returncode()
        self.run(["ip", "addr", "add", ip, "dev", "eth0"]).check_returncode()


@contextlib.contextmanager
def qemu_user_in(isons):
    opts = ["-m", "4G",
            "-netdev", f"user,id=slirp,hostfwd=tcp:127.0.0.1:{PORT}-:{PORT}",
            "-device", "virtio-net-pci,netdev=slirp"]
    with MbutoQemu(isons, "slirp-qemu", opts) as qemu:
        qemu.dhcp()
        with qemu.popen(SERVER):
            time.sleep(2)
            dbg("Guest server ready")
            yield (isons, "127.0.0.1")


@contextlib.contextmanager
def qemu_user_out(isons):
    opts = ["-m", "4G", "-netdev", "user,id=slirp",
            "-device", "virtio-net-pci,netdev=slirp"]
    with MbutoQemu(isons, "slirp-qemu", opts) as qemu:
        qemu.dhcp()
        with isons.popen(SERVER):
            time.sleep(2)
            yield (qemu, "10.0.2.2")


@contextlib.contextmanager
def passt(ns, sock, opts=[]):
    passt_cmd = [PASST, "-s", sock, "-1", "-f"] + opts
    with ns.popen(passt_cmd):
        yield
    os.remove(sock)


@contextlib.contextmanager
def passt_qemu(ns, name, opts=[]):
    sock = f"{name}.passt"
    with passt(ns, sock, opts):
        streamopts = f"id=passt,server=off,addr.type=unix,addr.path={sock}"
        qemu_opts = ["-m", "4G",
                     "-netdev", f"stream,{streamopts}",
                     "-device", "virtio-net-pci,netdev=passt"]
        with MbutoQemu(ns, name, qemu_opts) as qemu:
            qemu.dhcp()
            yield qemu


@contextlib.contextmanager
def passt_in(isons):
    passt_config_isons(isons)
    with passt_qemu(isons, "passt_in", ["-t", str(PORT)]) as qemu:
        with qemu.popen(SERVER):
            time.sleep(2)
            dbg("Guest server ready")
            yield (isons, "127.0.0.1")


@contextlib.contextmanager
def passt_out(isons):
    passt_config_isons(isons)
    with passt_qemu(isons, "passt_out") as qemu:
        with isons.popen(SERVER):
            time.sleep(2)
            yield (qemu, IP2)


@contextlib.contextmanager
def tap_qemu(ns, name, opts):
    opts += ["-netdev", "tap,id=tap,script=no,downscript=no",
             "-device", "virtio-net-pci,netdev=tap"]
    with MbutoQemu(ns, name, opts) as qemu:
        ns.run(["ip", "link", "set", "tap0", "up"]).check_returncode()
        ns.run(["ip", "addr", "add", f"{IP3}/24",
                "dev", "tap0"]).check_returncode()
        qemu.static(f"{IP4}/24")
        yield qemu


@contextlib.contextmanager
def qemu_tap_in(isons):
    opts = ["-m", "4G"]
    with tap_qemu(isons, "tap_in", opts) as qemu:
        with qemu.popen(SERVER):
            time.sleep(2)
            dbg("Guest server ready")
            yield (isons, IP4)


@contextlib.contextmanager
def qemu_tap_out(isons):
    opts = ["-m", "4G"]
    with tap_qemu(isons, "tap_out", opts) as qemu:
        with isons.popen(SERVER):
            time.sleep(2)
            yield (qemu, IP3)


def bench_one(name, setup):
    print(f"Running benchmark {name}...")
    with setup as (crun, addr):
        dbg("Setup for {}", name)
        ccmd = ["iperf3", "-c", addr, "-J", "-t", str(TIME), "-O", str(OMIT),
                "-i", str(INTERVAL), "-l", "1M"]
        out = crun.run(ccmd, stdout=subprocess.PIPE, stderr=None)
        dbg("Client completed for {}", name)
    out.check_returncode()
    return out.stdout


def stats(name, out):
    parse = json.loads(out)
    intervals = parse['intervals']
    sums = [i['sum'] for i in intervals]
    samples = [s['bits_per_second'] for s in sums if not s['omitted']]
    mean = statistics.fmean(samples)
    stdev = statistics.stdev(samples)
    spread = stdev / mean
    n = len(samples)
    stderr = stdev / math.sqrt(float(n))
    relerr = stderr / abs(mean)
    gbps = mean / 1000000000
    print(f"{name:16}:\t{gbps:8.2f} Gbps ± {relerr:.1%} " + \
          f"({n} intervals, {spread:.1%})")


ALL = ["loopback", "veth_in", "veth_out", "pasta_in", "pasta_out",
       "slirp4netns_in", "slirp4netns_out", "qemu_user_in", "qemu_user_out",
       "qemu_tap_in", "qemu_tap_out", "passt_in", "passt_out"]

if __name__ == "__main__":
    results = {}
    cases = sys.argv[1:]
    if not cases:
        cases = ALL
    for case in cases:
        with isolate("isolate.ctl") as isons:
            setup = globals()[case](isons)
            out = bench_one(case, setup)
            open(f"{case}.json", "wb").write(out)
            results[case] = out
    for name, out in results.items():
        stats(name, out)