HTTP load testing tool to study Rust language.

vegeta.rs

// usage:
// $ cargo build --bin vegeta --release
// $ ./target/release/vegeta -u http://localhost:8000/foo -r 5 -d 30
// generates 5 requests per second, during 30 seconds, targeting to http://localhost:8000/foo
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use histogram::Histogram;
use std::io::{self, Write};
use anyhow::Result;
use tokio::sync::mpsc::Sender;
use clap::Parser;

const BUFFER_1: usize = 1_000;
// you might need to increase ulimit. ex.) ulimit -n 1024
const BUFFER_2: usize = 1;

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// Url of the target to attack.
    #[clap(short, long, value_parser)]
    url: String,

    /// Number of requests per second, default to 3 requests/second
    #[clap(short, long, value_parser, default_value_t = 3)]
    rate: u64,

    /// Duration(sec) of the attacking, default to 10 seconds.
    /// Thus, total number of requests could be rate x duration
    #[clap(short, long, value_parser, default_value_t = 10)]
    duration: u64,
}


#[derive(Debug)]
enum Command {
    Get { key: String },
    Response { val: u128 }, // response time of a request
}

// This consists of three tasks: Cron Job, Attackers and Collector.
// Cron Job starts Attackers every second during the specified duration, DURATION.
// Attackers generates as many requests as specified by RATE.
// Collector collects results from Attackers.
#[tokio::main()]
async fn main() -> Result<()> {
    let args = Args::parse();

    let url: String = args.url;
    let rate: u64 = args.rate;
    let duration: u64 = args.duration;

    // Attackers -> Collector
    let (tx1, mut rx1) = mpsc::channel(BUFFER_1);
    // Cron Job -> Attackers
    let (tx2, mut rx2) = mpsc::channel(BUFFER_2);

    // 1. Collector task
    let handle = tokio::spawn(async move {
        let mut result: Vec<u64> = Vec::new();
        // from Attackers
        while let Some(cmd) = rx1.recv().await {
            use Command::*;
            match cmd {
                Response { val } => {
                    result.push(val as u64);
                }
                _ => {} // We can add other commands like HTTP status code.
            }
        }
        println!("Collector task is done"); // at this point, all the spawn should be done.
        result
    });

    // 2. Attackers to send HTTP requests to a web server.
    tokio::spawn(async move {
        while let Some(_cmd) = rx2.recv().await { // receive from Cron Job
            for _ in 0..rate { // Spawn and run tasks in parallel.
                tokio::spawn(worker(url.clone(), tx1.clone())); // it's not closure.
            }
        }
        // Note that the workers are still running at this point.
        println!("Attacker tasks are done")
    });

    // My bad, this is not right place to measure the duration of test.
    let now = Instant::now();

    // 3. Cron job to invoke attackers every one second.
    tokio::spawn(async move {
        let mut count = 0;
        while count < duration { // fire every second during DURATION.
            tokio::time::sleep(Duration::from_secs(1)).await;
            print!(".");
            io::stdout().flush().unwrap();

            let cmd = Command::Get { key: "invoke".into() };
            tx2.send(cmd).await.unwrap();

            count += 1;
        }
        println!("\nCron Job task is done")
    });

    let result: Vec<u64> = handle.await?;  // Collector is done.
    let actual_duration = now.elapsed().as_secs() as f32;

    // Let's compute and report the stats such as mean, median and so on.
    let mut h = Histogram::new();
    for i in result.iter() {  // a list of response times, [498, 127, 283...]
        h.increment(*i).unwrap();
    }
    println!("Percentiles: p50:{}ms p90:{}ms p99:{}ms p999:{}ms",
             h.percentile(50.0).unwrap(),
             h.percentile(90.0).unwrap(),
             h.percentile(99.0).unwrap(),
             h.percentile(99.9).unwrap(),
    );
    println!("Latency (ms): Min:{} Avg:{} Max:{} StdDev:{}",
             h.minimum().unwrap(),
             h.mean().unwrap(),
             h.maximum().unwrap(),
             h.stddev().unwrap(),
    );
    let requests = result.len() as f32;
    println!("{} total requests", requests);
    println!("{} req/sec", requests / actual_duration);
    Ok(())
}

// Here, implement your logic to attack the server.
async fn worker(url: String, tx1: Sender<Command>) {
    let now = Instant::now();

    // You can add GET params to make it more realistic use case.
    // You might want to use rand::thread_rng() if needed.
    let _ = reqwest::get(url)
        .await.unwrap()
        .text() // just ignore the response.
        .await.unwrap();

    let cmd = Command::Response { val: now.elapsed().as_millis() };
    tx1.send(cmd).await.unwrap();
}