ctf-ad-mirror-tool-rust/examples/benchmark.rs

//! Reproducible loopback comparison. Not a claim about production networks.
use anyhow::Result;
use mirror_tool::{
    config::{NetworkConfig, ReplicaPolicy, ServiceConfig, Team},
    engine::Engine,
};
use serde_json::json;
use std::{net::SocketAddr, time::Instant};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
};

async fn transfer(addr: SocketAddr, megabytes: usize) -> Result<f64> {
    let stream = TcpStream::connect(addr).await?;
    stream.set_nodelay(true)?;
    let (mut read, mut write) = stream.into_split();
    let start = Instant::now();
    let sender = async {
        let data = vec![0x5au8; 65536];
        for _ in 0..megabytes * 16 {
            write.write_all(&data).await?;
        }
        write.shutdown().await
    };
    let receiver = async {
        let mut data = vec![0; 65536];
        let mut total = 0;
        loop {
            let n = read.read(&mut data).await?;
            if n == 0 {
                break;
            }
            anyhow::ensure!(data[..n].iter().all(|b| *b == 0x5a), "byte corruption");
            total += n;
        }
        anyhow::ensure!(total == megabytes * 1048576, "truncated stream");
        Ok::<_, anyhow::Error>(())
    };
    tokio::try_join!(
        async {
            sender.await?;
            Ok::<_, anyhow::Error>(())
        },
        receiver
    )?;
    Ok((megabytes as f64 * 2.0) / start.elapsed().as_secs_f64())
}
async fn latency(addr: SocketAddr) -> Result<(f64, f64)> {
    let mut stream = TcpStream::connect(addr).await?;
    stream.set_nodelay(true)?;
    let mut times = vec![];
    let mut reply = [0u8; 32];
    for _ in 0..500 {
        let t = Instant::now();
        stream.write_all(&[42u8; 32]).await?;
        stream.read_exact(&mut reply).await?;
        anyhow::ensure!(reply == [42; 32], "latency reply corruption");
        times.push(t.elapsed().as_secs_f64() * 1e6);
    }
    stream.shutdown().await?;
    times.sort_by(f64::total_cmp);
    Ok((times[250], times[475]))
}
fn echo(listener: TcpListener) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        loop {
            let (stream, _) = listener.accept().await.unwrap();
            stream.set_nodelay(true).unwrap();
            tokio::spawn(async move {
                let (mut read, mut write) = stream.into_split();
                let _ = tokio::io::copy(&mut read, &mut write).await;
                let _ = write.shutdown().await;
            });
        }
    })
}
#[tokio::main]
async fn main() -> Result<()> {
    let megabytes = std::env::args()
        .nth(1)
        .and_then(|v| v.parse().ok())
        .unwrap_or(64);
    anyhow::ensure!(
        (1..=4096).contains(&megabytes),
        "payload must be 1..4096 MiB"
    );
    let mode = std::env::args().nth(2).unwrap_or_else(|| "none".into());
    anyhow::ensure!(
        matches!(mode.as_str(), "none" | "strict" | "isolate"),
        "mode: none, strict or isolate"
    );
    let listener = TcpListener::bind("127.0.0.1:0").await?;
    let direct = listener.local_addr()?;
    let upstream = echo(listener);
    let replica = if mode != "none" {
        Some(echo(
            TcpListener::bind((std::net::Ipv6Addr::LOCALHOST, direct.port())).await?,
        ))
    } else {
        None
    };
    let mut teams = vec![Team {
        hidden: false,
        id: "local".into(),
        name: "Loopback".into(),
        address: direct.ip(),
    }];
    if replica.is_some() {
        teams.push(Team {
            hidden: false,
            id: "replica".into(),
            name: "Replica".into(),
            address: std::net::Ipv6Addr::LOCALHOST.into(),
        });
    }
    let reserve = TcpListener::bind("127.0.0.1:0").await?;
    let proxy = reserve.local_addr()?;
    drop(reserve);
    let engine = Engine::new(NetworkConfig::default());
    let service = ServiceConfig {
        custom: false,
        custom_ip: None,
        id: "bench".into(),
        name: "Benchmark".into(),
        listen: proxy,
        upstream_port: direct.port(),
        primary: Some("local".into()),
        replicas: if replica.is_some() {
            vec!["replica".into()]
        } else {
            vec![]
        },
        enabled: true,
        auto_switch: false,
        replica_policy: if mode == "strict" {
            ReplicaPolicy::Strict
        } else {
            ReplicaPolicy::Isolate
        },
    };
    engine.upsert(&service, &teams).await?;
    transfer(direct, 4).await?;
    transfer(proxy, 4).await?;
    let mut direct_runs = vec![];
    let mut proxy_runs = vec![];
    for round in 0..3 {
        if round % 2 == 0 {
            direct_runs.push(transfer(direct, megabytes).await?);
            proxy_runs.push(transfer(proxy, megabytes).await?);
        } else {
            proxy_runs.push(transfer(proxy, megabytes).await?);
            direct_runs.push(transfer(direct, megabytes).await?);
        }
    }
    let direct_latency = latency(direct).await?;
    let proxy_latency = latency(proxy).await?;
    engine.shutdown().await;
    upstream.abort();
    let _ = upstream.await;
    if let Some(replica) = replica {
        replica.abort();
        let _ = replica.await;
    }
    println!(
        "{}",
        serde_json::to_string_pretty(
            &json!({"platform":std::env::consts::OS,"arch":std::env::consts::ARCH,"mode":mode,"payload_mib_per_run":megabytes,"throughput_unit":"MiB/s aggregate upload+echo download (replica bytes excluded)","direct":direct_runs,"proxy":proxy_runs,"direct_latency_us":{"p50":direct_latency.0,"p95":direct_latency.1},"proxy_latency_us":{"p50":proxy_latency.0,"p95":proxy_latency.1},"byte_integrity":"verified every primary reply byte", "stats":engine.snapshots()})
        )?
    );
    Ok(())
}