spawn

Author: shikhar

src/bench.rs

@@ -20,6 +20,7 @@ use s2_sdk::{
     },
 };
 use tokio::time::Instant;
+use tokio::sync::mpsc;
 use xxhash_rust::xxh3::Xxh3Default;
 
 use crate::error::{CliError, OpKind};
@@ -493,8 +494,36 @@ pub async fn run(
         write_done_records.clone(),
         bench_start,
     );
-    let mut write_stream = std::pin::pin!(write_stream);
-    let mut read_stream = std::pin::pin!(read_stream);
+
+    enum BenchEvent {
+        Write(Result<BenchWriteSample, CliError>),
+        Read(Result<BenchReadSample, CliError>),
+        WriteDone,
+        ReadDone,
+    }
+
+    let (tx, mut rx) = mpsc::unbounded_channel();
+    let write_tx = tx.clone();
+    let write_handle = tokio::spawn(async move {
+        let mut write_stream = std::pin::pin!(write_stream);
+        while let Some(sample) = write_stream.next().await {
+            if write_tx.send(BenchEvent::Write(sample)).is_err() {
+                return;
+            }
+        }
+        let _ = write_tx.send(BenchEvent::WriteDone);
+    });
+    let read_tx = tx.clone();
+    let read_handle = tokio::spawn(async move {
+        let mut read_stream = std::pin::pin!(read_stream);
+        while let Some(sample) = read_stream.next().await {
+            if read_tx.send(BenchEvent::Read(sample)).is_err() {
+                return;
+            }
+        }
+        let _ = read_tx.send(BenchEvent::ReadDone);
+    });
+    drop(tx);
 
     let deadline = bench_start + duration;
     let mut write_done = false;
@@ -511,47 +540,56 @@ pub async fn run(
             _ = tokio::signal::ctrl_c(), if !stop.load(Ordering::Relaxed) => {
                 stop.store(true, Ordering::Relaxed);
             }
-            result = write_stream.next(), if !write_done => {
-                match result {
-                    Some(Ok(sample)) => {
+            event = rx.recv() => {
+                match event {
+                    Some(BenchEvent::Write(Ok(sample))) => {
                         update_bench_bar(&write_bar, &sample);
                         all_ack_latencies.extend(sample.ack_latencies.iter().copied());
                         if let Some(hash) = sample.run_hash {
                             write_run_hash = Some(hash);
                         }
                         write_sample = Some(sample);
                     }
-                    Some(Err(e)) => {
+                    Some(BenchEvent::Write(Err(e))) => {
                         write_bar.finish_and_clear();
                         read_bar.finish_and_clear();
+                        stop.store(true, Ordering::Relaxed);
+                        write_handle.abort();
+                        read_handle.abort();
                         return Err(e);
                     }
-                    None => {
+                    Some(BenchEvent::WriteDone) => {
                         write_done = true;
                     }
-                }
-            }
-            result = read_stream.next(), if !read_done => {
-                match result {
-                    Some(Ok(sample)) => {
+                    Some(BenchEvent::Read(Ok(sample))) => {
                         update_bench_bar(&read_bar, &sample);
                         all_e2e_latencies.extend(sample.e2e_latencies.iter().copied());
                         if let Some(hash) = sample.run_hash {
                             read_run_hash = Some(hash);
                         }
                         read_sample = Some(sample);
                     }
-                    Some(Err(e)) => {
+                    Some(BenchEvent::Read(Err(e))) => {
                         write_bar.finish_and_clear();
                         read_bar.finish_and_clear();
+                        stop.store(true, Ordering::Relaxed);
+                        write_handle.abort();
+                        read_handle.abort();
                         return Err(e);
                     }
-                    None => read_done = true,
+                    Some(BenchEvent::ReadDone) => read_done = true,
+                    None => {
+                        write_done = true;
+                        read_done = true;
+                    }
                 }
             }
         }
     }
 
+    let _ = write_handle.await;
+    let _ = read_handle.await;
+
     write_bar.finish_and_clear();
     read_bar.finish_and_clear();
 

src/cli.rs

@@ -513,7 +513,7 @@ pub struct BenchArgs {
     #[arg(
         short = 'b',
         long,
-        default_value_t = 16*1024,
+        default_value_t = 8*1024,
         value_parser = clap::value_parser!(u32).range(128..1024*1024),
     )]
     pub record_size: u32,

src/main.rs

@@ -532,7 +532,7 @@ async fn run() -> Result<(), CliError> {
                 .map_err(|e| CliError::op(OpKind::Bench, e))?;
 
             eprintln!(
-                "Running for {} targeting {} MiB/s with {}B records",
+                "Running for {} targeting {} MiB/s with {}b records",
                 args.duration, args.target_mibps, args.record_size,
             );