From 4b910d266ecbc598745ec9a4df3d5c113f5b7928 Mon Sep 17 00:00:00 2001 From: Janakarajan Natarajan Date: Mon, 22 Jul 2024 22:04:35 +0000 Subject: [PATCH] Add support for catching signals This is used if the user ends an Aperf run using Ctrl+c or kill . We send the same signal to any child processes launched by aperf. --- Cargo.lock | 84 +++++++++++--------------------- Cargo.toml | 3 +- src/data.rs | 7 +++ src/data/java_profile.rs | 5 ++ src/data/perf_profile.rs | 6 +++ src/lib.rs | 102 ++++++++++++++++++++++++++++----------- 6 files changed, 122 insertions(+), 85 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c75ad8d0..a433c319 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -118,6 +118,7 @@ dependencies = [ "lazy_static", "libc", "log", + "nix", "perf-event2", "procfs", "rustix 0.38.28", @@ -549,6 +550,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chrono" version = "0.4.38" @@ -823,17 +830,6 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" -[[package]] -name = "errno" -version = "0.2.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f639046355ee4f37944e44f60642c6f3a7efa3cf6b78c78a0d989a8ce6c396a1" -dependencies = [ - "errno-dragonfly", - "libc", - "winapi", -] - [[package]] name = "errno" version = "0.3.8" @@ -844,16 +840,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "errno-dragonfly" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" -dependencies = [ - "cc", - "libc", -] - [[package]] name = "fastrand" version = "1.8.0" @@ -1210,12 +1196,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "io-lifetimes" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9448015e586b611e5d322f6703812bbca2f1e709d5773ecd38ddb4e3bb649504" - [[package]] name = "io-lifetimes" version = "1.0.10" @@ -1234,7 +1214,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adcf93614601c8129ddf72e2d5633df827ba6551541c6d8c59520a371475be1f" dependencies = [ "hermit-abi", - "io-lifetimes 1.0.10", + "io-lifetimes", "rustix 0.37.18", "windows-sys 0.48.0", ] @@ -1262,9 +1242,9 @@ checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" [[package]] name = "libc" -version = "0.2.150" +version = "0.2.155" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89d92a4743f9a61002fae18374ed11e7973f530cb3a3255fb354818118b2203c" +checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c" [[package]] name = "link-cplusplus" @@ -1275,12 +1255,6 @@ dependencies = [ "cc", ] -[[package]] -name = "linux-raw-sys" -version = "0.0.46" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d2456c373231a208ad294c33dc5bff30051eafd954cd4caae83a712b12854d" - [[package]] name = "linux-raw-sys" version = "0.3.6" @@ -1354,6 +1328,18 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags 2.4.0", + "cfg-if", + "cfg_aliases", + "libc", +] + [[package]] name = "ntapi" version = "0.3.7" @@ -1693,20 +1679,6 @@ dependencies = [ "semver", ] -[[package]] -name = "rustix" -version = "0.34.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2079c267b8394eb529872c3cf92e181c378b41fea36e68130357b52493701d2e" -dependencies = [ - "bitflags 1.3.2", - "errno 0.2.8", - "io-lifetimes 0.6.1", - "libc", - "linux-raw-sys 0.0.46", - "winapi", -] - [[package]] name = "rustix" version = "0.37.18" @@ -1714,8 +1686,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8bbfc1d1c7c40c01715f47d71444744a81669ca84e8b63e25a55e169b1f86433" dependencies = [ "bitflags 1.3.2", - "errno 0.3.8", - "io-lifetimes 1.0.10", + "errno", + "io-lifetimes", "libc", "linux-raw-sys 0.3.6", "windows-sys 0.48.0", @@ -1728,7 +1700,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72e572a5e8ca657d7366229cdde4bd14c4eb5499a9573d4d366fe1b599daa316" dependencies = [ "bitflags 2.4.0", - "errno 0.3.8", + "errno", "libc", "linux-raw-sys 0.4.12", "windows-sys 0.52.0", @@ -2177,11 +2149,11 @@ dependencies = [ [[package]] name = "timerfd" -version = "1.3.0" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29f85a7c965b8e7136952f59f2a359694c78f105b2d2ff99cf6c2c404bf7e33f" +checksum = "84e482e368cf7efa2c8b570f476e5b9fd9fd5e9b9219fc567832b05f13511091" dependencies = [ - "rustix 0.34.8", + "rustix 0.38.28", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 5276c5ff..23bd887a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ name = "aperf" path = "src/bin/aperf.rs" [dependencies] +nix = { version = "0.29.0", features = ["signal", "poll"] } clap = { version = "4.2.5", features = ["derive"] } serde = { version = "1.0", features = ["derive"] } chrono = { version = "0.4", features = ["serde"] } @@ -29,7 +30,7 @@ thiserror = "1.0" log = "0.4.21" env_logger = "0.10.0" lazy_static = "1.4.0" -timerfd = "1.3.0" +timerfd = "1.6.0" procfs = "0.12.0" ctor = "0.2.6" sysinfo = "0.26.2" diff --git a/src/data.rs b/src/data.rs index 940ddc8e..b14482f9 100644 --- a/src/data.rs +++ b/src/data.rs @@ -39,6 +39,7 @@ use kernel_config::KernelConfig; use log::trace; use meminfodata::{MeminfoData, MeminfoDataRaw}; use netstat::{Netstat, NetstatRaw}; +use nix::sys::{signal, signal::Signal}; use perf_profile::{PerfProfile, PerfProfileRaw}; use perf_stat::{PerfStat, PerfStatRaw}; use processes::{Processes, ProcessesRaw}; @@ -60,6 +61,7 @@ pub struct CollectorParams { pub run_name: String, pub profile: HashMap, pub tmp_dir: PathBuf, + pub signal: Signal, } impl CollectorParams { @@ -72,6 +74,7 @@ impl CollectorParams { run_name: String::new(), profile: HashMap::new(), tmp_dir: PathBuf::new(), + signal: signal::SIGTERM, } } } @@ -109,6 +112,10 @@ impl DataType { self.is_profile_option = true; } + pub fn set_signal(&mut self, signal: Signal) { + self.collector_params.signal = signal; + } + pub fn init_data_type(&mut self, param: InitParams) -> Result<()> { trace!("Initializing data type..."); let name = format!( diff --git a/src/data/java_profile.rs b/src/data/java_profile.rs index a5a207c5..46ec0d08 100644 --- a/src/data/java_profile.rs +++ b/src/data/java_profile.rs @@ -6,6 +6,7 @@ use crate::{PDError, PERFORMANCE_DATA, VISUALIZATION_DATA}; use anyhow::Result; use ctor::ctor; use log::{debug, error, trace}; +use nix::{sys::signal, unistd::Pid}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::io::Write; @@ -201,6 +202,10 @@ impl CollectData for JavaProfileRaw { } fn finish_data_collection(&mut self, params: CollectorParams) -> Result<()> { + for child in ASPROF_CHILDREN.lock().unwrap().iter() { + signal::kill(Pid::from_raw(child.id() as i32), params.signal)?; + } + trace!("Waiting for asprof profile collection to complete..."); while ASPROF_CHILDREN.lock().unwrap().len() > 0 { match ASPROF_CHILDREN.lock().unwrap().pop().unwrap().wait() { diff --git a/src/data/perf_profile.rs b/src/data/perf_profile.rs index 6a5780b5..45e4256b 100644 --- a/src/data/perf_profile.rs +++ b/src/data/perf_profile.rs @@ -6,6 +6,7 @@ use crate::{PDError, PERFORMANCE_DATA, VISUALIZATION_DATA}; use anyhow::Result; use ctor::ctor; use log::{error, trace}; +use nix::{sys::signal, unistd::Pid}; use serde::{Deserialize, Serialize}; use std::fs; use std::io::Write; @@ -79,6 +80,11 @@ impl CollectData for PerfProfileRaw { Some(_) => {} } + signal::kill( + Pid::from_raw(child.as_mut().unwrap().id() as i32), + params.signal, + )?; + trace!("Waiting for perf profile collection to complete..."); match child.as_mut().unwrap().wait() { Err(e) => { diff --git a/src/lib.rs b/src/lib.rs index a73bf95e..416efe20 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,9 +10,15 @@ use chrono::prelude::*; use data::TimeEnum; use flate2::{write::GzEncoder, Compression}; use log::{debug, error, info}; +use nix::poll::{poll, PollFd, PollFlags, PollTimeout}; +use nix::sys::{ + signal, + signalfd::{SfdFlags, SigSet, SignalFd}, +}; use serde::{Deserialize, Serialize}; use serde_json::{self}; use std::collections::HashMap; +use std::os::unix::io::AsFd; use std::path::{Path, PathBuf}; use std::sync::Mutex; use std::{fs, process, time}; @@ -243,6 +249,7 @@ impl PerformanceData { let mut current = time::Instant::now(); let end = current + time::Duration::from_secs(self.init_params.period); + // TimerFd let mut tfd = TimerFd::new()?; tfd.set_state( TimerState::Periodic { @@ -251,42 +258,81 @@ impl PerformanceData { }, SetTimeFlags::Default, ); + let timer_pollfd = PollFd::new(tfd.as_fd(), PollFlags::POLLIN); + + // SignalFd + let mut mask = SigSet::empty(); + mask.add(signal::SIGINT); + mask.add(signal::SIGTERM); + mask.thread_block()?; + let sfd = SignalFd::with_flags(&mask, SfdFlags::SFD_NONBLOCK)?; + let signal_pollfd = PollFd::new(sfd.as_fd(), PollFlags::POLLIN); + + let mut poll_fds = [timer_pollfd, signal_pollfd]; + let mut datatype_signal = signal::SIGTERM; + while current <= end { aperf_collect_data.time = TimeEnum::DateTime(Utc::now()); aperf_collect_data.data = HashMap::new(); - let ret = tfd.read(); - if ret > 1 { - error!("Missed {} interval(s)", ret - 1); + if poll(&mut poll_fds, PollTimeout::NONE)? <= 0 { + error!("Poll error."); + } + if let Some(ev) = poll_fds[0].revents() { + if ev.contains(PollFlags::POLLIN) { + let ret = tfd.read(); + if ret > 1 { + error!("Missed {} interval(s)", ret - 1); + } + debug!("Time elapsed: {:?}", start.elapsed()); + current += time::Duration::from_secs(ret * self.init_params.interval); + for (name, datatype) in self.collectors.iter_mut() { + if datatype.is_static { + continue; + } + + datatype.collector_params.elapsed_time = start.elapsed().as_secs(); + + aperf_collect_data.measure( + name.clone() + "-collect", + || -> Result<()> { + datatype.collect_data()?; + Ok(()) + }, + )?; + aperf_collect_data.measure(name.clone() + "-print", || -> Result<()> { + datatype.write_to_file()?; + Ok(()) + })?; + } + let data_collection_time = time::Instant::now() - current; + aperf_collect_data + .data + .insert("aperf".to_string(), data_collection_time.as_micros() as u64); + debug!("Collection time: {:?}", data_collection_time); + bincode::serialize_into( + self.aperf_stats_handle.as_ref().unwrap(), + &aperf_collect_data, + )?; + } } - debug!("Time elapsed: {:?}", start.elapsed()); - current += time::Duration::from_secs(ret * self.init_params.interval); - for (name, datatype) in self.collectors.iter_mut() { - if datatype.is_static { - continue; + if let Some(ev) = poll_fds[1].revents() { + if ev.contains(PollFlags::POLLIN) { + if let Ok(Some(siginfo)) = sfd.read_signal() { + if siginfo.ssi_signo == signal::SIGINT as u32 { + info!("Caught SIGINT. Exiting..."); + datatype_signal = signal::SIGINT; + } else if siginfo.ssi_signo == signal::SIGTERM as u32 { + info!("Caught SIGTERM. Exiting..."); + } else { + panic!("Caught an unknown signal: {}", siginfo.ssi_signo); + } + break; + } } - - datatype.collector_params.elapsed_time = start.elapsed().as_secs(); - - aperf_collect_data.measure(name.clone() + "-collect", || -> Result<()> { - datatype.collect_data()?; - Ok(()) - })?; - aperf_collect_data.measure(name.clone() + "-print", || -> Result<()> { - datatype.write_to_file()?; - Ok(()) - })?; } - let data_collection_time = time::Instant::now() - current; - aperf_collect_data - .data - .insert("aperf".to_string(), data_collection_time.as_micros() as u64); - debug!("Collection time: {:?}", data_collection_time); - bincode::serialize_into( - self.aperf_stats_handle.as_ref().unwrap(), - &aperf_collect_data, - )?; } for (_name, datatype) in self.collectors.iter_mut() { + datatype.set_signal(datatype_signal); datatype.finish_data_collection()?; } for (_name, datatype) in self.collectors.iter_mut() {