diff --git a/Cargo.lock b/Cargo.lock index 47c5e55e..27cf60ca 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1501,6 +1501,18 @@ dependencies = [ "memoffset", ] +[[package]] +name = "nix" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "bitflags 1.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "cc 1.0.45 (registry+https://github.com/rust-lang/crates.io-index)", + "cfg-if 0.1.10 (registry+https://github.com/rust-lang/crates.io-index)", + "libc 0.2.62 (registry+https://github.com/rust-lang/crates.io-index)", + "void 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "nodrop" version = "0.1.14" diff --git a/preload/Cargo.toml b/preload/Cargo.toml index 0e1317e8..acbd5658 100644 --- a/preload/Cargo.toml +++ b/preload/Cargo.toml @@ -21,6 +21,7 @@ tikv-jemallocator = { path = "../jemallocator", default-features = false } tikv-jemalloc-sys = { path = "../jemallocator/jemalloc-sys", default-features = false } goblin = "0.0.24" smallvec = { version = "1", features = ["union"] } +nix = "0.17" [dependencies.nwind] git = "https://github.com/koute/not-perf.git" diff --git a/preload/src/init.rs b/preload/src/init.rs index cc69477a..537a01c8 100644 --- a/preload/src/init.rs +++ b/preload/src/init.rs @@ -1,21 +1,28 @@ +use nix::sys::signal; use std::env; +use std::thread; +use std::sync::atomic::{AtomicBool, Ordering}; use crate::global::on_exit; use crate::logger; use crate::opt; +use crate::spin_lock::SpinLock; use crate::utils::generate_filename; +static SIGNALWAIT_THREAD_HANDLE: SpinLock>> = + SpinLock::new(None); + fn initialize_logger() { static mut SYSCALL_LOGGER: logger::SyscallLogger = logger::SyscallLogger::empty(); static mut FILE_LOGGER: logger::FileLogger = logger::FileLogger::empty(); - let log_level = if let Ok( value ) = env::var( "MEMORY_PROFILER_LOG" ) { + let log_level = if let Ok(value) = env::var("MEMORY_PROFILER_LOG") { match value.as_str() { "trace" => log::LevelFilter::Trace, "debug" => log::LevelFilter::Debug, "info" => log::LevelFilter::Info, "warn" => log::LevelFilter::Warn, "error" => log::LevelFilter::Error, - _ => log::LevelFilter::Off + _ => log::LevelFilter::Off, } } else { log::LevelFilter::Off @@ -23,77 +30,100 @@ fn initialize_logger() { let pid = unsafe { libc::getpid() }; - if let Ok( value ) = env::var( "MEMORY_PROFILER_LOGFILE" ) { - let path = generate_filename( &value, None ); - let rotate_at = env::var( "MEMORY_PROFILER_LOGFILE_ROTATE_WHEN_BIGGER_THAN" ).ok().and_then( |value| value.parse().ok() ); + if let Ok(value) = env::var("MEMORY_PROFILER_LOGFILE") { + let path = generate_filename(&value, None); + let rotate_at = env::var("MEMORY_PROFILER_LOGFILE_ROTATE_WHEN_BIGGER_THAN") + .ok() + .and_then(|value| value.parse().ok()); unsafe { - if let Ok(()) = FILE_LOGGER.initialize( path, rotate_at, log_level, pid ) { - log::set_logger( &FILE_LOGGER ).unwrap(); + if let Ok(()) = FILE_LOGGER.initialize(path, rotate_at, log_level, pid) { + log::set_logger(&FILE_LOGGER).unwrap(); } } } else { unsafe { - SYSCALL_LOGGER.initialize( log_level, pid ); - log::set_logger( &SYSCALL_LOGGER ).unwrap(); + SYSCALL_LOGGER.initialize(log_level, pid); + log::set_logger(&SYSCALL_LOGGER).unwrap(); } } - log::set_max_level( log_level ); + log::set_max_level(log_level); } fn initialize_atexit_hook() { - info!( "Setting atexit hook..." ); + info!("Setting atexit hook..."); unsafe { - let result = libc::atexit( on_exit ); + let result = libc::atexit(on_exit); if result != 0 { - error!( "Cannot set the at-exit hook" ); + error!("Cannot set the at-exit hook"); } } } fn initialize_signal_handlers() { - extern "C" fn sigusr_handler( signal: libc::c_int ) { - let signal_name = match signal { - libc::SIGUSR1 => "SIGUSR1", - libc::SIGUSR2 => "SIGUSR2", - _ => "???" - }; - - info!( "Signal handler triggered with signal: {} ({})", signal_name, signal ); - crate::global::toggle(); + if !opt::get().register_sigusr1 && !opt::get().register_sigusr2 { + info!("Skip signal register."); + return; } - + let mut sigset = signal::SigSet::empty(); if opt::get().register_sigusr1 { - info!( "Registering SIGUSR1 handler..." ); - unsafe { - libc::signal( libc::SIGUSR1, sigusr_handler as libc::sighandler_t ); - } + sigset.add(signal::Signal::SIGUSR1); + info!("Registering SIGUSR1 handler..."); } - if opt::get().register_sigusr2 { - info!( "Registering SIGUSR2 handler..." ); - unsafe { - libc::signal( libc::SIGUSR2, sigusr_handler as libc::sighandler_t ); - } + sigset.add(signal::Signal::SIGUSR2); + info!("Registering SIGUSR2 handler..."); } + sigset.thread_block().expect("Register signal failed!"); + + let mut thread_handle = SIGNALWAIT_THREAD_HANDLE.lock(); + static SIG_THREAD_RUNNING: AtomicBool = AtomicBool::new(false); + let new_handle = thread::Builder::new() + .name("Sigwait".into()) + .spawn(move || loop { + match sigset.wait() { + Ok(sig) => { + let signal_name = match sig { + signal::Signal::SIGUSR1 => "SIGUSR1", + signal::Signal::SIGUSR2 => "SIGUSR2", + _ => "???", + }; + + info!( + "Signal handler triggered with signal: {} ({})", + signal_name, sig + ); + crate::global::toggle(); + } + Err(e) => error!("Signal wait error: {}", e), + } + }) + .expect("Failed to start Sigwait thread"); + + while SIG_THREAD_RUNNING.load(Ordering::SeqCst) == false { + thread::yield_now(); + } + + *thread_handle = Some(new_handle); } pub fn startup() { initialize_logger(); - info!( "Version: {}", env!( "CARGO_PKG_VERSION" ) ); + info!("Version: {}", env!("CARGO_PKG_VERSION")); unsafe { opt::initialize(); } initialize_atexit_hook(); + + initialize_signal_handlers(); + if !opt::get().disabled_by_default { crate::global::toggle(); } - initialize_signal_handlers(); - - env::remove_var( "LD_PRELOAD" ); - info!( "Startup initialization finished" ); + env::remove_var("LD_PRELOAD"); + info!("Startup initialization finished"); }