rok-cli 0.6.1

Developer CLI for rok-based Axum applications
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

//! `rok bench` — built-in HTTP benchmarking via oha/wrk integration (M7.10).

pub fn run(url: &str, requests: Option<u64>, duration: Option<u64>, concurrency: Option<u64>) -> anyhow::Result<()> {
    let n = requests.unwrap_or(1000);
    let c = concurrency.unwrap_or(10);
    let dur = duration;

    // Prefer oha, fall back to wrk, then ab
    if let Some(tool) = find_bench_tool() {
        println!("rok bench — using {tool}");
        run_tool(&tool, url, n, c, dur)
    } else {
        println!("rok bench — no benchmark tool found.");
        println!();
        println!("Install one of the following:");
        println!("  oha     cargo install oha");
        println!("  wrk     (brew install wrk / apt install wrk)");
        println!("  ab      (brew install httpd / apt install apache2-utils)");
        println!();
        println!("Then re-run: rok bench {url}");
        Ok(())
    }
}

fn find_bench_tool() -> Option<String> {
    for tool in ["oha", "wrk", "ab"] {
        let found = std::process::Command::new(tool)
            .arg("--version")
            .output()
            .map(|o| o.status.success())
            .unwrap_or(false);
        // For ab, version flag is different
        let found = found || std::process::Command::new(tool)
            .arg("-V")
            .output()
            .map(|o| o.status.success() || !o.stderr.is_empty())
            .unwrap_or(false);
        if found {
            return Some(tool.to_string());
        }
    }
    None
}

fn run_tool(tool: &str, url: &str, n: u64, c: u64, dur: Option<u64>) -> anyhow::Result<()> {
    let mut cmd = std::process::Command::new(tool);
    match tool {
        "oha" => {
            cmd.arg("-n").arg(n.to_string());
            cmd.arg("-c").arg(c.to_string());
            if let Some(d) = dur {
                cmd.arg("-z").arg(format!("{d}s"));
            }
            cmd.arg(url);
        }
        "wrk" => {
            cmd.arg("-t").arg((c / 4).max(1).to_string());
            cmd.arg("-c").arg(c.to_string());
            if let Some(d) = dur {
                cmd.arg("-d").arg(format!("{d}s"));
            } else {
                cmd.arg("-d").arg("10s");
            }
            cmd.arg(url);
        }
        "ab" => {
            cmd.arg("-n").arg(n.to_string());
            cmd.arg("-c").arg(c.to_string());
            cmd.arg(url);
        }
        _ => {}
    }

    let status = cmd.status()?;
    if !status.success() {
        anyhow::bail!("{tool} exited with status {status}");
    }
    Ok(())
}