Add nix crate to improve signal handling in multi-thread application.

Author: unpluggedcoder

Cargo.lock

@@ -15,6 +15,7 @@ lazy_static = "1"
 memmap = "0.7"
 log = "0.4"
 glob = "0.2"
+nix = "0.17"
 jemallocator = { path = "../jemallocator", default-features = false }
 jemalloc-sys = { path = "../jemallocator/jemalloc-sys", default-features = false }
 

preload/Cargo.toml

@@ -1,3 +1,4 @@
+use nix::sys::signal;
 use std::env;
 
 use crate::global::on_exit;
@@ -8,80 +9,94 @@ use crate::utils::generate_filename;
 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
     };
 
     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!");
+
+    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");
 }
 
 pub fn startup() {
     initialize_logger();
-    info!( "Version: {}", env!( "CARGO_PKG_VERSION" ) );
+    info!("Version: {}", env!("CARGO_PKG_VERSION"));
 
     unsafe {
         opt::initialize();
@@ -94,6 +109,6 @@ pub fn startup() {
 
     initialize_signal_handlers();
 
-    env::remove_var( "LD_PRELOAD" );
-    info!( "Startup initialization finished" );
+    env::remove_var("LD_PRELOAD");
+    info!("Startup initialization finished");
 }