feat(cli): handle interrupt as graceful shutdown

Cargo.lock

@@ -52,6 +52,12 @@ dependencies = [
  "windows-sys 0.61.2",
 ]
 
+[[package]]
+name = "bitflags"
+version = "2.11.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3"
+
 [[package]]
 name = "block-buffer"
 version = "0.10.4"
@@ -61,12 +67,27 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "block2"
+version = "0.6.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5"
+dependencies = [
+ "objc2",
+]
+
 [[package]]
 name = "cfg-if"
 version = "1.0.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
 
+[[package]]
+name = "cfg_aliases"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
+
 [[package]]
 name = "clap"
 version = "4.6.1"
@@ -123,6 +144,17 @@ dependencies = [
  "typenum",
 ]
 
+[[package]]
+name = "ctrlc"
+version = "3.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e0b1fab2ae45819af2d0731d60f2afe17227ebb1a1538a236da84c93e9a60162"
+dependencies = [
+ "dispatch2",
+ "nix",
+ "windows-sys 0.61.2",
+]
+
 [[package]]
 name = "digest"
 version = "0.10.7"
@@ -134,6 +166,18 @@ dependencies = [
  "subtle",
 ]
 
+[[package]]
+name = "dispatch2"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38"
+dependencies = [
+ "bitflags",
+ "block2",
+ "libc",
+ "objc2",
+]
+
 [[package]]
 name = "generic-array"
 version = "0.14.7"
