harmony/iobench/src/main.rs

use std::fs;
use std::io::{self, Write};
use std::process::{Command, Stdio};
use std::thread;
use std::time::Duration;

use chrono::Local;
use clap::Parser;
use serde::{Deserialize, Serialize};

/// A simple yet powerful I/O benchmarking tool using fio.
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Target for the benchmark.
    /// Formats:
    ///   - localhost (default)
    ///   - ssh/{user}@{host}
    ///   - ssh/{user}@{host}:{port}
    ///   - k8s/{namespace}/{pod}
    #[arg(short, long, default_value = "localhost")]
    target: String,

    #[arg(short, long, default_value = ".")]
    benchmark_dir: String,

    /// Comma-separated list of tests to run.
    /// Available tests: read, write, randread, randwrite,
    /// multiread, multiwrite, multirandread, multirandwrite.
    #[arg(long, default_value = "read,write,randread,randwrite,multiread,multiwrite,multirandread,multirandwrite")]
    tests: String,

    /// Duration of each test in seconds.
    #[arg(long, default_value_t = 15)]
    duration: u64,

    /// Output directory for results.
    /// Defaults to ./iobench-{current_datetime}.
    #[arg(long)]
    output_dir: Option<String>,

    /// The size of the test file for fio.
    #[arg(long, default_value = "1G")]
    size: String,

    /// The block size for I/O operations.
    #[arg(long, default_value = "4k")]
    block_size: String,
}

#[derive(Debug, Serialize, Deserialize)]
struct FioOutput {
    jobs: Vec<FioJobResult>,
}

#[derive(Debug, Serialize, Deserialize)]
struct FioJobResult {
    jobname: String,
    read: FioMetrics,
    write: FioMetrics,
}

#[derive(Debug, Serialize, Deserialize)]
struct FioMetrics {
    bw: f64,
    iops: f64,
    clat_ns: LatencyMetrics,
}

#[derive(Debug, Serialize, Deserialize)]
struct LatencyMetrics {
    mean: f64,
    stddev: f64,
}

#[derive(Debug, Serialize)]
struct BenchmarkResult {
    test_name: String,
    iops: f64,
    bandwidth_kibps: f64,
    latency_mean_ms: f64,
    latency_stddev_ms: f64,
}

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

    let output_dir = args.output_dir.unwrap_or_else(|| {
        format!("./iobench-{}", Local::now().format("%Y-%m-%d-%H%M%S"))
    });
    fs::create_dir_all(&output_dir)?;

    let tests_to_run: Vec<&str> = args.tests.split(',').collect();
    let mut results = Vec::new();

    for test in tests_to_run {
        println!("--------------------------------------------------");
        println!("Running test: {}", test);

        let (rw, numjobs) = match test {
            "read" => ("read", 1),
            "write" => ("write", 1),
            "randread" => ("randread", 1),
            "randwrite" => ("randwrite", 1),
            "multiread" => ("read", 4),
            "multiwrite" => ("write", 4),
            "multirandread" => ("randread", 4),
            "multirandwrite" => ("randwrite", 4),
            _ => {
                eprintln!("Unknown test: {}. Skipping.", test);
                continue;
            }
        };

        let test_name = format!("{}-{}-sync-test", test, args.block_size);
        let fio_command = format!(
            "fio --filename={}/iobench_testfile --direct=1 --fsync=1 --rw={} --bs={} --numjobs={} --iodepth=1 --runtime={} --time_based --group_reporting --name={} --size={} --output-format=json",
            args.benchmark_dir, rw, args.block_size, numjobs, args.duration, test_name, args.size
        );

        println!("Executing command:\n{}\n", fio_command);

        let output = match run_command(&args.target, &fio_command) {
            Ok(out) => out,
            Err(e) => {
                eprintln!("Failed to execute command for test {}: {}", test, e);
                continue;
            }
        };


        let result = parse_fio_output(&output, &test_name, rw);
        // TODO store raw fio output and print it
        match result {
            Ok(res) => {
                results.push(res);
            }
            Err(e) => {
                eprintln!("Error parsing fio output for test {}: {}", test, e);
                eprintln!("Raw output:\n{}", output);
            }
        }

        println!("{output}");
        println!("Test {} completed.", test);
        // A brief pause to let the system settle before the next test.
        thread::sleep(Duration::from_secs(2));
    }
    
    // Cleanup the test file on the target
    println!("--------------------------------------------------");
    println!("Cleaning up test file on target...");
    let cleanup_command = "rm -f ./iobench_testfile";
    if let Err(e) = run_command(&args.target, cleanup_command) {
        eprintln!("Warning: Failed to clean up test file on target: {}", e);
    } else {
        println!("Cleanup successful.");
    }


    if results.is_empty() {
        println!("\nNo benchmark results to display.");
        return Ok(());
    }

    // Output results to a CSV file for easy analysis
    let csv_path = format!("{}/summary.csv", output_dir);
    let mut wtr = csv::Writer::from_path(&csv_path)?;
    for result in &results {
        wtr.serialize(result)?;
    }
    wtr.flush()?;

    println!("\nBenchmark summary saved to {}", csv_path);
    println!("\n--- Benchmark Results Summary ---");
    println!("{:<25} {:>10} {:>18} {:>20} {:>22}", "Test Name", "IOPS", "Bandwidth (KiB/s)", "Latency Mean (ms)", "Latency StdDev (ms)");
    println!("{:-<98}", "");
    for result in results {
        println!("{:<25} {:>10.2} {:>18.2} {:>20.4} {:>22.4}", result.test_name, result.iops, result.bandwidth_kibps, result.latency_mean_ms, result.latency_stddev_ms);
    }

    Ok(())
}

