add time argument implement

Author: shinnku-nikaidou

src/args.rs

@@ -9,5 +9,5 @@ pub struct Args {
     pub url: String,
 
     #[arg(short = 't', long = "time", default_value_t = 60)]
-    pub time: u32,
+    pub time: u64,
 }

src/main.rs

@@ -1,27 +1,49 @@
 mod args;
+use std::time::Duration;
+
 use args::Args;
 use clap::Parser;
 use reqwest::Client;
+use tokio::sync::broadcast;
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let args = Args::parse();
     let mut handles = Vec::new();
     let client = Client::new();
+    let (shutdown_tx, _) = broadcast::channel(1);
 
     for _ in 0..args.concurrent_count {
         let url = args.url.clone();
         let client = client.clone();
+        let mut shutdown_rx = shutdown_tx.subscribe();
+
         let handle = tokio::spawn(async move {
-            if let Ok(resp) = client.get(&url).send().await {
-                let _ = resp.bytes().await;
+            loop {
+                tokio::select! {
+                    biased;
+                    _ = shutdown_rx.recv() => {
+                        break;
+                    }
+                    result = client.get(&url).send() => {
+                        if let Ok(resp) = result {
+                            let _ = resp.bytes().await;
+                        }
+                        tokio::task::yield_now().await;
+                    }
+                }
             }
         });
         handles.push(handle);
     }
 
+    tokio::time::sleep(Duration::from_secs(args.time)).await;
+    let _ = shutdown_tx.send(());
+
     for handle in handles {
-        handle.await.unwrap();
+        if let Err(e) = handle.await {
+            eprintln!("Task exited with error: {:?}", e);
+        }
     }
 
     Ok(())