From cd6dacbb5c0acc8bdd3faa8f17417e355aa9b558 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Thu, 7 May 2026 22:08:41 +0700 Subject: [PATCH] feat(cli): handle interrupt as graceful shutdown --- Cargo.lock | 72 +++++++++++++++++++++++++ crates/irtt-cli/Cargo.toml | 1 + crates/irtt-cli/src/main.rs | 101 +++++++++++++++++++++++++++++++++--- 3 files changed, 166 insertions(+), 8 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e55760a..80dea9b 100644 --- a/Cargo.lock +++ b/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" diff --git a/crates/irtt-cli/Cargo.toml b/crates/irtt-cli/Cargo.toml index d4b9c8a..acaf64a 100644 --- a/crates/irtt-cli/Cargo.toml +++ b/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 } diff --git a/crates/irtt-cli/src/main.rs b/crates/irtt-cli/src/main.rs index b5f9b8a..d61e610 100644 --- a/crates/irtt-cli/src/main.rs +++ b/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> { +fn install_signal_handler(shutdown_requested: Arc) -> Result<(), ctrlc::Error> { + ctrlc::set_handler(move || { + shutdown_requested.store(true, Ordering::Relaxed); + }) +} + +fn run(args: CliArgs, shutdown_requested: &AtomicBool) -> Result<(), Box> { 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> { #[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> { 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> { Ok(()) } +fn is_shutdown_requested(shutdown_requested: &AtomicBool) -> bool { + shutdown_requested.load(Ordering::Relaxed) +} + +fn should_continue_run( + continuous: bool, + next_send_deadline: Option, + shutdown_requested: &AtomicBool, +) -> bool { + !is_shutdown_requested(shutdown_requested) && (continuous || next_send_deadline.is_some()) +} + +fn should_send_probe(next_send_deadline: Option, 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::*;