fn run_command(target: &str, command: &str) -> io::Result<String> {
    let (program, args) = if target == "localhost" {
        ("sudo", vec!["sh".to_string(), "-c".to_string(), command.to_string()])
    } else if target.starts_with("ssh/") {
        let target_str = target.strip_prefix("ssh/").unwrap();
        let ssh_target;
        let mut ssh_args = vec!["-o".to_string(), "StrictHostKeyChecking=no".to_string()];
        let port_parts: Vec<&str> = target_str.split(':').collect();
        if port_parts.len() == 2 {
            ssh_target = port_parts[0].to_string();
            ssh_args.push("-p".to_string());
            ssh_args.push(port_parts[1].to_string());
        } else {
            ssh_target = target_str.to_string();
        }
        ssh_args.push(ssh_target);
        ssh_args.push(format!("sudo sh -c '{}'", command));
        ("ssh", ssh_args)
    } else if target.starts_with("k8s/") {
        let parts: Vec<&str> = target.strip_prefix("k8s/").unwrap().split('/').collect();
        if parts.len() != 2 {
            return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid k8s target format. Expected k8s/{namespace}/{pod}"));
        }
        let namespace = parts[0];
        let pod = parts[1];
        ("kubectl", vec!["exec".to_string(), "-n".to_string(), namespace.to_string(), pod.to_string(), "--".to_string(), "sh".to_string(), "-c".to_string(), command.to_string()])
    } else {
        return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid target format"));
    };

    let mut cmd = Command::new(program);
    cmd.args(&args);
    cmd.stdout(Stdio::piped()).stderr(Stdio::piped());

    let child = cmd.spawn()?;
    let output = child.wait_with_output()?;

    if !output.status.success() {
        eprintln!("Command failed with status: {}", output.status);
        io::stderr().write_all(&output.stderr)?;
        return Err(io::Error::new(io::ErrorKind::Other, "Command execution failed"));
    }

    String::from_utf8(output.stdout)
        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}

fn parse_fio_output(output: &str, test_name: &str, rw: &str) -> Result<BenchmarkResult, String> {
    let fio_data: FioOutput = serde_json::from_str(output)
        .map_err(|e| format!("Failed to deserialize fio JSON: {}", e))?;

    let job_result = fio_data.jobs.iter()
        .find(|j| j.jobname == test_name)
        .ok_or_else(|| format!("Could not find job result for '{}' in fio output", test_name))?;

    let metrics = if rw.contains("read") {
        &job_result.read
    } else {
        &job_result.write
    };

    Ok(BenchmarkResult {
        test_name: test_name.to_string(),
        iops: metrics.iops,
        bandwidth_kibps: metrics.bw,
        latency_mean_ms: metrics.clat_ns.mean / 1_000_000.0,
        latency_stddev_ms: metrics.clat_ns.stddev / 1_000_000.0,
    })
}