// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::Result;
use netbench::stats::{Initialize, Stats};
use once_cell::sync::OnceCell;
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::{
    collections::{BTreeSet, HashMap},
    io::BufRead,
    path::PathBuf,
};
use structopt::StructOpt;

static VEGA: &str = include_str!("./vega.json");
static VEGA_PARSED: OnceCell<serde_json::Value> = OnceCell::new();

fn vega_template() -> serde_json::Value {
    VEGA_PARSED
        .get_or_init(|| serde_json::from_str(VEGA).unwrap())
        .clone()
}

#[derive(Debug, Default, StructOpt)]
pub struct Report {
    pub inputs: Vec<PathBuf>,
    #[structopt(short, long)]
    pub output: Option<PathBuf>,
}

impl Report {
    pub fn run(&self) -> Result<()> {
        let mut stats_table = vec![];
        let mut stream_table = vec![];
        let mut signals = vec![];
        let mut names = vec![];
        let mut scenario_names = BTreeSet::new();
        let mut stream_ids = HashMap::new();
        let mut trace_ids = vec![];
        let mut pids = vec![];

        for (pid, input) in self.inputs.iter().enumerate() {
            let pid = pid as u64;
            let input = std::fs::File::open(input)?;
            let mut input = std::io::BufReader::new(input);

            let mut first = String::new();
            input.read_line(&mut first)?;
            let Initialize {
                driver,
                scenario,
                traces,
                ..
            } = serde_json::from_str(&first)?;

            let name = driver
                .split('/')
                .last()
                .unwrap()
                .trim_start_matches("netbench-driver-")
                .to_string();

            let scenario_name = scenario
                .split('/')
                .last()
                .unwrap()
                .trim_end_matches(".json");

            scenario_names.insert(scenario_name.to_string());

            pids.push(format!("!indata('data$hidden', 'name', {name:?})"));
            names.push(name);

            let input = serde_json::de::IoRead::new(input);
            let input = serde_json::StreamDeserializer::new(input);
            let mut prev_x = 0;
            for event in input {
                let Stats {
                    time,
                    cpu,
                    cycles,
                    instructions,
                    branches,
                    context_switches,
                    memory,
                    virtual_memory,
                    allocs,
                    reallocs,
                    deallocs,
                    syscalls,
                    connections,
                    accept,
                    send,
                    receive,
                    connect_time,
                    profiles,
                } = event?;

                let x = time.as_millis() as u64;

                macro_rules! emit {
                    ($name:ident, $value:expr) => {{
                        let mut y = $value as f64;
                        if !f64::is_normal(y) {
                            y = 0.0;
                        }
                        stats_table.push(Row {
                            x,
                            y,
                            pid,
                            stat: Stat::$name as _,
                            stream_id: None,
                        });
                    }};
                    ($name:ident, $value:expr, $id:expr) => {{
                        let mut y = $value as f64;
                        if !f64::is_normal(y) {
                            y = 0.0;
                        }
                        let next_id = stream_ids.len() as u64;
                        let stream_id = stream_ids.entry((pid, $id)).or_insert_with(|| {
                            // area charts need another point so zero it out first
                            stream_table.push(Row {
                                x: prev_x,
                                y: 0.0,
                                pid,
                                stat: Stat::StreamSendBytes as _,
                                stream_id: Some(next_id),
                            });
                            stream_table.push(Row {
                                x: prev_x,
                                y: 0.0,
                                pid,
                                stat: Stat::StreamReceiveBytes as _,
                                stream_id: Some(next_id),
                            });
                            next_id
                        });
                        stream_table.push(Row {
                            x,
                            y,
                            pid,
                            stat: Stat::$name as _,
                            stream_id: Some(*stream_id),
                        });
                    };};
                }

                emit!(Cpu, cpu);
                emit!(Memory, memory);
                emit!(VirtualMemory, virtual_memory);
                emit!(Cycles, cycles);
                emit!(Instructions, instructions);
                emit!(Branches, branches);
                emit!(ContextSwitches, context_switches);
                emit!(Syscalls, syscalls);
                emit!(Connections, connections);
                emit!(Accept, accept);
                emit!(AllocBytes, allocs.total);
                emit!(AllocCount, allocs.count);
                emit!(ReallocBytes, reallocs.total);
                emit!(ReallocCount, reallocs.count);
                emit!(DeallocBytes, deallocs.total);
                emit!(DeallocCount, deallocs.count);

                {
                    let mut bytes = 0;
                    let mut count = 0;
                    for (id, s) in send {
                        bytes += s.total;
                        count += s.count;
                        emit!(StreamSendBytes, s.total, id);
                    }
                    emit!(SendBytes, bytes);
                    emit!(SendCount, count);
                    emit!(SendBytesPerCpu, bytes as f64 / cpu as f64);
                    emit!(SendBytesPerInstruction, bytes as f64 / instructions as f64);
                }

                {
                    let mut bytes = 0;
                    let mut count = 0;
                    for (id, s) in receive {
                        bytes += s.total;
                        count += s.count;
                        emit!(StreamReceiveBytes, s.total, id);
                    }
                    emit!(ReceiveBytes, bytes);
                    emit!(ReceiveCount, count);
                    emit!(ReceiveBytesPerCpu, bytes as f64 / cpu as f64);
                    emit!(
                        ReceiveBytesPerInstruction,
                        bytes as f64 / instructions as f64
                    );
                }

                {
                    let mut y = connect_time.average();

                    if !f64::is_normal(y) {
                        y = 0.0;
                    }

                    // convert micros to seconds
                    y /= 1_000_000.0;

                    stats_table.push(Row {
                        x,
                        y,
                        pid,
                        stat: Stat::ConnectTime as _,
                        stream_id: None,
                    });
                }

                for (trace_id, hist) in profiles {
                    let trace = &traces[trace_id as usize];
                    let trace_id = if let Some(id) = trace_ids.iter().position(|v| v == trace) {
                        id as u64
                    } else {
                        let id = trace_ids.len() as u64;
                        trace_ids.push(trace.to_string());
                        id
                    };

                    // offset the stat id with the built-in names
                    let stat = Stat::NAMES.len() as u64 + trace_id;

                    let mut y = hist.stat.average();
                    if !f64::is_normal(y) {
                        y = 0.0;
                    }

                    // convert micros to seconds
                    let y = y / 1_000_000.0;

                    stats_table.push(Row {
                        x,
                        y,
                        pid,
                        stat,
                        stream_id: None,
                    });

                    /*
                     // TODO figure out how to visualize multiple histograms over time
                    for bucket in hist.buckets {
                        profile_hist_table.push(Bucket {
                            x,
                            pid,
                            trace_id,
                            lower: bucket.lower,
                            upper: bucket.upper,
                            count: bucket.count,
                        });
                    }
                    */
                }

                prev_x = x;
            }
        }

        stats_table.sort_by(|a, b| {
            a.x.cmp(&b.x)
                .then(a.pid.cmp(&b.pid))
                .then(a.stat.cmp(&b.stat))
        });
        stream_table.sort_by(|a, b| {
            a.x.cmp(&b.x)
                .then(a.pid.cmp(&b.pid))
                .then(a.stat.cmp(&b.stat))
                .then(a.stream_id.cmp(&b.stream_id))
        });

        let mut view_names = Stat::NAMES
            .iter()
            .map(|v| v.to_string())
            .collect::<Vec<_>>();

        for trace in &trace_ids {
            view_names.push(format!("trace - {trace}"));
        }

        // expose an option to select the view
        signals.push(json!({
            "name": "ui$view",
            "value": &view_names[0],
            "bind": {
                "input": "select",
                "name": "View",
                "options": view_names,
            },
        }));

        // translate the view name into an index
        signals.push(json!({
            "name": "sig$view",
            "value": "0",
            "update": format!("indexof({view_names:?},ui$view)"),
        }));

        signals.push(json!({
            "name": "pids",
            "value": "[]",
            "update": format!("[{}]", pids.join(",")),
        }));

        let mut stream_counts = vec![0u64; pids.len()];
        for (pid, _) in stream_ids.keys() {
            stream_counts[*pid as usize] += 1;
        }

        let mut stream_count_expr = "0".to_string();
        for (id, count) in stream_counts.iter().enumerate() {
            use core::fmt::Write;
            let _ = write!(stream_count_expr, "+(pids[{id}]?{count}:0)");
        }

        signals.push(json!({
            "name": "sig$streamCount",
            "update": stream_count_expr,
        }));

        signals.push(json!({
            "name": "sig$streamTypes",
            "update": format!("{{{}:(ui$sendStreams?-1:0),{}:(ui$recvStreams?1:0)}}", Stat::StreamSendBytes as u64, Stat::StreamReceiveBytes as u64),
        }));

        let mut output: serde_json::Value = vega_template();

        let root = output.as_object_mut().unwrap();

        root.insert(
            "title".to_string(),
            scenario_names
                .iter()
                .map(|v| v.as_str())
                .collect::<Vec<_>>()
                .join(", ")
                .into(),
        );

        root.get_mut("signals")
            .unwrap()
            .as_array_mut()
            .unwrap()
            .extend(signals);

        let data = root.get_mut("data").unwrap().as_array_mut().unwrap();

        data.insert(
            0,
            json!({
                "name": "data$stats",
                "values": &stats_table,
            }),
        );
        data.insert(
            1,
            json!({
                "name": "data$streams",
                "values": &stream_table,
            }),
        );
        data.insert(
            2,
            json!({
                "name": "data$drivers",
                "values": names.iter().map(|name|{
                    json!({ "name": name })
                }).collect::<Vec<_>>()
            }),
        );

        if let Some(out_path) = self.output.as_ref() {
            if let Some(parent) = out_path.parent() {
                std::fs::create_dir_all(parent)?;
            }
            let mut out_file = std::fs::File::create(out_path)?;
            serde_json::to_writer(&mut out_file, &output)?;
        } else {
            serde_json::to_writer(std::io::stdout(), &output)?;
        }

        Ok(())
    }
}

macro_rules! stat {
    (enum Stat { $($name:ident = $desc:expr),* $(,)? }) => {
        #[repr(u64)]
        enum Stat {
            $(
                $name,
            )*
        }

        impl Stat {
            const NAMES: &'static [&'static str] = &[$($desc,)*];
        }
    };
}

stat!(
    enum Stat {
        Cpu = "cpu %",
        Memory = "memory (bytes)",
        VirtualMemory = "virtual memory (bytes)",
        Cycles = "cycles",
        Instructions = "instructions",
        Branches = "branches",
        ContextSwitches = "context-switches",
        Syscalls = "syscalls",
        Connections = "connections",
        ConnectTime = "connect-time",
        Accept = "accept (streams)",
        AllocBytes = "alloc (bytes)",
        AllocCount = "alloc (count)",
        ReallocBytes = "realloc (bytes)",
        ReallocCount = "realloc (count)",
        DeallocBytes = "dealloc (bytes)",
        DeallocCount = "dealloc (count)",
        SendBytes = "send (bytes)",
        SendCount = "send (count)",
        SendBytesPerCpu = "send (bytes/cpu %)",
        SendBytesPerInstruction = "send (bytes/instruction)",
        ReceiveBytes = "receive (bytes)",
        ReceiveCount = "receive (count)",
        ReceiveBytesPerCpu = "receive (bytes/cpu %)",
        ReceiveBytesPerInstruction = "receive (bytes/instruction)",
        StreamSendBytes = "stream send",
        StreamReceiveBytes = "stream receive",
    }
);

