Files
blob-store-app/load_test/main.rs
2024-05-11 14:36:50 -04:00

80 lines
2.0 KiB
Rust

use std::sync::Arc;
use std::sync::Mutex;
use clap::Parser;
use kdam::tqdm;
use kdam::Bar;
use kdam::BarExt;
use rand::Rng;
use reqwest::StatusCode;
#[derive(Parser, Debug, Clone)]
#[command(version, about, long_about = None)]
struct Args {
#[arg(long)]
file_size: usize,
#[arg(long)]
num_threads: usize,
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let args = Args::parse();
let pb = Arc::new(Mutex::new(tqdm!()));
let mut handles = vec![];
for _ in 0..args.num_threads {
let pb = pb.clone();
let args = args.clone();
handles.push(std::thread::spawn(move || {
run_loop(pb, args).unwrap();
}));
}
for handle in handles {
handle.join().unwrap();
}
Ok(())
}
fn run_loop(pb: Arc<Mutex<Bar>>, args: Args) -> Result<(), Box<dyn std::error::Error>> {
let client = reqwest::blocking::Client::new();
let mut rng = rand::thread_rng();
let mut rand_data = vec![0u8; args.file_size];
rng.fill(&mut rand_data[..]);
let hints = vec![
"foo", "bar", "baz", "qux", "quux", "corge", "grault", "garply", "waldo", "fred", "plugh",
"xyzzy", "thud",
];
loop {
// tweak a byte in the data
let idx = rng.gen_range(0..rand_data.len());
rand_data[idx] = rng.gen();
let hint = hints[rng.gen_range(0..hints.len())];
let form = reqwest::blocking::multipart::Form::new()
.text("content_type", "text/plain")
.part(
"compression_hint",
reqwest::blocking::multipart::Part::text(hint),
)
.part(
"data",
reqwest::blocking::multipart::Part::bytes(rand_data.clone()),
);
let resp = client
.post("http://localhost:7692/store")
.multipart(form)
.send()?;
// update progress bar
let mut pb = pb.lock().unwrap();
if resp.status() == StatusCode::CREATED {
pb.update(1)?;
}
}
}