info

load_test/main.rs

@@ -1,20 +1,33 @@
 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![];
-    let num_shards = 8;
-    for _ in 0..num_shards {
+    for _ in 0..args.num_threads {
         let pb = pb.clone();
+        let args = args.clone();
         handles.push(std::thread::spawn(move || {
-            run_loop(pb).unwrap();
+            run_loop(pb, args).unwrap();
         }));
     }
 
@@ -25,10 +38,10 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
-fn run_loop(pb: Arc<Mutex<Bar>>) -> Result<(), Box<dyn std::error::Error>> {
+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; 1024 * 1024];
+    let mut rand_data = vec![0u8; args.file_size];
     rng.fill(&mut rand_data[..]);
 
     loop {
@@ -50,7 +63,8 @@ fn run_loop(pb: Arc<Mutex<Bar>>) -> Result<(), Box<dyn std::error::Error>> {
 
         // update progress bar
         let mut pb = pb.lock().unwrap();
-        pb.update(1)?;
-        if resp.status() != 200 && resp.status() != 201 {}
+        if resp.status() == StatusCode::CREATED {
+            pb.update(1)?;
+        }
     }
 }