#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
struct Row {
    x: u64,
    y: f64,
    #[serde(rename = "p")]
    pid: u64,
    #[serde(rename = "s")]
    stat: u64,
    #[serde(rename = "i", skip_serializing_if = "Option::is_none")]
    stream_id: Option<u64>,
}

#[derive(Clone, Copy, Debug, Deserialize, Serialize)]
struct Bucket {
    x: u64,
    #[serde(rename = "p")]
    pid: u64,
    #[serde(rename = "t")]
    trace_id: u64,
    #[serde(rename = "l")]
    lower: f64,
    #[serde(rename = "u")]
    upper: f64,
    #[serde(rename = "c")]
    count: u64,
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{report::Report, Result};
use serde_json::json;
use std::{
    collections::BTreeMap,
    path::{Path, PathBuf},
};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub struct ReportTree {
    input_dir: PathBuf,
    out_dir: PathBuf,
}

static INDEX_HTML: &str = include_str!("./report_tree.html");

type ScenarioMap = BTreeMap<String, Report>;

impl ReportTree {
    pub fn run(&self) -> Result<()> {
        let mut client_scenarios: ScenarioMap = Default::default();
        let mut server_scenarios: ScenarioMap = Default::default();

        for scenario in self.input_dir.read_dir()? {
            let scenario = scenario?;
            let path = scenario.path();
            let scenario_name = if let Some(name) = path_name(&path) {
                name
            } else {
                continue;
            };

            for driver in scenario.path().read_dir()? {
                let driver = driver?.path();

                macro_rules! push_scenario {
                    ($target:ident, $name:literal) => {{
                        let endpoint = driver.join(concat!($name, ".json"));
                        if endpoint.exists() {
                            $target
                                .entry(scenario_name.to_string())
                                .or_insert_with(|| Report {
                                    output: Some(
                                        self.out_dir
                                            .join(scenario_name)
                                            .join(concat!($name, "s.json")),
                                    ),
                                    ..Default::default()
                                })
                                .inputs
                                .push(endpoint);
                        }
                    }};
                }

                push_scenario!(client_scenarios, "client");
                push_scenario!(server_scenarios, "server");
            }
        }

        std::fs::create_dir_all(&self.out_dir)?;

        let index = {
            let template = handlebars::Handlebars::new();

            template.render_template(
                INDEX_HTML,
                &json!({
                    "clients": render_scenarios(client_scenarios)?,
                    "servers": render_scenarios(server_scenarios)?,
                }),
            )?
        };

        std::fs::write(self.out_dir.join("index.html"), index)?;

        Ok(())
    }
}

fn render_scenarios(scenarios: ScenarioMap) -> Result<Vec<String>> {
    let mut names = vec![];
    for (name, report) in scenarios {
        names.push(name);
        report.run()?;
    }
    Ok(names)
}

fn path_name(path: &Path) -> Option<&str> {
    let stem = path.file_stem()?.to_str()?;

    // filter out any hidden files
    if stem.starts_with('.') {
        return None;
    }

    Some(stem)
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{procinfo::Proc, Result};
use netbench::{
    stats::{Bucket, Histogram, Initialize, Print, Stat, Stats, StreamId},
    units::ByteExt as _,
};
use serde_json::json;
use std::{
    collections::HashMap,
    io,
    io::BufRead,
    process::{Command, Stdio},
    time::Duration,
};

static PROGRAM: &str = include_str!("./netbench.bt");

#[derive(Debug)]
struct Report {
    count: u64,
    interval: Duration,
    prev: Stats,
    current: Stats,
    send: HashMap<StreamId, Stat>,
    receive: HashMap<StreamId, Stat>,
    profiles: HashMap<u64, Histogram>,
    pending_profile: Option<u64>,
    proc: Option<Proc>,
}

impl Report {
    fn new(interval: Duration) -> Self {
        Self {
            interval,
            count: 0,
            prev: Default::default(),
            current: Default::default(),
            send: Default::default(),
            receive: Default::default(),
            profiles: Default::default(),
            pending_profile: None,
            proc: None,
        }
    }

    fn push(&mut self, line: &str) {
        let line = line.trim();

        if let Some(id) = self.pending_profile {
            if let Some(bucket) = parse_hist_line(line) {
                self.profiles.entry(id).or_default().buckets.push(bucket);
                return;
            } else {
                self.pending_profile = None;
            }
        }

        if line.is_empty() {
            return;
        }

        if line == "===" {
            self.dump();
            return;
        }

        if let Some(count) = line.strip_prefix("@: ") {
            self.count = count.parse().unwrap();
            return;
        }

        if let Some(trace) = line.strip_prefix("@P[") {
            let trace = trace.trim_end_matches("]:");
            if let Ok(trace) = trace.parse() {
                self.pending_profile = Some(trace);
                return;
            }
        }

        if let Some(line) = line.strip_prefix("@p[") {
            let (id, line) = line.split_once("]: ").unwrap();
            let id = id.parse().unwrap();

            let stat = BpfParse::parse(line).unwrap();
            self.profiles.entry(id).or_default().stat = stat;
            return;
        }

        macro_rules! stat {
            ($prefix:literal, $name:ident) => {
                if let Some(value) = line.strip_prefix(concat!("@", $prefix, ": ")) {
                    self.current.$name = BpfParse::parse(value).unwrap();
                    return;
                }
            };
        }

        stat!("c", cycles);
        stat!("i", instructions);
        stat!("b", branches);
        stat!("C", context_switches);
        stat!("a", allocs);
        stat!("R", reallocs);
        stat!("d", deallocs);
        stat!("S", syscalls);
        stat!("O", connections);
        stat!("A", accept);
        stat!("h", connect_time);

        macro_rules! try_map {
            ($prefix:literal, $on_value:expr) => {
                if let Some(line) = line.strip_prefix(concat!("@", $prefix, "[")) {
                    let mut on_value = $on_value;
                    let (id, line) = line.split_once("]: ").unwrap();
                    let (conn, id) = id.split_once(", ").unwrap();
                    let connection_id = conn.parse().unwrap();
                    let id = id.parse().unwrap();
                    let id = StreamId { connection_id, id };
                    let value = BpfParse::parse(line).unwrap();
                    on_value(id, value);
                    return;
                }
            };
        }

        macro_rules! stream_stat {
            ($prefix:literal, $name:ident) => {
                try_map!($prefix, |id, value| self.$name.insert(id, value));
            };
        }

        stream_stat!("s", send);
        stream_stat!("r", receive);

        if let Some(pid) = line.strip_prefix("cpid=") {
            let pid = pid.parse().unwrap();
            self.proc = Some(Proc::new(pid));

            // dump the initial numbers
            self.dump();
            return;
        }

        eprintln!("> {line}");
    }

    fn dump(&mut self) {
        if let Some(proc) = self.proc.as_mut() {
            proc.load(&mut self.current);
        }

        self.entry().print().unwrap();

        // reset the values
        core::mem::swap(&mut self.prev, &mut self.current);
        self.current.allocs = Default::default();
        self.current.reallocs = Default::default();
        self.current.deallocs = Default::default();
    }

    fn entry(&mut self) -> Stats {
        let current = &self.current;
        let prev = &self.prev;
        let cycles = current.cycles.saturating_sub(prev.cycles);
        let instructions = current.instructions.saturating_sub(prev.instructions);
        let branches = current.branches.saturating_sub(prev.branches);
        let context_switches = current
            .context_switches
            .saturating_sub(prev.context_switches);
        let syscalls = current.syscalls.saturating_sub(prev.syscalls);
        let connections = current.connections.saturating_sub(prev.connections);
        let accept = current.accept.saturating_sub(prev.accept);

        let time = self.interval.as_millis() as u64 * self.count;
        let time = Duration::from_millis(time);

        Stats {
            time,
            cpu: current.cpu,
            cycles,
            instructions,
            branches,
            context_switches,
            syscalls,
            connections,
            accept,
            memory: current.memory,
            virtual_memory: current.virtual_memory,
            allocs: current.allocs,
            reallocs: current.reallocs,
            deallocs: current.deallocs,
            connect_time: current.connect_time,
            send: core::mem::take(&mut self.send),
            receive: core::mem::take(&mut self.receive),
            profiles: core::mem::take(&mut self.profiles),
        }
    }
}

trait BpfParse: Sized {
    fn parse(s: &str) -> Result<Self>;
}

impl BpfParse for u64 {
    fn parse(s: &str) -> Result<Self> {
        Ok(s.parse()?)
    }
}

impl BpfParse for Stat {
    fn parse(s: &str) -> Result<Self> {
        let mut parts = s.split(", ");

        macro_rules! part {
            ($name:ident) => {{
                let err = concat!("missing ", stringify!($name));
                let value = parts.next().ok_or(err)?;
                let value = value
                    .strip_prefix(concat!(stringify!($name), " "))
                    .ok_or(err)?;
                let value = value.parse().map_err(|_| err)?;
                value
            }};
        }

        let count = part!(count);
        // skip average since it can be computed with the other values
        let _average = parts.next();
        let total = part!(total);

        Ok(Self { count, total })
    }
}

fn parse_hist_line(line: &str) -> Option<Bucket> {
    let mut parts = line.split_whitespace();

    let lower = parts.next()?;
    let upper = parts.next()?;
    let count = parts.next()?;

    let lower = lower.trim_start_matches('[');
    let lower = lower.trim_end_matches(',');
    let lower = parse_hist_bound(lower)?;

    let upper = upper.trim_end_matches(')');
    let upper = parse_hist_bound(upper)?;

    let count = count.parse().ok()?;

    Some(Bucket {
        lower,
        upper,
        count,
    })
}

fn parse_hist_bound(s: &str) -> Option<u64> {
    // try parsing without any suffix
    if let Ok(v) = s.parse() {
        return Some(v);
    }

    let number_index = s
        .char_indices()
        .find_map(|(idx, c)| {
            if !(c.is_numeric() || c == '.') {
                Some(idx)
            } else {
                None
            }
        })
        .unwrap_or(s.len());

    let mut v: f64 = s[..number_index].parse().ok()?;

    let suffix = s[number_index..].trim();

    v *= *match suffix {
        "" => 1.bytes(),
        "K" | "k" => 1.kilobytes(),
        "Ki" | "ki" => 1.kibibytes(),
        "M" | "m" => 1.megabytes(),
        "Mi" | "mi" => 1.mebibytes(),
        "G" | "g" => 1.gigabytes(),
        "Gi" | "gi" => 1.gibibytes(),
        "T" | "t" => 1.terabytes(),
        "Ti" | "ti" => 1.tebibytes(),
        _ => return None,
    } as f64;

    Some(v as u64)
}

pub fn try_run(args: &crate::Args) -> Result<Option<()>> {
    let mut command = if let Ok(bpftrace) = find_bpftrace() {
        eprintln!("collecting stats with bpftrace");
        Command::new(bpftrace)
    } else {
        return Ok(None);
    };

    let driver = &args.driver;
    let interval = args.interval;
    let scenario = args.scenario()?;
    let scenario_path = &args.scenario;

    let program = {
        let template = handlebars::Handlebars::new();
        template.render_template(
            PROGRAM,
            &json!({
                "bin": &driver,
                "interval_ms": interval.as_millis() as u64,
                "libc": libc_location(driver)?.unwrap_or_else(|| driver.to_string()),
                "hardware": detect_hardware_events()?
            }),
        )?
    };

    command
        .arg("-c")
        .arg(driver)
        .arg("-e")
        .arg(program)
        .env("TRACE", "disabled")
        .env("SCENARIO", scenario_path)
        .stdout(Stdio::piped());

    let mut proc = command.spawn()?;

    Initialize {
        pid: proc.id() as _,
        driver: driver.to_string(),
        scenario: scenario_path.to_string(),
        traces: scenario.traces.to_vec(),
        ..Default::default()
    }
    .print()?;

    let output = proc.stdout.take().unwrap();
    let handle = std::thread::spawn(move || {
        let output = io::BufReader::new(output);
        let mut report = Report::new(interval);
        for line in output.lines() {
            if let Ok(line) = line {
                report.push(&line);
            } else {
                break;
            }
        }
        report.dump();
    });

    proc.wait()?;

    let _ = handle.join();

    Ok(Some(()))
}

fn find_bpftrace() -> Result<String> {
    let out = Command::new("which").arg("bpftrace").output()?;
    if out.status.success() {
        let out = core::str::from_utf8(&out.stdout)?;
        let out = out.trim();
        Ok(out.to_string())
    } else {
        Err("missing bpftrace".into())
    }
}

// finds the libc location for the given program
fn libc_location(cmd: &str) -> Result<Option<String>> {
    let out = Command::new("ldd").arg(cmd).output()?;
    let out = core::str::from_utf8(&out.stdout)?;
    for line in out.lines() {
        let line = line.trim();
        if line.starts_with("libc") {
            let (_, path) = line.split_once("=>").unwrap();
            let (path, _) = path.split_once('(').unwrap();
            let path = path.trim();
            return Ok(Some(path.to_string()));
        }
    }

    Ok(None)
}

fn detect_hardware_events() -> Result<bool> {
    let out = Command::new("perf").arg("list").arg("hw").output()?;
    let out = core::str::from_utf8(&out.stdout)?;

    if out.contains("cycles") && out.contains("branches") && out.contains("instructions") {
        return Ok(true);
    }

    Ok(false)
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{procinfo::Proc, Args, Result};
use netbench::stats::{Initialize, Print, Stats};
use std::{
    process::Command,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    time::Duration,
};

pub fn run(args: &Args) -> Result<()> {
    let mut command = Command::new(&args.driver);

    let driver = &args.driver;
    let interval = args.interval;
    let scenario_path = &args.scenario;
    let scenario = args.scenario()?;

    command
        .env("TRACE", "disabled")
        .env("SCENARIO", scenario_path);

    let mut proc = command.spawn()?;
    let info = Proc::new(proc.id());

    Initialize {
        pid: proc.id() as _,
        driver: driver.to_string(),
        scenario: scenario_path.to_string(),
        traces: scenario.traces.to_vec(),
        ..Default::default()
    }
    .print()?;

    let is_open = Arc::new(AtomicBool::new(true));
    let is_open_handle = is_open.clone();

    let handle = std::thread::spawn(move || {
        collect(info, interval, is_open_handle);
    });

    proc.wait()?;

    is_open.store(false, Ordering::Relaxed);

    let _ = handle.join();

    Ok(())
}

fn collect(mut proc: Proc, interval: Duration, is_open: Arc<AtomicBool>) {
    let mut stats = Stats::default();

    loop {
        proc.load(&mut stats);
        stats.print().unwrap();

        if !is_open.load(Ordering::Relaxed) {
            return;
        }

        std::thread::sleep(interval);
        stats.time += interval;
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{scenario::Scenario, units::parse_duration, Result};
use std::time::Duration;
use structopt::StructOpt;

mod bpftrace;
mod generic;
mod procinfo;

#[derive(Debug, StructOpt)]
pub struct Args {
    pub driver: String,

    #[structopt(long, short, env = "SCENARIO")]
    pub scenario: String,

    #[structopt(long, short, parse(try_from_str=parse_duration), default_value = "1s")]
    pub interval: Duration,
}

impl Args {
    pub fn scenario(&self) -> Result<Scenario> {
        Scenario::open(std::path::Path::new(&self.scenario))
    }
}

fn main() -> Result<()> {
    let args = Args::from_args();

    // try to use bpftrace
    if let Some(()) = bpftrace::try_run(&args)? {
        return Ok(());
    }

    // fall back to the generic collector
    generic::run(&args)?;

    Ok(())
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::stats::Stats;
use sysinfo::{CpuRefreshKind, Pid, ProcessExt, RefreshKind, System, SystemExt};

#[derive(Debug)]
pub struct Proc {
    pid: Pid,
    system: System,
}

impl Proc {
    pub fn new(pid: u32) -> Self {
        Self {
            pid: Pid::from(pid as i32),
            system: System::new_with_specifics(
                RefreshKind::new()
                    .with_cpu(CpuRefreshKind::new().with_cpu_usage())
                    .with_memory(),
            ),
        }
    }

    pub fn load(&mut self, stats: &mut Stats) {
        self.system.refresh_process(self.pid);
        if let Some(proc) = self.system.process(self.pid) {
            stats.cpu = proc.cpu_usage();
            // memory is returned in KB
            stats.memory = proc.memory() * 1000;
            stats.virtual_memory = proc.virtual_memory() * 1000;
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use probe::probe;

pub struct Allocator {
    inner: mimalloc::MiMalloc,
}

impl Allocator {
    pub const fn new() -> Self {
        Self {
            inner: mimalloc::MiMalloc,
        }
    }
}

#[inline(never)]
#[allow(unused_variables)]
fn alloc(size: usize) {
    probe!(netbench, netbench__alloc, size);
}

#[inline(never)]
#[allow(unused_variables)]
fn dealloc(size: usize) {
    probe!(netbench, netbench__dealloc, size);
}

#[inline(never)]
#[allow(unused_variables)]
fn realloc(prev_size: usize, new_size: usize) {
    probe!(netbench, netbench__realloc, prev_size, new_size);
}

unsafe impl std::alloc::GlobalAlloc for Allocator {
    #[inline]
    unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
        alloc(layout.size());
        self.inner.alloc(layout)
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
        dealloc(layout.size());
        self.inner.dealloc(ptr, layout)
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
        realloc(layout.size(), new_size);
        self.inner.realloc(ptr, layout, new_size)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{
    client::{self, AddressMap},
    multiplex, scenario, trace,
    units::Byte,
    Error, Result,
};
use std::{net::IpAddr, ops::Deref, path::Path, str::FromStr, sync::Arc, time::Duration};
use structopt::StructOpt;

mod alloc;
pub use alloc::Allocator;

const TRACE_VALUES: &[&str] = &["disabled", "throughput", "stdio"];

#[derive(Debug, StructOpt)]
pub struct Server {
    #[structopt(short, long, default_value = "::")]
    pub ip: IpAddr,

    #[structopt(short, long, default_value = "4433", env = "PORT")]
    pub port: u16,

    #[structopt(long, default_value = "netbench")]
    pub application_protocols: Vec<String>,

    #[structopt(long, default_value = "0", env = "SERVER_ID")]
    pub server_id: usize,

    #[structopt(long, default_value = "throughput", possible_values = TRACE_VALUES, env = "TRACE")]
    pub trace: Vec<String>,

    #[structopt(long, short = "V")]
    pub verbose: bool,

    #[structopt(long, default_value = "8KiB")]
    pub rx_buffer: Byte,

    #[structopt(long, default_value = "8KiB")]
    pub tx_buffer: Byte,

    #[structopt(env = "SCENARIO")]
    pub scenario: Scenario,

    #[structopt(long)]
    pub nagle: bool,

    #[structopt(long, env = "MULTITHREADED")]
    pub multithreaded: Option<Option<bool>>,

    /// Forces multiplex mode for the driver
    ///
    /// Without this, the requirement is inferred based on the scenario
    #[structopt(long, env = "MULTIPLEX")]
    multiplex: Option<Option<bool>>,
}

impl Server {
    pub fn runtime(&self) -> tokio::runtime::Runtime {
        let multithreaded = match self.multithreaded {
            Some(Some(v)) => v,
            Some(None) => true,
            None => false,
        };
        if multithreaded {
            tokio::runtime::Builder::new_multi_thread()
        } else {
            tokio::runtime::Builder::new_current_thread()
        }
        .enable_all()
        .build()
        .unwrap()
    }

    pub fn scenario(&self) -> Arc<scenario::Server> {
        let id = self.server_id;
        self.scenario.servers[id].clone()
    }

    pub fn certificate(&self) -> (&Arc<scenario::Certificate>, &Arc<scenario::Certificate>) {
        let id = self.server_id;
        let server = &self.scenario.servers[id];
        let cert = &self.scenario.certificates[server.certificate as usize];
        let private_key = &self.scenario.certificates[server.private_key as usize];
        (cert, private_key)
    }

    pub fn trace(&self) -> impl trace::Trace + Clone {
        traces(&self.trace[..], self.verbose, &self.scenario.traces)
    }

    pub fn multiplex(&self) -> Option<multiplex::Config> {
        // TODO infer this based on the scenario requirements
        if is_multiplex_enabled(self.multiplex) {
            // TODO load this from the scenario configuration
            Some(multiplex::Config::default())
        } else {
            None
        }
    }
}

#[derive(Debug, StructOpt)]
pub struct Client {
    #[structopt(long, default_value = "netbench")]
    pub application_protocols: Vec<String>,

    #[structopt(short, long, default_value = "::", env = "LOCAL_IP")]
    pub local_ip: IpAddr,

    #[structopt(long, default_value = "0", env = "CLIENT_ID")]
    pub client_id: usize,

    #[structopt(long, default_value = "throughput", possible_values = TRACE_VALUES, env = "TRACE")]
    pub trace: Vec<String>,

    #[structopt(long, short = "V")]
    pub verbose: bool,

    #[structopt(long, default_value = "8KiB")]
    pub rx_buffer: Byte,

    #[structopt(long, default_value = "8KiB")]
    pub tx_buffer: Byte,

    #[structopt(env = "SCENARIO")]
    pub scenario: Scenario,

    #[structopt(long)]
    pub nagle: bool,

    #[structopt(long, env = "MULTITHREADED")]
    pub multithreaded: Option<Option<bool>>,

    /// Forces multiplex mode for the driver
    ///
    /// Without this, the requirement is inferred based on the scenario
    #[structopt(long, env = "MULTIPLEX")]
    multiplex: Option<Option<bool>>,
}

impl Client {
    pub fn runtime(&self) -> tokio::runtime::Runtime {
        let multithreaded = match self.multithreaded {
            Some(Some(v)) => v,
            Some(None) => true,
            None => false,
        };
        if multithreaded {
            tokio::runtime::Builder::new_multi_thread()
        } else {
            tokio::runtime::Builder::new_current_thread()
        }
        .enable_all()
        .build()
        .unwrap()
    }

    pub fn scenario(&self) -> Arc<scenario::Client> {
        let id = self.client_id;
        self.scenario.clients[id].clone()
    }

    pub fn certificate_authorities(&self) -> impl Iterator<Item = Arc<scenario::Certificate>> + '_ {
        let id = self.client_id;
        let certs = &self.scenario.certificates;
        self.scenario.clients[id]
            .certificate_authorities
            .iter()
            .copied()
            .map(move |ca| certs[ca as usize].clone())
    }

    pub async fn address_map(&self) -> Result<AddressMap> {
        let id = self.client_id as u64;
        AddressMap::new(&self.scenario, id, &mut Resolver).await
    }

    pub fn trace(&self) -> impl trace::Trace + Clone {
        traces(&self.trace[..], self.verbose, &self.scenario.traces)
    }

    pub fn multiplex(&self) -> Option<multiplex::Config> {
        // TODO infer this based on the scenario requirements
        if is_multiplex_enabled(self.multiplex) {
            // TODO load this from the scenario configuration
            Some(multiplex::Config::default())
        } else {
            None
        }
    }
}

fn is_multiplex_enabled(opt: Option<Option<bool>>) -> bool {
    match opt {
        Some(Some(v)) => v,
        Some(None) => true,
        None => false,
    }
}

struct Resolver;

impl Resolver {
    fn get(&self, key: String) -> Result<String> {
        let host =
            std::env::var(&key).map_err(|_| format!("missing {key} environment variable"))?;
        Ok(host)
    }
}

impl client::Resolver for Resolver {
    fn server(&mut self, id: u64) -> Result<String> {
        self.get(format!("SERVER_{id}"))
    }

    fn router(&mut self, router_id: u64, server_id: u64) -> Result<String> {
        self.get(format!("ROUTER_{router_id}_SERVER_{server_id}"))
    }
}

#[derive(Clone, Debug)]
pub struct Scenario(Arc<scenario::Scenario>);

impl FromStr for Scenario {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self> {
        let scenario = scenario::Scenario::open(Path::new(s))?;
        Ok(Self(Arc::new(scenario)))
    }
}

impl Deref for Scenario {
    type Target = scenario::Scenario;

    fn deref(&self) -> &Self::Target {
        self.0.deref()
    }
}

fn traces(trace: &[String], verbose: bool, traces: &Arc<Vec<String>>) -> impl trace::Trace + Clone {
    let enabled = !trace.iter().any(|v| v == "disabled");

    let throughput = if enabled && trace.iter().any(|v| v == "throughput") {
        let trace = trace::Throughput::default();
        trace.reporter(Duration::from_secs(1));
        Some(trace)
    } else {
        None
    };

    let stdio = if enabled && trace.iter().any(|v| v == "stdio") {
        let mut trace = trace::StdioLogger::new(traces.clone());
        trace.verbose(verbose);
        Some(trace)
    } else {
        None
    };

    let usdt = trace::Usdt::default();

    (usdt, (throughput, stdio))
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{duplex, multiplex, scenario, Result};
use netbench_driver::Allocator;
use std::{collections::HashSet, future::Future, net::SocketAddr, pin::Pin, sync::Arc};
use structopt::StructOpt;
use tokio::{
    io::{self, AsyncWriteExt},
    net::TcpStream,
};
use tokio_native_tls::{
    native_tls::{Certificate, TlsConnector},
    TlsStream,
};

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

fn main() -> Result<()> {
    let args = Client::from_args();
    let runtime = args.opts.runtime();
    runtime.block_on(args.run())
}

#[derive(Debug, StructOpt)]
pub struct Client {
    #[structopt(flatten)]
    opts: netbench_driver::Client,
}

impl Client {
    pub async fn run(&self) -> Result<()> {
        let addresses = self.opts.address_map().await?;
        let scenario = self.opts.scenario();

        let mut trace = self.opts.trace();
        let mut checkpoints = HashSet::new();
        let mut timer = netbench::timer::Tokio::default();

        if let Some(config) = self.opts.multiplex() {
            let client = self.multiplex_client(config)?;
            let client = netbench::Client::new(client, &scenario, &addresses);
            client.run(&mut trace, &mut checkpoints, &mut timer).await?;
        } else {
            let client = self.duplex_client()?;
            let client = netbench::Client::new(client, &scenario, &addresses);
            client.run(&mut trace, &mut checkpoints, &mut timer).await?;
        }

        Ok(())
    }

    fn duplex_client(&self) -> Result<ClientImpl> {
        let mut builder = TlsConnector::builder();
        for ca in self.opts.certificate_authorities() {
            let ca = Certificate::from_pem(ca.pem.as_bytes())?;
            builder.add_root_certificate(ca);
        }
        let connector = builder.build()?;
        let connector: tokio_native_tls::TlsConnector = connector.into();
        let connector = Arc::new(connector);

        Ok(ClientImpl {
            connector,
            id: 0,
            rx_buffer: *self.opts.rx_buffer as _,
            tx_buffer: *self.opts.tx_buffer as _,
            nagle: self.opts.nagle,
        })
    }

    fn multiplex_client(&self, config: multiplex::Config) -> Result<MultiplexClientImpl> {
        let client = self.duplex_client()?;
        Ok(MultiplexClientImpl { client, config })
    }
}

type Stream = io::BufStream<TcpStream>;
type Connection<'a> = netbench::Driver<'a, duplex::Connection<TlsStream<Stream>>>;
type MultiplexConnection<'a> = netbench::Driver<'a, multiplex::Connection<TlsStream<Stream>>>;

#[derive(Clone, Debug)]
struct ClientImpl {
    connector: Arc<tokio_native_tls::TlsConnector>,
    id: u64,
    rx_buffer: usize,
    tx_buffer: usize,
    nagle: bool,
}

impl ClientImpl {
    fn id(&mut self) -> u64 {
        let id = self.id;
        self.id = id + 1;
        id
    }
}

impl<'a> netbench::client::Client<'a> for ClientImpl {
    type Connect = Pin<Box<dyn Future<Output = Result<Self::Connection>> + 'a>>;
    type Connection = Connection<'a>;

    fn connect(
        &mut self,
        addr: SocketAddr,
        server_name: &str,
        server_conn_id: u64,
        scenario: &'a Arc<scenario::Connection>,
    ) -> Self::Connect {
        let id = self.id();
        let connector = self.connector.clone();
        let server_name = server_name.to_string();
        let rx_buffer = self.rx_buffer;
        let tx_buffer = self.tx_buffer;
        let nagle = self.nagle;

        let fut = async move {
            let conn = TcpStream::connect(addr).await?;

            if !nagle {
                let _ = conn.set_nodelay(true);
            }

            let conn = io::BufStream::with_capacity(rx_buffer, tx_buffer, conn);

            let mut conn = connector.connect(&server_name, conn).await?;

            // The native-tls crate does not expose the server name on the server so we need to
            // write the connection id for now.
            conn.write_u64(server_conn_id).await?;

            let conn = Box::pin(conn);
            let conn = duplex::Connection::new(id, conn);
            let conn: Self::Connection = netbench::Driver::new(scenario, conn);

            Result::Ok(conn)
        };

        Box::pin(fut)
    }
}

#[derive(Clone, Debug)]
struct MultiplexClientImpl {
    config: multiplex::Config,
    client: ClientImpl,
}

impl<'a> netbench::client::Client<'a> for MultiplexClientImpl {
    type Connect = Pin<Box<dyn Future<Output = Result<Self::Connection>> + 'a>>;
    type Connection = MultiplexConnection<'a>;

    fn connect(
        &mut self,
        addr: SocketAddr,
        server_name: &str,
        server_conn_id: u64,
        scenario: &'a Arc<scenario::Connection>,
    ) -> Self::Connect {
        let id = self.client.id();
        let config = self.config.clone();
        let connector = self.client.connector.clone();
        let server_name = server_name.to_string();
        let rx_buffer = self.client.rx_buffer;
        let tx_buffer = self.client.tx_buffer;
        let nagle = self.client.nagle;

        let fut = async move {
            let conn = TcpStream::connect(addr).await?;

            if !nagle {
                let _ = conn.set_nodelay(true);
            }

            let conn = io::BufStream::with_capacity(rx_buffer, tx_buffer, conn);

            let mut conn = connector.connect(&server_name, conn).await?;

            // The native-tls crate does not expose the server name on the server so we need to
            // write the connection id for now.
            conn.write_u64(server_conn_id).await?;

            let conn = Box::pin(conn);
            let conn = multiplex::Connection::new(id, conn, config);
            let conn: Self::Connection = netbench::Driver::new(scenario, conn);

            Result::Ok(conn)
        };

        Box::pin(fut)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{multiplex, scenario, Result, Timer};
use netbench_driver::Allocator;
use std::{collections::HashSet, sync::Arc};
use structopt::StructOpt;
use tokio::{
    io::{self, AsyncReadExt},
    net::{TcpListener, TcpStream},
    spawn,
};
use tokio_native_tls::native_tls::{Identity, TlsAcceptor};

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

fn main() -> Result<()> {
    let args = Server::from_args();
    let runtime = args.opts.runtime();
    runtime.block_on(args.run())
}

#[derive(Debug, StructOpt)]
pub struct Server {
    #[structopt(flatten)]
    opts: netbench_driver::Server,
}

impl Server {
    pub async fn run(&self) -> Result<()> {
        let scenario = self.opts.scenario();
        let buffer = (*self.opts.rx_buffer as usize, *self.opts.tx_buffer as usize);

        let server = self.server().await?;

        let trace = self.opts.trace();
        let config = self.opts.multiplex();
        let ident = self.identity()?;
        let acceptor = TlsAcceptor::builder(ident).build()?;
        let acceptor: tokio_native_tls::TlsAcceptor = acceptor.into();
        let acceptor = Arc::new(acceptor);

        let mut conn_id = 0;
        loop {
            let (connection, _addr) = server.accept().await?;

            if !self.opts.nagle {
                let _ = connection.set_nodelay(true);
            }

            let scenario = scenario.clone();
            let id = conn_id;
            conn_id += 1;
            let acceptor = acceptor.clone();
            let trace = trace.clone();
            let config = config.clone();
            spawn(async move {
                if let Err(err) =
                    handle_connection(acceptor, connection, id, scenario, trace, config, buffer)
                        .await
                {
                    eprintln!("error: {err}");
                }
            });
        }

        async fn handle_connection(
            acceptor: Arc<tokio_native_tls::TlsAcceptor>,
            connection: TcpStream,
            conn_id: u64,
            scenario: Arc<scenario::Server>,
            mut trace: impl netbench::Trace,
            config: Option<multiplex::Config>,
            (rx_buffer, tx_buffer): (usize, usize),
        ) -> Result<()> {
            let connection = io::BufStream::with_capacity(rx_buffer, tx_buffer, connection);

            let mut timer = netbench::timer::Tokio::default();
            let before = timer.now();

            let connection = acceptor.accept(connection).await?;

            let now = timer.now();
            trace.connect(now, conn_id, now - before);

            let mut connection = Box::pin(connection);

            let server_idx = connection.read_u64().await?;
            let scenario = scenario
                .connections
                .get(server_idx as usize)
                .ok_or("invalid connection id")?;

            let mut checkpoints = HashSet::new();

            if let Some(config) = config {
                let conn = netbench::multiplex::Connection::new(conn_id, connection, config);
                let conn = netbench::Driver::new(scenario, conn);
                conn.run(&mut trace, &mut checkpoints, &mut timer).await?;
            } else {
                let conn = netbench::duplex::Connection::new(conn_id, connection);
                let conn = netbench::Driver::new(scenario, conn);
                conn.run(&mut trace, &mut checkpoints, &mut timer).await?;
            }

            Ok(())
        }
    }

    async fn server(&self) -> Result<TcpListener> {
        let server = TcpListener::bind((self.opts.ip, self.opts.port)).await?;
        Ok(server)
    }

    fn identity(&self) -> Result<Identity> {
        let (_, private_key) = self.opts.certificate();
        let cert = &private_key.pkcs12;
        let ident = Identity::from_pkcs12(cert, "")?;
        Ok(ident)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::Result;
use netbench_driver::Allocator;
use s2n_quic::provider::io;
use std::collections::HashSet;
use structopt::StructOpt;

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

fn main() -> Result<()> {
    let args = Client::from_args();
    let runtime = args.opts.runtime();
    runtime.block_on(args.run())
}

#[derive(Debug, StructOpt)]
pub struct Client {
    #[structopt(flatten)]
    opts: netbench_driver::Client,

    #[structopt(long, default_value = "9001", env = "MAX_MTU")]
    max_mtu: u16,

    #[structopt(long, env = "DISABLE_GSO")]
    disable_gso: bool,
}

impl Client {
    pub async fn run(&self) -> Result<()> {
        let addresses = self.opts.address_map().await?;
        let scenario = self.opts.scenario();

        let client = self.client()?;
        let client = netbench::Client::new(client, &scenario, &addresses);
        let mut trace = self.opts.trace();
        let mut checkpoints = HashSet::new();
        let mut timer = netbench::timer::Tokio::default();
        let mut client = client.run(&mut trace, &mut checkpoints, &mut timer).await?;

        client.wait_idle().await?;

        Ok(())
    }

    fn client(&self) -> Result<s2n_quic::Client> {
        let mut tls = s2n_quic::provider::tls::default::Client::builder()
            // handle larger cert chains
            .with_max_cert_chain_depth(10)?
            .with_application_protocols(
                self.opts.application_protocols.iter().map(String::as_bytes),
            )?
            .with_key_logging()?;

        for ca in self.opts.certificate_authorities() {
            tls = tls.with_certificate(ca.pem.as_str())?;
        }

        let tls = tls.build()?;

        let mut io_builder = io::Default::builder()
            .with_max_mtu(self.max_mtu)?
            .with_receive_address((self.opts.local_ip, 0u16).into())?;

        if self.disable_gso {
            io_builder = io_builder.with_gso_disabled()?;
        }

        let io = io_builder.build()?;

        let client = s2n_quic::Client::builder()
            .with_io(io)?
            .with_tls(tls)?
            .start()
            .unwrap();

        Ok(client)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{scenario, timer::Timestamp, Result, Timer};
use netbench_driver::Allocator;
use s2n_quic::{provider::io, Connection};
use std::{collections::HashSet, sync::Arc};
use structopt::StructOpt;
use tokio::spawn;

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

fn main() -> Result<()> {
    let args = Server::from_args();
    let runtime = args.opts.runtime();
    runtime.block_on(args.run())
}

#[derive(Debug, StructOpt)]
pub struct Server {
    #[structopt(flatten)]
    opts: netbench_driver::Server,

    #[structopt(long, default_value = "9001", env = "MAX_MTU")]
    max_mtu: u16,

    #[structopt(long, env = "DISABLE_GSO")]
    disable_gso: bool,
}

impl Server {
    pub async fn run(&self) -> Result<()> {
        let scenario = self.opts.scenario();
        let trace = self.opts.trace();

        let mut server = self.server(trace.clone())?;

        while let Some(connection) = server.accept().await {
            let scenario = scenario.clone();
            let trace = trace.clone();
            spawn(async move {
                if let Err(error) = handle_connection(connection, scenario, trace).await {
                    eprintln!("error: {error:#}");
                }
            });
        }

        return Err("server shut down unexpectedly".into());

        async fn handle_connection(
            connection: Connection,
            scenario: Arc<scenario::Server>,
            mut trace: impl netbench::Trace,
        ) -> Result<()> {
            let server_name = connection.server_name()?.ok_or("missing server name")?;
            let scenario = scenario.on_server_name(&server_name)?;
            let conn =
                netbench::Driver::new(scenario, netbench::s2n_quic::Connection::new(connection));

            let mut checkpoints = HashSet::new();
            let mut timer = netbench::timer::Tokio::default();

            conn.run(&mut trace, &mut checkpoints, &mut timer).await?;

            Ok(())
        }
    }

    fn server(&self, trace: impl netbench::Trace + Send + 'static) -> Result<s2n_quic::Server> {
        let (certificate, private_key) = self.opts.certificate();
        let certificate = certificate.pem.as_str();
        let private_key = private_key.pem.as_str();

        let tls = s2n_quic::provider::tls::default::Server::builder()
            .with_certificate(certificate, private_key)?
            .with_application_protocols(
                self.opts.application_protocols.iter().map(String::as_bytes),
            )?
            .with_key_logging()?
            .build()?;

        let mut io_builder = io::Default::builder()
            .with_max_mtu(self.max_mtu)?
            .with_receive_address((self.opts.ip, self.opts.port).into())?;

        if self.disable_gso {
            io_builder = io_builder.with_gso_disabled()?;
        }

        let io = io_builder.build()?;

        let server = s2n_quic::Server::builder()
            .with_io(io)?
            .with_tls(tls)?
            .with_event(EventTracer::new(trace))?
            .start()
            .unwrap();

        Ok(server)
    }
}

struct EventTracer<T> {
    trace: T,
    timer: netbench::timer::Tokio,
}

impl<T> EventTracer<T> {
    fn new(trace: T) -> Self {
        Self {
            trace,
            timer: Default::default(),
        }
    }
}

impl<T: 'static + Send + netbench::Trace> s2n_quic::provider::event::Subscriber for EventTracer<T> {
    type ConnectionContext = Timestamp;

    #[inline]
    fn create_connection_context(
        &mut self,
        _meta: &s2n_quic::provider::event::events::ConnectionMeta,
        _info: &s2n_quic::provider::event::ConnectionInfo,
    ) -> Timestamp {
        self.timer.now()
    }

    #[inline]
    fn on_handshake_status_updated(
        &mut self,
        start: &mut Timestamp,
        meta: &s2n_quic::provider::event::events::ConnectionMeta,
        event: &s2n_quic::provider::event::events::HandshakeStatusUpdated,
    ) {
        use s2n_quic::provider::event::events::HandshakeStatus;

        // record the difference of when we started the connection and completed the handshake
        if let HandshakeStatus::Complete { .. } = event.status {
            let now = self.timer.now();
            self.trace.connect(now, meta.id, now - *start)
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{duplex, multiplex, scenario, Result};
use netbench_driver::Allocator;
use s2n_tls::{
    config::{Builder, Config},
    error::Error,
    security::DEFAULT_TLS13,
};
use s2n_tls_tokio::{TlsConnector, TlsStream};
use std::{collections::HashSet, future::Future, net::SocketAddr, pin::Pin, sync::Arc};
use structopt::StructOpt;
use tokio::net::TcpStream;

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

fn main() -> Result<()> {
    let args = Client::from_args();
    let runtime = args.opts.runtime();
    runtime.block_on(args.run())
}

#[derive(Debug, StructOpt)]
pub struct Client {
    #[structopt(flatten)]
    opts: netbench_driver::Client,
}

impl Client {
    pub async fn run(&self) -> Result<()> {
        let addresses = self.opts.address_map().await?;
        let scenario = self.opts.scenario();

        let mut trace = self.opts.trace();
        let mut checkpoints = HashSet::new();
        let mut timer = netbench::timer::Tokio::default();

        if let Some(config) = self.opts.multiplex() {
            let client = self.multiplex_client(config)?;
            let client = netbench::Client::new(client, &scenario, &addresses);
            client.run(&mut trace, &mut checkpoints, &mut timer).await?;
        } else {
            let client = self.duplex_client()?;
            let client = netbench::Client::new(client, &scenario, &addresses);
            client.run(&mut trace, &mut checkpoints, &mut timer).await?;
        }

        Ok(())
    }

    fn duplex_client(&self) -> Result<ClientImpl> {
        let connector = self.connector()?;
        Ok(ClientImpl {
            connector,
            id: 0,
            nagle: self.opts.nagle,
        })
    }

    fn multiplex_client(&self, config: multiplex::Config) -> Result<MultiplexClientImpl> {
        let client = self.duplex_client()?;
        Ok(MultiplexClientImpl { config, client })
    }

    fn connector(&self) -> Result<Arc<s2n_tls_tokio::TlsConnector<Config>>> {
        let connector = TlsConnector::new(self.config()?.build()?);
        let connector: s2n_tls_tokio::TlsConnector<Config> = connector;
        let connector = Arc::new(connector);
        Ok(connector)
    }

    fn config(&self) -> Result<Builder, Error> {
        let mut builder = Config::builder();
        builder.set_security_policy(&DEFAULT_TLS13)?;
        for ca in self.opts.certificate_authorities() {
            builder.trust_pem(ca.pem.as_bytes())?;
        }
        Ok(builder)
    }
}

type Connection<'a> = netbench::Driver<'a, duplex::Connection<TlsStream<TcpStream>>>;
type MultiplexConnection<'a> = netbench::Driver<'a, multiplex::Connection<TlsStream<TcpStream>>>;

#[derive(Clone)]
struct ClientImpl {
    connector: Arc<s2n_tls_tokio::TlsConnector<Config>>,
    id: u64,
    nagle: bool,
}

impl ClientImpl {
    fn id(&mut self) -> u64 {
        let id = self.id;
        self.id = id + 1;
        id
    }
}

impl<'a> netbench::client::Client<'a> for ClientImpl {
    type Connect = Pin<Box<dyn Future<Output = Result<Self::Connection>> + 'a>>;
    type Connection = Connection<'a>;

    fn connect(
        &mut self,
        addr: SocketAddr,
        server_name: &str,
        _server_conn_id: u64,
        scenario: &'a Arc<scenario::Connection>,
    ) -> Self::Connect {
        let id = self.id();
        let connector = self.connector.clone();
        let nagle = self.nagle;
        let server_name = server_name.to_string();

        let fut = async move {
            let conn = TcpStream::connect(addr).await?;

            if !nagle {
                let _ = conn.set_nodelay(true);
            }

            let conn = connector.connect(&server_name, conn).await?;
            let conn = Box::pin(conn);
            let conn = duplex::Connection::new(id, conn);
            let conn: Self::Connection = netbench::Driver::new(scenario, conn);

            Result::Ok(conn)
        };

        Box::pin(fut)
    }
}

#[derive(Clone)]
struct MultiplexClientImpl {
    config: multiplex::Config,
    client: ClientImpl,
}

impl<'a> netbench::client::Client<'a> for MultiplexClientImpl {
    type Connect = Pin<Box<dyn Future<Output = Result<Self::Connection>> + 'a>>;
    type Connection = MultiplexConnection<'a>;

    fn connect(
        &mut self,
        addr: SocketAddr,
        server_name: &str,
        _server_conn_id: u64,
        scenario: &'a Arc<scenario::Connection>,
    ) -> Self::Connect {
        let id = self.client.id();
        let config = self.config.clone();
        let connector = self.client.connector.clone();
        let nagle = self.client.nagle;
        let server_name = server_name.to_string();

        let fut = async move {
            let conn = TcpStream::connect(addr).await?;

            if !nagle {
                let _ = conn.set_nodelay(true);
            }

            let conn = connector.connect(&server_name, conn).await?;
            let conn = Box::pin(conn);
            let conn = multiplex::Connection::new(id, conn, config);
            let conn: Self::Connection = netbench::Driver::new(scenario, conn);

            Result::Ok(conn)
        };

        Box::pin(fut)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{duplex, multiplex, scenario, Driver, Result, Timer};
use netbench_driver::Allocator;
use s2n_tls::{
    config::{Builder, Config},
    error::Error,
    security::DEFAULT_TLS13,
};
use s2n_tls_tokio::TlsAcceptor;
use std::{collections::HashSet, sync::Arc};
use structopt::StructOpt;
use tokio::{
    net::{TcpListener, TcpStream},
    spawn,
};

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

fn main() -> Result<()> {
    let args = Server::from_args();
    let runtime = args.opts.runtime();
    runtime.block_on(args.run())
}

#[derive(Debug, StructOpt)]
pub struct Server {
    #[structopt(flatten)]
    opts: netbench_driver::Server,
}

impl Server {
    pub async fn run(&self) -> Result<()> {
        let scenario = self.opts.scenario();

        let server = self.server().await?;

        let trace = self.opts.trace();
        let config = self.opts.multiplex();

        let acceptor = TlsAcceptor::new(self.config()?.build()?);
        let acceptor: s2n_tls_tokio::TlsAcceptor<Config> = acceptor;
        let acceptor = Arc::new(acceptor);

        let mut conn_id = 0;
        loop {
            let (connection, _addr) = server.accept().await?;

            if !self.opts.nagle {
                let _ = connection.set_nodelay(true);
            }

            let scenario = scenario.clone();
            let id = conn_id;
            conn_id += 1;
            let acceptor = acceptor.clone();
            let trace = trace.clone();
            let config = config.clone();
            spawn(async move {
                if let Err(err) =
                    handle_connection(acceptor, connection, id, scenario, trace, config).await
                {
                    eprintln!("error: {err}");
                }
            });
        }

        async fn handle_connection(
            acceptor: Arc<s2n_tls_tokio::TlsAcceptor<Config>>,
            connection: TcpStream,
            conn_id: u64,
            scenario: Arc<scenario::Server>,
            mut trace: impl netbench::Trace,
            config: Option<multiplex::Config>,
        ) -> Result<()> {
            let mut timer = netbench::timer::Tokio::default();
            let before = timer.now();

            let connection = acceptor.accept(connection).await?;

            let now = timer.now();
            trace.connect(now, conn_id, now - before);

            let server_name = connection
                .as_ref()
                .server_name()
                .ok_or("missing server name")?;
            let scenario = scenario.on_server_name(server_name)?;

            let connection = Box::pin(connection);

            let mut checkpoints = HashSet::new();

            if let Some(config) = config {
                let conn = multiplex::Connection::new(conn_id, connection, config);
                let conn = Driver::new(scenario, conn);
                conn.run(&mut trace, &mut checkpoints, &mut timer).await?;
            } else {
                let conn = duplex::Connection::new(conn_id, connection);
                let conn = Driver::new(scenario, conn);
                conn.run(&mut trace, &mut checkpoints, &mut timer).await?;
            }

            Ok(())
        }
    }

    async fn server(&self) -> Result<TcpListener> {
        let server = TcpListener::bind((self.opts.ip, self.opts.port)).await?;
        Ok(server)
    }

    fn config(&self) -> Result<Builder, Error> {
        let (cert, private_key) = self.opts.certificate();

        let mut builder = Config::builder();
        builder.set_security_policy(&DEFAULT_TLS13)?;
        builder.load_pem(cert.pem.as_bytes(), private_key.pem.as_bytes())?;

        Ok(builder)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{duplex, multiplex, scenario, Result};
use netbench_driver::Allocator;
use std::{collections::HashSet, future::Future, net::SocketAddr, pin::Pin, sync::Arc};
use structopt::StructOpt;
use tokio::{
    io::{self, AsyncWriteExt},
    net::TcpStream,
};

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    Client::from_args().run().await
}

#[derive(Debug, StructOpt)]
pub struct Client {
    #[structopt(flatten)]
    opts: netbench_driver::Client,
}

impl Client {
    pub async fn run(&self) -> Result<()> {
        let addresses = self.opts.address_map().await?;
        let scenario = self.opts.scenario();

        let mut trace = self.opts.trace();
        let mut checkpoints = HashSet::new();
        let mut timer = netbench::timer::Tokio::default();

        if let Some(config) = self.opts.multiplex() {
            let client = self.multiplex_client(config)?;
            let client = netbench::Client::new(client, &scenario, &addresses);
            client.run(&mut trace, &mut checkpoints, &mut timer).await?;
        } else {
            let client = self.duplex_client()?;
            let client = netbench::Client::new(client, &scenario, &addresses);
            client.run(&mut trace, &mut checkpoints, &mut timer).await?;
        }

        Ok(())
    }

    fn duplex_client(&self) -> Result<ClientImpl> {
        Ok(ClientImpl {
            id: 0,
            nagle: self.opts.nagle,
            rx_buffer: *self.opts.rx_buffer as _,
            tx_buffer: *self.opts.tx_buffer as _,
        })
    }

    fn multiplex_client(&self, config: multiplex::Config) -> Result<MultiplexClientImpl> {
        let client = self.duplex_client()?;
        Ok(MultiplexClientImpl { config, client })
    }
}

type Stream = io::BufStream<TcpStream>;
type Connection<'a> = netbench::Driver<'a, duplex::Connection<Stream>>;
type MultiplexConnection<'a> = netbench::Driver<'a, multiplex::Connection<Stream>>;

#[derive(Debug)]
struct ClientImpl {
    id: u64,
    nagle: bool,
    rx_buffer: usize,
    tx_buffer: usize,
}

impl ClientImpl {
    fn id(&mut self) -> u64 {
        let id = self.id;
        self.id = id + 1;
        id
    }
}

impl<'a> netbench::client::Client<'a> for ClientImpl {
    type Connect = Pin<Box<dyn Future<Output = Result<Self::Connection>> + 'a>>;
    type Connection = Connection<'a>;

    fn connect(
        &mut self,
        addr: SocketAddr,
        _server_name: &str,
        server_conn_id: u64,
        scenario: &'a Arc<scenario::Connection>,
    ) -> Self::Connect {
        let id = self.id();
        let rx_buffer = self.rx_buffer;
        let tx_buffer = self.tx_buffer;
        let nagle = self.nagle;

        let fut = async move {
            let conn = TcpStream::connect(addr).await?;

            if !nagle {
                let _ = conn.set_nodelay(true);
            }

            let conn = io::BufStream::with_capacity(rx_buffer, tx_buffer, conn);
            let mut conn = Box::pin(conn);

            // tell the server which connection ID to use
            conn.write_u64(server_conn_id).await?;

            let conn = duplex::Connection::new(id, conn);
            let conn: Self::Connection = netbench::Driver::new(scenario, conn);

            Result::Ok(conn)
        };

        Box::pin(fut)
    }
}
#[derive(Debug)]
struct MultiplexClientImpl {
    config: multiplex::Config,
    client: ClientImpl,
}

impl<'a> netbench::client::Client<'a> for MultiplexClientImpl {
    type Connect = Pin<Box<dyn Future<Output = Result<Self::Connection>> + 'a>>;
    type Connection = MultiplexConnection<'a>;

    fn connect(
        &mut self,
        addr: SocketAddr,
        _server_name: &str,
        server_conn_id: u64,
        scenario: &'a Arc<scenario::Connection>,
    ) -> Self::Connect {
        let id = self.client.id();
        let config = self.config.clone();
        let rx_buffer = self.client.rx_buffer;
        let tx_buffer = self.client.tx_buffer;
        let nagle = self.client.nagle;

        let fut = async move {
            let conn = TcpStream::connect(addr).await?;

            if !nagle {
                let _ = conn.set_nodelay(true);
            }

            let conn = io::BufStream::with_capacity(rx_buffer, tx_buffer, conn);
            let mut conn = Box::pin(conn);

            // tell the server which connection ID to use
            conn.write_u64(server_conn_id).await?;

            let conn = multiplex::Connection::new(id, conn, config);
            let conn: Self::Connection = netbench::Driver::new(scenario, conn);

            Result::Ok(conn)
        };

        Box::pin(fut)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{multiplex, scenario, Result, Timer};
use netbench_driver::Allocator;
use std::{collections::HashSet, sync::Arc};
use structopt::StructOpt;
use tokio::{
    io::{self, AsyncReadExt},
    net::{TcpListener, TcpStream},
    spawn,
};

#[global_allocator]
static ALLOCATOR: Allocator = Allocator::new();

#[tokio::main(flavor = "current_thread")]
async fn main() -> Result<()> {
    Server::from_args().run().await
}

#[derive(Debug, StructOpt)]
pub struct Server {
    #[structopt(flatten)]
    opts: netbench_driver::Server,
}

impl Server {
    pub async fn run(&self) -> Result<()> {
        let scenario = self.opts.scenario();
        let buffer = (*self.opts.rx_buffer as usize, *self.opts.tx_buffer as usize);

        let server = self.server().await?;
        let trace = self.opts.trace();
        let config = self.opts.multiplex();

        let mut conn_id = 0;
        loop {
            let (connection, _addr) = server.accept().await?;

            if !self.opts.nagle {
                let _ = connection.set_nodelay(true);
            }

            let scenario = scenario.clone();
            let id = conn_id;
            conn_id += 1;
            let trace = trace.clone();
            let config = config.clone();
            spawn(async move {
                if let Err(err) =
                    handle_connection(connection, id, scenario, trace, config, buffer).await
                {
                    eprintln!("error: {err}");
                }
            });
        }

        async fn handle_connection(
            connection: TcpStream,
            conn_id: u64,
            scenario: Arc<scenario::Server>,
            mut trace: impl netbench::Trace,
            config: Option<multiplex::Config>,
            (rx_buffer, tx_buffer): (usize, usize),
        ) -> Result<()> {
            let mut timer = netbench::timer::Tokio::default();

            // TCP servers don't have a connect time
            trace.connect(timer.now(), conn_id, Default::default());

            let connection = io::BufStream::with_capacity(rx_buffer, tx_buffer, connection);
            let mut connection = Box::pin(connection);

            let server_idx = connection.read_u64().await?;
            let scenario = scenario
                .connections
                .get(server_idx as usize)
                .ok_or("invalid connection id")?;

            let mut checkpoints = HashSet::new();

            if let Some(config) = config {
                let conn = netbench::Driver::new(
                    scenario,
                    netbench::multiplex::Connection::new(conn_id, connection, config),
                );

                conn.run(&mut trace, &mut checkpoints, &mut timer).await?;
            } else {
                let conn = netbench::Driver::new(
                    scenario,
                    netbench::duplex::Connection::new(conn_id, connection),
                );

                conn.run(&mut trace, &mut checkpoints, &mut timer).await?;
            }

            Ok(())
        }
    }

    async fn server(&self) -> Result<TcpListener> {
        let server = TcpListener::bind((self.opts.ip, self.opts.port)).await?;
        Ok(server)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench::{units::*, Error, Result};
use std::collections::BTreeMap;

#[derive(Debug, Default)]
pub struct Registry {
    definitions: BTreeMap<&'static str, Definition>,
}

impl Registry {
    pub fn define<T>(&mut self, name: &'static str, docs: &'static [&'static str], default: &T)
    where
        T: TryFromValue,
    {
        self.definitions
            .insert(name, Definition::new::<T>(docs, default.display()));
    }

    pub fn clap_args(&self) -> impl Iterator<Item = clap::Arg> + '_ {
        self.definitions.iter().map(|(name, def)| {
            clap::Arg::with_name(name)
                .long(name)
                .value_name(def.value_name)
                .help(def.help)
                .default_value(&def.default)
                .takes_value(def.takes_value)
                .long_help(&def.long_help)
        })
    }

    pub fn load_overrides(&self, matches: &clap::ArgMatches) -> Overrides {
        let mut overrides = Overrides::default();
        for (name, _) in self.definitions.iter() {
            if let Some(value) = matches.value_of(name) {
                overrides
                    .values
                    .insert(name, Override::String(value.to_string()));
            } else if matches.is_present(name) {
                overrides.values.insert(name, Override::Enabled);
            }
        }
        overrides
    }
}

#[derive(Default)]
pub struct Overrides {
    values: BTreeMap<&'static str, Override>,
    errors: BTreeMap<&'static str, Error>,
}

impl Overrides {
    pub fn resolve<T>(&mut self, name: &'static str, default: T) -> T
    where
        T: TryFromValue,
    {
        if let Some(value) = self.values.get(name) {
            match T::try_from_value(value) {
                Ok(value) => {
                    return value;
                }
                Err(err) => {
                    self.errors.insert(name, err);
                }
            }
        }

        default
    }

    pub fn errors(&self) -> impl Iterator<Item = String> + '_ {
        self.errors
            .iter()
            .map(|(name, error)| format!("{name}: {error}\n"))
    }
}

#[derive(Debug)]
struct Definition {
    help: &'static str,
    long_help: String,
    default: String,
    value_name: &'static str,
    takes_value: bool,
}

impl Definition {
    fn new<T: TryFromValue>(docs: &[&'static str], default: String) -> Self {
        Self {
            help: docs[0],
            long_help: docs.iter().map(|v| v.trim()).collect::<Vec<_>>().join("\n"),
            default,
            value_name: T::VALUE_NAME,
            takes_value: T::TAKES_VALUE,
        }
    }
}

#[derive(Debug)]
pub enum Override {
    Enabled,
    String(String),
}

pub trait TryFromValue: Sized {
    const VALUE_NAME: &'static str;
    const TAKES_VALUE: bool = true;

    fn try_from_value(value: &Override) -> Result<Self>;
    fn display(&self) -> String;
}

impl TryFromValue for bool {
    const VALUE_NAME: &'static str = "BOOL";
    const TAKES_VALUE: bool = false;

    fn try_from_value(value: &Override) -> Result<Self> {
        match value {
            Override::Enabled => Ok(true),
            Override::String(v) => match v.as_str() {
                "true" | "TRUE" | "1" | "yes" | "YES" => Ok(true),
                "false" | "FALSE" | "0" | "no" | "NO" => Ok(false),
                _ => Err(format!("invalid bool: {v:?}").into()),
            },
        }
    }

    fn display(&self) -> String {
        self.to_string()
    }
}

impl TryFromValue for u64 {
    const VALUE_NAME: &'static str = "COUNT";

    fn try_from_value(value: &Override) -> Result<Self> {
        match value {
            Override::Enabled => Err("missing value".into()),
            Override::String(v) => Ok(v.parse()?),
        }
    }

    fn display(&self) -> String {
        self.to_string()
    }
}

impl TryFromValue for Duration {
    const VALUE_NAME: &'static str = "TIME";

    fn try_from_value(value: &Override) -> Result<Self> {
        match value {
            Override::Enabled => Err("missing value".into()),
            Override::String(v) => {
                let v: humantime::Duration = v.parse()?;
                Ok(*v)
            }
        }
    }

    fn display(&self) -> String {
        if *self == Self::ZERO {
            return "0s".to_owned();
        }
        format!("{self:?}")
    }
}

impl<T: TryFromValue> TryFromValue for Option<T> {
    const VALUE_NAME: &'static str = T::VALUE_NAME;

    fn try_from_value(value: &Override) -> Result<Self> {
        if matches!(value, Override::String(v) if v == "NONE") {
            return Ok(None);
        }

        T::try_from_value(value).map(Some)
    }

    fn display(&self) -> String {
        if let Some(value) = self.as_ref() {
            value.display()
        } else {
            "NONE".to_owned()
        }
    }
}

macro_rules! try_from_value {
    ($name:ty, $value_name:literal) => {
        impl TryFromValue for $name {
            const VALUE_NAME: &'static str = $value_name;

            fn try_from_value(value: &Override) -> Result<Self> {
                match value {
                    Override::Enabled => Err("missing value".into()),
                    Override::String(v) => Ok(v.parse()?),
                }
            }

            fn display(&self) -> String {
                self.to_string()
            }
        }
    };
}

try_from_value!(Byte, "BYTES");
try_from_value!(Rate, "RATE");

impl<T: TryFromValue> TryFromValue for Vec<T> {
    const VALUE_NAME: &'static str = T::VALUE_NAME;

    fn try_from_value(value: &Override) -> Result<Self> {
        let mut out = vec![];
        if let Override::String(v) = value {
            for value in v.split(',') {
                let value = value.trim();
                let value = Override::String(value.to_owned());
                out.push(T::try_from_value(&value)?);
            }
        }
        Ok(out)
    }

    fn display(&self) -> String {
        let mut out = vec![];
        for value in self {
            out.push(value.display());
        }
        out.join(",")
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench_scenarios::prelude::*;

config!({
    /// The number of separate connections to create
    let connections: u64 = 1000;
});

pub fn scenario(config: Config) -> Scenario {
    let Config { connections } = config;

    Scenario::build(|scenario| {
        let server = scenario.create_server();

        scenario.create_client(|client| {
            for _ in 0..connections {
                client.connect_to(&server, |conn| {
                    conn.open_bidirectional_stream(
                        |local| {
                            local.send(1.bytes());
                            local.receive(1.bytes());
                        },
                        |remote| {
                            remote.receive(1.bytes());
                            remote.send(1.bytes());
                        },
                    );
                });
            }
        });
    })
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use clap::{App, Arg};
pub use netbench::Result;
use std::path::Path;

pub mod config;
pub mod prelude;

pub trait Configs {
    fn registry() -> config::Registry;
    fn new(overrides: &mut config::Overrides) -> Self;
    fn write(self, out_dir: &Path) -> Result<()>;
}

#[doc(hidden)]
pub fn main<C: Configs>() -> Result<()> {
    let app = App::new("netbench scenarios")
        .after_help(LONG_ABOUT.trim())
        .arg(
            Arg::with_name("out_dir")
                .value_name("OUT_DIR")
                .default_value("target/netbench")
                .takes_value(true),
        );

    let map = C::registry();
    let args = map.clap_args().collect::<Vec<_>>();
    let matches = app.args(&args).get_matches();
    let mut overrides = map.load_overrides(&matches);

    let configs = C::new(&mut overrides);

    let mut has_error = false;
    for error in overrides.errors() {
        eprintln!("{error}");
        has_error = true;
    }

    if has_error {
        return Err("bailing due to errors".into());
    }

    let out_dir = matches.value_of("out_dir").unwrap();
    let out_dir = Path::new(out_dir);

    std::fs::create_dir_all(out_dir)?;
    configs.write(out_dir)?;

    Ok(())
}

const LONG_ABOUT: &str = r#"
FORMATS:
    BYTES
        42b         ->    42 bits
        42          ->    42 bytes
        42B         ->    42 bytes
        42K         ->    42000 bytes
        42Kb        ->    42000 bits
        42KB        ->    42000 bytes
        42KiB       ->    43008 bytes

    COUNT
        42          ->    42 units

    RATE
        42bps       ->    42 bits per second
        42Mbps      ->    42 megabits per second
        42MBps      ->    42 megabytes per second
        42MiBps     ->    42 mebibytes per second
        42MB/50ms   ->    42 megabytes per 50 milliseconds

    TIME
        42ms         ->    42 milliseconds
        42s          ->    42 seconds
        1s42ms       ->    1 second + 42 milliseconds
"#;

#[macro_export]
macro_rules! scenarios {
    ($($name:ident),* $(,)?) => {
        $(
            mod $name;
        )*

        #[derive(Debug)]
        pub struct Configs {
            $(
                $name: $name::Config,
            )*
        }

        impl $crate::Configs for Configs {
            fn registry() -> $crate::config::Registry {
                let mut registry = $crate::config::Registry::default();
                $(
                    $name::Config::register(&mut registry);
                )*
                registry
            }

            fn new(overrides: &mut $crate::config::Overrides) -> Self {
                Self {
                    $(
                        $name: $name::Config::new(overrides),
                    )*
                }
            }

            fn write(self, out_dir: &std::path::Path) -> $crate::Result<()> {
                $({
                    let path = out_dir.join(concat!(stringify!($name), ".json"));
                    let mut f = std::fs::File::create(&path)?;
                    let s = $name::scenario(self.$name);
                    s.write(&mut f)?;
                    eprintln!("created: {}", path.display());
                })*

                Ok(())
            }
        }

        fn main() -> $crate::Result<()> {
            $crate::main::<Configs>()
        }
    };
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

netbench_scenarios::scenarios!(connect, ping, request_response);


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench_scenarios::prelude::*;

config!({
    /// The number of concurrent connections to create
    let connections: u64 = 1;

    /// The number of concurrent streams to ping on
    let streams: u64 = 1;

    /// The amount of time to spend pinging for each size
    let time: Duration = 15.seconds();

    /// The amount of data to send in each ping
    let size: Vec<Byte> = vec![
        1000.bytes(),
        10_000.bytes(),
        100_000.bytes(),
        1_000_000.bytes(),
    ];
});

pub fn scenario(config: Config) -> Scenario {
    let Config {
        connections,
        streams,
        size: sizes,
        time,
    } = config;

    Scenario::build(|scenario| {
        let server = scenario.create_server();

        scenario.create_client(|client| {
            for size in sizes.iter().copied() {
                let ping = format!("ping {size}");
                let pong = format!("pong {size}");
                client.scope(|client| {
                    for _ in 0..connections {
                        client.spawn(|client| {
                            client.connect_to(&server, |conn| {
                                conn.scope(|conn| {
                                    for _ in 0..streams {
                                        conn.spawn(|conn| {
                                            conn.open_bidirectional_stream(
                                                |local| {
                                                    local.iterate(time, |local| {
                                                        local.profile(&ping, |local| {
                                                            local.send(size);
                                                            local.receive(size);
                                                        });
                                                    });
                                                },
                                                |remote| {
                                                    remote.iterate(time, |remote| {
                                                        remote.profile(&pong, |remote| {
                                                            remote.receive(size);
                                                            remote.send(size);
                                                        });
                                                    });
                                                },
                                            );
                                        });
                                    }
                                });
                            });
                        });
                    }
                });
            }
        });
    })
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub use crate::config;
pub use netbench::{
    scenario::{builder, Scenario},
    units::*,
    Result,
};

#[macro_export]
macro_rules! config {
    ({$($(#[doc = $doc:literal])* let $name:ident: $ty:ty = $default:expr;)*}) => {
        #[derive(Debug)]
        pub struct Config {
            $(
                $(#[doc = $doc])*
                $name: $ty,
            )*
        }

        impl Config {
            pub fn register(registry: &mut $crate::config::Registry) {
                $(
                    {
                        let name = concat!(module_path!(), ".", stringify!($name)).split("::").last().unwrap();
                        static DOCS: &[&str] = &[$($doc),*];
                        let default: $ty = $default;
                        registry.define(name, DOCS, &default)
                    }
                )*
            }

            pub fn new(overrides: &mut $crate::config::Overrides) -> Config {
                Config {
                    $(
                        $name: {
                            let name = concat!(module_path!(), ".", stringify!($name)).split("::").last().unwrap();
                            overrides.resolve(name, $default)
                        },
                    )*
                }
            }
        }
    };
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use netbench_scenarios::prelude::*;

config!({
    /// The size of the client's request to the server
    let request_size: Byte = 1.kilobytes();

    /// The size of the server's response to the client
    let response_size: Byte = 10.megabytes();

    /// How long the server will take to respond to the request
    let response_delay: Duration = 0.seconds();

    /// The number of requests to make
    let count: u64 = 1;

    /// The number of separate connections to create
    let connections: u64 = 1;

    /// Specifies if the requests should be performed in parallel
    let parallel: bool = false;

    /// The rate at which the client sends data
    let client_send_rate: Option<Rate> = None;

    /// The rate at which the client receives data
    let client_receive_rate: Option<Rate> = None;

    /// The rate at which the server sends data
    let server_send_rate: Option<Rate> = None;

    /// The rate at which the server receives data
    let server_receive_rate: Option<Rate> = None;

    /// The number of bytes that must be received before the next request
    let response_unblock: Byte = 0.bytes();
});

pub fn scenario(config: Config) -> Scenario {
    let Config {
        request_size,
        response_size,
        count,
        connections,
        parallel,
        client_send_rate,
        client_receive_rate,
        server_send_rate,
        server_receive_rate,
        response_delay,
        response_unblock,
    } = config;
    let response_unblock = response_unblock.min(response_size);

    type Checkpoint = Option<
        builder::checkpoint::Checkpoint<builder::Client, builder::Local, builder::checkpoint::Park>,
    >;

    let request = |conn: &mut builder::connection::Builder<builder::Client>,
                   checkpoint: &mut Checkpoint| {
        let (park, unpark) = conn.checkpoint();

        if let Some(park) = checkpoint.take() {
            conn.park(park);
        }

        conn.open_bidirectional_stream(
            |local| {
                if let Some(rate) = client_send_rate {
                    local.set_send_rate(rate);
                }
                if let Some(rate) = client_receive_rate {
                    local.set_receive_rate(rate);
                }
                local.send(request_size);

                if *response_unblock > 0 {
                    local.receive(response_unblock);
                    local.unpark(unpark);
                    local.receive(response_size - response_unblock);
                } else {
                    local.receive(response_size);
                }
            },
            |remote| {
                if let Some(rate) = server_send_rate {
                    remote.set_send_rate(rate);
                }
                if let Some(rate) = server_receive_rate {
                    remote.set_receive_rate(rate);
                }
                remote.receive(request_size);

                if response_delay != Duration::ZERO {
                    remote.sleep(response_delay);
                }

                remote.send(response_size);
            },
        );

        if *response_unblock > 0 {
            *checkpoint = Some(park)
        }
    };

    Scenario::build(|scenario| {
        let server = scenario.create_server();

        scenario.create_client(|client| {
            for _ in 0..connections {
                client.connect_to(&server, |conn| {
                    if parallel {
                        conn.scope(|scope| {
                            let mut prev_checkpoint = None;
                            for _ in 0..count {
                                scope.spawn(|conn| {
                                    request(conn, &mut prev_checkpoint);
                                });
                            }
                        });
                    } else {
                        for _ in 0..count {
                            request(conn, &mut None);
                        }
                    }
                });
            }
        });
    })
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    connection::{self, Connection},
    provider::*,
};
use core::{
    fmt,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use s2n_quic_transport::endpoint::{connect, handle::Connector};

mod builder;
mod providers;

pub use builder::*;
pub use connect::Connect;
pub use providers::*;

/// A QUIC client endpoint, capable of opening connections
#[derive(Clone)]
pub struct Client {
    connector: Connector,
    local_addr: s2n_quic_core::inet::SocketAddress,
}

impl fmt::Debug for Client {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Client")
            .field("local_addr", &self.local_addr().ok())
            .finish()
    }
}

impl Client {
    /// Starts listening on the provided socket
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use std::error::Error;
    /// # use s2n_quic::Client;
    /// #
    /// # fn main() -> Result<(), Box<dyn Error>> {
    /// let Client = Client::bind("0.0.0.0:0")?;
    /// #
    /// #    Ok(())
    /// # }
    /// ```
    pub fn bind<T>(socket: T) -> Result<Self, StartError>
    where
        T: io::TryInto,
    {
        let client = Self::builder()
            .with_io(socket)
            .map_err(StartError::new)?
            .start()?;
        Ok(client)
    }

    /// Returns a [`Builder`] which is able to configure the [`Client`] components.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use std::error::Error;
    /// use std::path::Path;
    /// use s2n_quic::Client;
    ///
    /// #
    /// # fn main() -> Result<(), Box<dyn Error>> {
    /// let client = Client::builder()
    ///     .with_tls(Path::new("./certs/cert.pem"))?
    ///     .with_io("0.0.0.0:0")?
    ///     .start()?;
    /// #
    /// #    Ok(())
    /// # }
    /// ```
    pub fn builder() -> Builder<impl ClientProviders> {
        Builder::default()
    }

    /// Establishes a connection to the specified endpoint
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use std::error::Error;
    /// use s2n_quic::Client;
    /// use std::{net::SocketAddr, path::Path};
    ///
    /// # async fn connect() -> Result<(), Box<dyn Error>> {
    /// let client = Client::builder()
    ///     .with_tls(Path::new("./certs/cert.pem"))?
    ///     .with_io("0.0.0.0:0")?
    ///     .start()?;
    ///
    /// let addr: SocketAddr = "127.0.0.1:443".parse()?;
    /// let connection = client.connect(addr.into()).await?;
    /// #
    /// #    Ok(())
    /// # }
    /// ```
    pub fn connect(&self, connect: Connect) -> ConnectionAttempt {
        let attempt = self.connector.connect(connect);
        ConnectionAttempt(attempt)
    }

    /// Wait for the client endpoint to finish handling all outstanding connections
    ///
    /// Notifies the endpoint of application interest in closing the endpoint. The
    /// call waits for **all** outstanding connections to finish before returning.
    ///
    /// Note: The endpoint will continue to accept new connection attempts. If there
    /// are other client handles with active connections, then this call will never
    /// return.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use std::error::Error;
    /// use s2n_quic::Client;
    /// use std::{net::SocketAddr, path::Path};
    ///
    /// # async fn connect() -> Result<(), Box<dyn Error>> {
    /// let mut client = Client::builder()
    ///     .with_tls(Path::new("./certs/cert.pem"))?
    ///     .with_io("0.0.0.0:0")?
    ///     .start()?;
    ///
    /// let addr: SocketAddr = "127.0.0.1:443".parse()?;
    /// let connection = client.connect(addr.into()).await?;
    ///
    /// client.wait_idle().await?;
    /// #
    /// #    Ok(())
    /// # }
    /// ```
    pub async fn wait_idle(&mut self) -> Result<(), connection::Error> {
        futures::future::poll_fn(|cx| self.connector.poll_close(cx)).await
    }

    /// Returns the local address that this listener is bound to.
    ///
    /// This can be useful, for example, when binding to port `0` to figure out which
    /// port was actually bound.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use std::error::Error;
    /// # use s2n_quic::Client;
    /// #
    /// # fn main() -> Result<(), Box<dyn Error>> {
    /// let client = Client::bind("0.0.0.0:0")?;
    ///
    /// let local_addr = client.local_addr()?;
    /// assert_ne!(local_addr.port(), 0);
    /// #    Ok(())
    /// # }
    /// ```
    pub fn local_addr(&self) -> Result<std::net::SocketAddr, std::io::Error> {
        Ok(self.local_addr.into())
    }
}

#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct ConnectionAttempt(connect::Attempt);

impl Future for ConnectionAttempt {
    type Output = Result<Connection, connection::Error>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match Pin::new(&mut self.0).poll(cx) {
            Poll::Ready(Ok(conn)) => Poll::Ready(Ok(Connection::new(conn))),
            Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
            Poll::Pending => Poll::Pending,
        }
    }
}