From 461a1385d9dc47a78cfaa729f4e265f2deb8c9e6 Mon Sep 17 00:00:00 2001
From: zeroXbrock <2791467+zeroXbrock@users.noreply.github.com>
Date: Wed, 11 Dec 2024 16:11:19 -0800
Subject: [PATCH] make CTRL-C handling extra-graceful (2-stage spam
 termination)

---
 crates/core/src/spammer/spammer_trait.rs | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/crates/core/src/spammer/spammer_trait.rs b/crates/core/src/spammer/spammer_trait.rs
index 69fec83..7aa6440 100644
--- a/crates/core/src/spammer/spammer_trait.rs
+++ b/crates/core/src/spammer/spammer_trait.rs
@@ -44,9 +44,11 @@ where
 
         let quit_clone = quit.clone();
         tokio::task::spawn(async move {
-            let _ = tokio::signal::ctrl_c().await;
-            let mut quit = quit_clone.lock().unwrap();
-            *quit = true;
+            loop {
+                let _ = tokio::signal::ctrl_c().await;
+                let mut quit = quit_clone.lock().unwrap();
+                *quit = true;
+            }
         });
 
         async move {
@@ -67,6 +69,13 @@ where
             let mut cursor = self.on_spam(scenario).await?.take(num_periods);
 
             while let Some(trigger) = cursor.next().await {
+                if *quit.lock().expect("lock failure") {
+                    println!("CTRL-C received, stopping spam and collecting results...");
+                    let mut quit = quit.lock().expect("lock failure");
+                    *quit = false;
+                    break;
+                }
+
                 let trigger = trigger.to_owned();
                 let payloads = scenario.prepare_spam(tx_req_chunks[tick]).await?;
                 let spam_tasks = scenario
@@ -93,12 +102,12 @@ where
                         break;
                     }
                     if *quit.lock().expect("lock failure") {
-                        println!("CTRL-C received, stopping spam...");
+                        println!("CTRL-C received, stopping result collection...");
                         break;
                     }
                     block_counter += 1;
                 }
-                println!("done spamming. run_id={}", run_id);
+                println!("done. run_id={}", run_id);
             }
 
             Ok(())