@@ -164,6 +208,7 @@ name = "irtt-cli"
 version = "0.1.0"
 dependencies = [
  "clap",
+ "ctrlc",
  "irtt-client",
  "irtt-proto",
  "irtt-stats",
@@ -218,6 +263,33 @@ dependencies = [
  "digest",
 ]
 
+[[package]]
+name = "nix"
+version = "0.31.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3"
+dependencies = [
+ "bitflags",
+ "cfg-if",
+ "cfg_aliases",
+ "libc",
+]
+
+[[package]]
+name = "objc2"
+version = "0.6.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f"
+dependencies = [
+ "objc2-encode",
+]
+
+[[package]]
+name = "objc2-encode"
+version = "4.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33"
+
 [[package]]
 name = "once_cell_polyfill"
 version = "1.70.2"

crates/irtt-cli/Cargo.toml

@@ -11,6 +11,7 @@ stats = ["dep:irtt-stats"]
 
 [dependencies]
 clap = { version = "4.5", features = ["derive"] }
+ctrlc = "3.4"
 irtt-client = { path = "../irtt-client" }
 irtt-proto = { path = "../irtt-proto" }
 irtt-stats = { path = "../irtt-stats", optional = true }

crates/irtt-cli/src/main.rs

@@ -1,6 +1,10 @@
 use std::{
     io::{self, Write},
     process::ExitCode,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
     thread,
     time::{Duration, Instant},
 };
@@ -18,7 +22,13 @@ const MAX_SLEEP: Duration = Duration::from_millis(20);
 
 fn main() -> ExitCode {
     let args = CliArgs::parse();
-    match run(args) {
+    let shutdown_requested = Arc::new(AtomicBool::new(false));
+    if let Err(err) = install_signal_handler(Arc::clone(&shutdown_requested)) {
+        eprintln!("irtt-rs: failed to install signal handler: {err}");
+        return ExitCode::FAILURE;
+    }
+
+    match run(args, shutdown_requested.as_ref()) {
         Ok(()) => ExitCode::SUCCESS,
         Err(err) => {
             eprintln!("irtt-rs: {err}");
@@ -27,7 +37,13 @@ fn main() -> ExitCode {
     }
 }
 
-fn run(args: CliArgs) -> Result<(), Box<dyn std::error::Error>> {
+fn install_signal_handler(shutdown_requested: Arc<AtomicBool>) -> Result<(), ctrlc::Error> {
+    ctrlc::set_handler(move || {
+        shutdown_requested.store(true, Ordering::Relaxed);
+    })
+}
+
+fn run(args: CliArgs, shutdown_requested: &AtomicBool) -> Result<(), Box<dyn std::error::Error>> {
     let mode = args.output;
     let continuous = args.is_continuous();
     let mut stdout = io::LineWriter::new(io::stdout().lock());
@@ -40,16 +56,23 @@ fn run(args: CliArgs) -> Result<(), Box<dyn std::error::Error>> {
         #[cfg(feature = "stats")]
         stats: &mut stats,
     };
+
+    if is_shutdown_requested(shutdown_requested) {
+        return Ok(());
+    }
+
     let mut client = Client::connect(args.to_client_config())?;
 
+    if is_shutdown_requested(shutdown_requested) {
+        return Ok(());
+    }
+
     let open = client.open(ClientTimestamp::now())?;
     output.print_event(open_event(&open))?;
 
-    while continuous || client.next_send_deadline().is_some() {
-        if client
-            .next_send_deadline()
-            .is_some_and(|deadline| deadline <= Instant::now())
-        {
+    let mut interrupted = false;
+    while should_continue_run(continuous, client.next_send_deadline(), shutdown_requested) {
+        if should_send_probe(client.next_send_deadline(), shutdown_requested) {
             let events = client.send_probe()?;
             output.print_events(&events)?;
         }
@@ -60,10 +83,20 @@ fn run(args: CliArgs) -> Result<(), Box<dyn std::error::Error>> {
         let events = client.poll_timeouts(ClientTimestamp::now())?;
         output.print_events(&events)?;
 
+        if is_shutdown_requested(shutdown_requested) {
+            interrupted = true;
+            break;
+        }
+
         sleep_until_next_send(client.next_send_deadline());
     }
+    interrupted |= is_shutdown_requested(shutdown_requested);
 
-    if !continuous {
+    if interrupted {
+        eprintln!("interrupted, closing session...");
+    }
+
+    if should_drain_final(continuous, interrupted) {
         drain_final_replies(&mut client, &mut output)?;
     }
 
@@ -77,6 +110,27 @@ fn run(args: CliArgs) -> Result<(), Box<dyn std::error::Error>> {
     Ok(())
 }
 
+fn is_shutdown_requested(shutdown_requested: &AtomicBool) -> bool {
+    shutdown_requested.load(Ordering::Relaxed)
+}
+
+fn should_continue_run(
+    continuous: bool,
+    next_send_deadline: Option<Instant>,
+    shutdown_requested: &AtomicBool,
+) -> bool {
+    !is_shutdown_requested(shutdown_requested) && (continuous || next_send_deadline.is_some())
+}
+
+fn should_send_probe(next_send_deadline: Option<Instant>, shutdown_requested: &AtomicBool) -> bool {
+    !is_shutdown_requested(shutdown_requested)
+        && next_send_deadline.is_some_and(|deadline| deadline <= Instant::now())
+}
+
+fn should_drain_final(continuous: bool, interrupted: bool) -> bool {
+    !continuous || interrupted
+}
+
 fn open_event(outcome: &OpenOutcome) -> &ClientEvent {
     match outcome {
         OpenOutcome::Started { event, .. } | OpenOutcome::NoTestCompleted { event, .. } => event,
@@ -179,6 +233,37 @@ fn final_drain_duration(probe_timeout: Duration) -> Duration {
     probe_timeout.min(MAX_FINAL_DRAIN)
 }
 
+#[cfg(test)]
+mod shutdown_tests {
+    use super::*;
+
+    #[test]
+    fn shutdown_flag_stops_run_loop() {
+        let shutdown = AtomicBool::new(false);
+        assert!(should_continue_run(true, Some(Instant::now()), &shutdown));
+
+        shutdown.store(true, Ordering::Relaxed);
+        assert!(!should_continue_run(true, Some(Instant::now()), &shutdown));
+    }
+
+    #[test]
+    fn shutdown_flag_suppresses_due_probe_send() {
+        let shutdown = AtomicBool::new(false);
+        let due = Instant::now() - Duration::from_millis(1);
+        assert!(should_send_probe(Some(due), &shutdown));
+
+        shutdown.store(true, Ordering::Relaxed);
+        assert!(!should_send_probe(Some(due), &shutdown));
+    }
+
+    #[test]
+    fn interrupted_continuous_run_uses_final_drain_before_close() {
+        assert!(should_drain_final(true, true));
+        assert!(!should_drain_final(true, false));
+        assert!(should_drain_final(false, false));
+    }
+}
+
 #[cfg(all(test, feature = "stats"))]
 mod tests {
     use super::*;