Fix how we handle control-c in stress

MystenLabs · Sep 19, 2022 · 87306c7 · 87306c7
1 parent c4083b9
commit 87306c7
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/sui-benchmark/Cargo.toml b/crates/sui-benchmark/Cargo.toml
@@ -35,6 +35,7 @@ duration-str = "0.4.0"
 hdrhistogram = "7.5.1"
 comfy-table = "6.1.0"
 bcs = "0.1.3"
+tokio-util = "0.7.4"
 sui-core = { path = "../sui-core" }
 sui-config = { path = "../sui-config" }
 sui-types = { path = "../sui-types" }

diff --git a/crates/sui-benchmark/src/drivers/bench_driver.rs b/crates/sui-benchmark/src/drivers/bench_driver.rs
@@ -18,6 +18,7 @@ use prometheus::IntCounterVec;
 use prometheus::Registry;
 use sui_core::authority_aggregator::AuthorityAggregator;
 use tokio::sync::OnceCell;
+use tokio_util::sync::CancellationToken;
 
 use crate::drivers::driver::Driver;
 use crate::drivers::HistogramWrapper;
@@ -127,15 +128,20 @@ pub struct BenchWorker {
 pub struct BenchDriver {
     pub stat_collection_interval: u64,
     pub start_time: Instant,
+    pub token: CancellationToken,
 }
 
 impl BenchDriver {
     pub fn new(stat_collection_interval: u64) -> BenchDriver {
         BenchDriver {
             stat_collection_interval,
             start_time: Instant::now(),
+            token: CancellationToken::new(),
         }
     }
+    pub fn terminate(&self) {
+        self.token.cancel()
+    }
     pub fn update_progress(
         start_time: Instant,
         interval: Interval,
@@ -224,6 +230,7 @@ impl Driver<BenchmarkStats> for BenchDriver {
                 ),
         });
         for (i, worker) in bench_workers.into_iter().enumerate() {
+            let cloned_token = self.token.clone();
             let request_delay_micros = 1_000_000 / worker.target_qps;
             let mut free_pool = worker.payload;
             let progress = progress.clone();
@@ -254,7 +261,7 @@ impl Driver<BenchmarkStats> for BenchDriver {
                 let mut stat_start_time: Instant = Instant::now();
                 loop {
                     tokio::select! {
-                            _ = tokio::signal::ctrl_c() => {
+                            _ = cloned_token.cancelled() => {
                                 break;
                             }
                             _ = stat_interval.tick() => {
@@ -491,7 +498,14 @@ impl Driver<BenchmarkStats> for BenchDriver {
             benchmark_stat
         });
         drop(tx);
-        let _res: Vec<_> = try_join_all(tasks).await.unwrap().into_iter().collect();
+        let all_tasks = try_join_all(tasks);
+        let _res = tokio::select! {
+            _ = tokio::signal::ctrl_c() => {
+                self.terminate();
+                vec![]
+            }
+            res = all_tasks => res.unwrap().into_iter().collect()
+        };
         let benchmark_stat = stat_task.await.unwrap();
         Ok(benchmark_stat)
     }