From fb23a01c2c21cae0eebea2c71fb8442ed09ad691 Mon Sep 17 00:00:00 2001 From: Markus Stange Date: Thu, 2 May 2024 18:19:05 -0400 Subject: [PATCH] Gracefully handle Ctrl+C when the server is running. --- Cargo.lock | 31 +++---- samply/Cargo.toml | 2 +- samply/src/linux/profiler.rs | 52 ++++------- samply/src/mac/profiler.rs | 15 +--- samply/src/server.rs | 13 ++- samply/src/shared/ctrl_c.rs | 164 +++++++++++++++++++++++++++++++++++ samply/src/shared/mod.rs | 1 + 7 files changed, 210 insertions(+), 68 deletions(-) create mode 100644 samply/src/shared/ctrl_c.rs diff --git a/Cargo.lock b/Cargo.lock index 30a20474..0e6f91e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -441,6 +441,16 @@ version = "0.8.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "248e3bacc7dc6baa3b21e405ee045c3047101a49145e7e9eca583ab4c2ca5345" +[[package]] +name = "ctrlc" +version = "3.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "672465ae37dc1bc6380a6547a8883d5dd397b0f1faaad4f265726cc7042a5345" +dependencies = [ + "nix", + "windows-sys 0.52.0", +] + [[package]] name = "debugid" version = "0.8.0" @@ -1750,6 +1760,7 @@ dependencies = [ "byteorder", "clap", "crossbeam-channel", + "ctrlc", "debugid", "dirs", "flate2", @@ -1784,7 +1795,6 @@ dependencies = [ "serde", "serde_derive", "serde_json", - "signal-hook", "sysctl", "tempfile", "thiserror", @@ -1921,25 +1931,6 @@ dependencies = [ "serde", ] -[[package]] -name = "signal-hook" -version = "0.3.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8621587d4798caf8eb44879d42e56b9a93ea5dcd315a6487c357130095b62801" -dependencies = [ - "libc", - "signal-hook-registry", -] - -[[package]] -name = "signal-hook-registry" -version = "1.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" -dependencies = [ - "libc", -] - [[package]] name = "slab" version = "0.4.9" diff --git a/samply/Cargo.toml b/samply/Cargo.toml index 538d33ba..564db680 100644 --- a/samply/Cargo.toml +++ b/samply/Cargo.toml @@ -45,12 +45,12 @@ dirs = "5.0.0" once_cell = "1.17" fxhash = "0.2.1" mio = { version = "0.8.11", features = ["os-ext", "os-poll"] } +ctrlc = "3.4.4" [target.'cfg(any(target_os = "android", target_os = "macos", target_os = "linux"))'.dependencies] libc = "0.2.71" crossbeam-channel = "0.5.12" -signal-hook = "0.3.9" [target.'cfg(target_os = "macos")'.dependencies] diff --git a/samply/src/linux/profiler.rs b/samply/src/linux/profiler.rs index e66138c3..978766fd 100644 --- a/samply/src/linux/profiler.rs +++ b/samply/src/linux/profiler.rs @@ -5,6 +5,7 @@ use linux_perf_data::linux_perf_event_reader::{ CpuMode, Endianness, Mmap2FileId, Mmap2InodeAndVersion, Mmap2Record, RawData, }; use nix::sys::wait::WaitStatus; +use tokio::sync::oneshot; use std::collections::HashMap; use std::fs::File; @@ -13,8 +14,6 @@ use std::ops::Deref; use std::os::unix::process::ExitStatusExt; use std::path::Path; use std::process::ExitStatus; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; use std::thread; use std::time::{Duration, SystemTime}; @@ -27,6 +26,7 @@ use crate::linux_shared::{ ConvertRegs, Converter, EventInterpretation, MmapRangeOrVec, OffCpuIndicator, }; use crate::server::{start_server_main, ServerProps}; +use crate::shared::ctrl_c::CtrlC; use crate::shared::recording_props::{ProcessLaunchProps, ProfileCreationProps, RecordingProps}; #[cfg(target_arch = "x86_64")] @@ -51,16 +51,10 @@ pub fn start_recording( iteration_count, } = process_launch_props; - // Ignore SIGINT in our process while the child process is running. The - // signal will still reach the child process, because Ctrl+C sends the - // SIGINT signal to all processes in the foreground process group. - let should_terminate_on_ctrl_c = Arc::new(AtomicBool::new(false)); - #[cfg(unix)] - signal_hook::flag::register_conditional_default( - signal_hook::consts::SIGINT, - should_terminate_on_ctrl_c.clone(), - ) - .expect("cannot register signal handler"); + // Ignore Ctrl+C while the subcommand is running. The signal still reaches the process + // under observation while we continue to record it. (ctrl+c will send the SIGINT signal + // to all processes in the foreground process group). + let mut ctrl_c_receiver = CtrlC::observe_oneshot(); // Start a new process for the launched command and get its pid. // The command will not start running until we tell it to. @@ -96,12 +90,12 @@ pub fn start_recording( // Tell the main thread to tell the child process to begin executing. profile_another_pid_reply_sender.send(true).unwrap(); - // Create a stop flag which always stays false. We won't stop profiling until the + // Create a stop receiver which is never notified. We won't stop profiling until the // child process is done. // If Ctrl+C is pressed, it will reach the child process, and the child process // will act on it and maybe terminate. If it does, profiling stops too because // the main thread's wait() call below will exit. - let stop_flag = Arc::new(AtomicBool::new(false)); + let (_stop_sender, stop_receiver) = oneshot::channel(); // Start profiling the process. run_profiler( @@ -111,7 +105,7 @@ pub fn start_recording( time_limit, profile_another_pid_request_receiver, profile_another_pid_reply_sender, - stop_flag, + stop_receiver, ); }); @@ -200,9 +194,8 @@ pub fn start_recording( .send(SamplerRequest::StopProfilingOncePerfEventsExhausted) .unwrap(); - // The child has quit. - // From now on, we want to terminate if the user presses Ctrl+C. - should_terminate_on_ctrl_c.store(true, std::sync::atomic::Ordering::SeqCst); + // The launched subprocess is done. From now on, we want to terminate if the user presses Ctrl+C. + ctrl_c_receiver.close(); // Now wait for the observer thread to quit. It will keep running until all // perf events are closed, which happens if all processes which the events @@ -236,14 +229,7 @@ pub fn start_profiling_pid( server_props: Option, ) { // When the first Ctrl+C is received, stop recording. - // The server launches after the recording finishes. On the second Ctrl+C, terminate the server. - let stop = Arc::new(AtomicBool::new(false)); - #[cfg(unix)] - signal_hook::flag::register_conditional_default(signal_hook::consts::SIGINT, stop.clone()) - .expect("cannot register signal handler"); - #[cfg(unix)] - signal_hook::flag::register(signal_hook::consts::SIGINT, stop.clone()) - .expect("cannot register signal handler"); + let ctrl_c_receiver = CtrlC::observe_oneshot(); // Create a channel for the observer thread to notify the main thread once // profiling has been initialized. @@ -254,7 +240,6 @@ pub fn start_profiling_pid( let output_file = recording_props.output_file.clone(); let observer_thread = thread::spawn({ - let stop = stop.clone(); move || { let interval = recording_props.interval; let time_limit = recording_props.time_limit; @@ -277,7 +262,7 @@ pub fn start_profiling_pid( time_limit, profile_another_pid_request_receiver, profile_another_pid_reply_sender, - stop, + ctrl_c_receiver, ) } }); @@ -301,15 +286,14 @@ pub fn start_profiling_pid( .unwrap(); // Now wait for the observer thread to quit. It will keep running until the - // stop flag has been set to true by Ctrl+C, or until all perf events are closed, + // CtrlC receiver has been notified, or until all perf events are closed, // which happens if all processes which the events are attached to have quit. observer_thread .join() .expect("couldn't join observer thread"); - // From now on we want Ctrl+C to always quit our process. The stop flag might still be - // false if the observer thread finished because the observed processes terminated. - stop.store(true, Ordering::SeqCst); + // From now on, pressing Ctrl+C will kill our process, because the observer will have + // dropped its CtrlC receiver by now. if let Some(server_props) = server_props { let libinfo_map = crate::profile_json_preparse::parse_libinfo_map_from_profile_file( @@ -535,7 +519,7 @@ fn run_profiler( _time_limit: Option, more_processes_request_receiver: Receiver, more_processes_reply_sender: Sender, - stop: Arc, + mut stop_receiver: oneshot::Receiver<()>, ) { // eprintln!("Running..."); @@ -544,7 +528,7 @@ fn run_profiler( let mut total_lost_events = 0; let mut last_timestamp = 0; loop { - if stop.load(Ordering::SeqCst) { + if stop_receiver.try_recv().is_ok() { break; } diff --git a/samply/src/mac/profiler.rs b/samply/src/mac/profiler.rs index bfc2913d..1adbeafe 100644 --- a/samply/src/mac/profiler.rs +++ b/samply/src/mac/profiler.rs @@ -6,8 +6,6 @@ use std::collections::HashMap; use std::fs::File; use std::io::BufWriter; use std::process::ExitStatus; -use std::sync::atomic::AtomicBool; -use std::sync::Arc; use std::thread; use std::time::Duration; @@ -16,6 +14,7 @@ use super::process_launcher::{MachError, ReceivedStuff, TaskAccepter}; use super::sampler::{JitdumpOrMarkerPath, Sampler, TaskInit}; use super::time::get_monotonic_timestamp; use crate::server::{start_server_main, ServerProps}; +use crate::shared::ctrl_c::CtrlC; use crate::shared::recording_props::{ProcessLaunchProps, ProfileCreationProps, RecordingProps}; pub fn start_profiling_pid( @@ -56,16 +55,10 @@ pub fn start_recording( sampler.run() }); - // Ignore SIGINT while the subcommand is running. The signal still reaches the process + // Ignore Ctrl+C while the subcommand is running. The signal still reaches the process // under observation while we continue to record it. (ctrl+c will send the SIGINT signal // to all processes in the foreground process group). - let should_terminate_on_ctrl_c = Arc::new(AtomicBool::new(false)); - #[cfg(unix)] - signal_hook::flag::register_conditional_default( - signal_hook::consts::SIGINT, - should_terminate_on_ctrl_c.clone(), - ) - .expect("cannot register signal handler"); + let mut ctrl_c_receiver = CtrlC::observe_oneshot(); let (mut task_accepter, task_launcher) = TaskAccepter::new(&command_name, &args, &env_vars)?; @@ -161,7 +154,7 @@ pub fn start_recording( } // The launched subprocess is done. From now on, we want to terminate if the user presses Ctrl+C. - should_terminate_on_ctrl_c.store(true, std::sync::atomic::Ordering::SeqCst); + ctrl_c_receiver.close(); accepter_sender .send(()) diff --git a/samply/src/server.rs b/samply/src/server.rs index e7bfaffd..9c0e401b 100644 --- a/samply/src/server.rs +++ b/samply/src/server.rs @@ -21,6 +21,8 @@ use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::Arc; +use crate::shared::ctrl_c::CtrlC; + #[derive(Clone, Debug)] pub struct ServerProps { pub port_selection: PortSelection, @@ -146,7 +148,7 @@ async fn start_server( } } - // Run this server for... forever! + // Run this server until it stops. if let Err(e) = server.await { eprintln!("server error: {e}"); } @@ -232,9 +234,16 @@ async fn run_server( template_values: Arc>, path_prefix: String, ) -> Result<(), Box> { + let mut ctrl_c_receiver = CtrlC::observe_oneshot(); + // We start a loop to continuously accept incoming connections loop { - let (stream, _) = listener.accept().await?; + let (stream, _) = tokio::select! { + stream_and_addr_res = listener.accept() => stream_and_addr_res?, + ctrl_c_result = &mut ctrl_c_receiver => { + return Ok(ctrl_c_result?); + } + }; // Use an adapter to access something implementing `tokio::io` traits as if they implement // `hyper::rt` IO traits. diff --git a/samply/src/shared/ctrl_c.rs b/samply/src/shared/ctrl_c.rs new file mode 100644 index 00000000..a94211b2 --- /dev/null +++ b/samply/src/shared/ctrl_c.rs @@ -0,0 +1,164 @@ +use std::sync::{Arc, Mutex, OnceLock}; + +use tokio::sync::oneshot; + +static INSTANCE: OnceLock>> = OnceLock::new(); + +/// Provides Ctrl+C notifications, and allows suppressing the automatic termination +/// of the process. +/// +/// Ctrl+C can only be suppressed "once at a time". Imagine the following scenario: +/// +/// - You call `CtrlC::observe_oneshot()` and store the receiver in a variable +/// that remains alive for the upcoming long-running activity. +/// - The long-running is running. +/// - The user presses Ctrl+C. This sends a message to the receiver. +/// - You take your time reacting to the message, or you don't check it often enough. +/// - In the meantime, the user presses Ctrl+C again. +/// +/// This second press terminates the process, which is what the user expects. +pub struct CtrlC; + +impl CtrlC { + /// Returns a new [`oneshot::Receiver`] which will receive a message once + /// Ctrl+C is pressed. + /// + /// Suspends the automatic termination for *one* Ctrl+C. + /// + /// But the Ctrl+C after that will terminate again - unless `observe_oneshot` + /// has been called again since then. + /// + /// Furthermore, once the receiver is dropped, Ctrl+C will also terminate the process. + /// + /// ## Usage + /// + /// ### Example 1: Suspend automatic termination for a given scope + /// + /// ``` + /// let mut ctrl_c_receiver = CtrlC::observe_oneshot(); + /// + /// // do something + /// // [...] + /// + /// ctrl_c_receiver.close(); // Restores automatic termination behavior + /// ``` + /// + /// ### Example 2: Suspend automatic termination and check if Ctrl+C was pressed + /// + /// ``` + /// let mut ctrl_c_receiver = CtrlC::observe_oneshot(); + /// + /// // do something + /// // [...] + /// + /// match ctrl_c_receiver.try_recv() { + /// Ok(()) => { + /// // Ctrl+C was pressed once. If it had been pressed another time then we wouldn't + /// // be here because the process would already have terminated. + /// } + /// Err(TryRecvError::Empty) => { + /// // Ctrl+C was not pressed. + /// } + /// Err(TryRecvError::Closed) => { + /// // Someone else has called `CtrlC::observe_oneshot()` in the meantime and swapped + /// // out our handler. + /// // When our handler was active, Ctrl+C was not pressed. + /// } + /// } + /// ctrl_c_receiver.close(); // Restores automatic termination behavior + /// // Alternatively, just drop ctrl_c_receiver, or let it go out of scope. + /// ``` + /// + /// ### Example 3: Keep checking for Ctrl+C in a loop + /// + /// ``` + /// let mut ctrl_c_receiver = CtrlC::observe_oneshot(); + /// + /// loop { + /// if ctrl_c_receiver.try_recv().is_ok() { + /// // Ctrl+C was pressed once. Exit the loop. + /// // (If Ctrl+C had been pressed more than once then we wouldn't + /// // be here because the process would already have terminated.) + /// break; + /// } + /// + /// // do something + /// // [...] + /// } + /// + /// ctrl_c_receiver.close(); // Restores automatic termination behavior + /// // Alternatively, just drop ctrl_c_receiver, or let it go out of scope. + /// ``` + /// + /// ### Example 4: Loop on a future and stop early if Ctrl+C is pressed + /// + /// ``` + /// let mut ctrl_c_receiver = CtrlC::observe_oneshot(); + /// + /// loop { + /// tokio::select! { + /// ctrl_c_result = &mut ctrl_c_receiver => { + /// match ctrl_c_result { + /// Ok(()) => { + /// // Ctrl+C was pressed once. If it had been pressed another time then we wouldn't + /// // be here because the process would already have terminated. + /// } + /// Err(e) => { + /// // Someone else has called `CtrlC::observe_oneshot()` in the meantime and swapped + /// // out our handler. + /// // When our handler was active, Ctrl+C was not pressed. + /// } + /// } + /// } + /// something_else = some_other_future => { + /// // [...] + /// } + /// } + /// + /// // do something + /// // [...] + /// } + /// + /// ctrl_c_receiver.close(); // Restores automatic termination behavior + /// // Alternatively, just drop ctrl_c_receiver, or let it go out of scope. + /// ``` + pub fn observe_oneshot() -> oneshot::Receiver<()> { + let (tx, rx) = oneshot::channel(); + CtrlCState::get().lock().unwrap().current_sender = Some(tx); + rx + } +} + +struct CtrlCState { + current_sender: Option>, +} + +impl CtrlCState { + pub fn get() -> &'static Arc> { + INSTANCE.get_or_init(|| { + ctrlc::set_handler(|| { + let sender = CtrlCState::get().lock().unwrap().current_sender.take(); + if let Some(sender) = sender { + if let Ok(()) = sender.send(()) { + // The receiver still existed. Trust that it will handle this Ctrl+C. + // Do not terminate this process. + return; + } + } + + // We get here if there is no current handler installed, or if the + // receiver has been destroyed. + // Terminate the process. + terminate_for_ctrl_c(); + }) + .expect("Couldn't install Ctrl+C handler"); + Arc::new(Mutex::new(CtrlCState { + current_sender: None, + })) + }) + } +} + +fn terminate_for_ctrl_c() -> ! { + std::process::exit(1) +} diff --git a/samply/src/shared/mod.rs b/samply/src/shared/mod.rs index 20175182..83d5984a 100644 --- a/samply/src/shared/mod.rs +++ b/samply/src/shared/mod.rs @@ -1,4 +1,5 @@ pub mod context_switch; +pub mod ctrl_c; pub mod jit_category_manager; pub mod jit_function_add_marker; pub mod jit_function_recycler;