Merge branch 'handlers-mt'

Author: Byron

Diff for: Cargo.lock

@@ -191,6 +191,9 @@ regex = { version = "1.6.0", optional = true, default-features = false, features
 # For internal use to allow pure-Rust builds without openssl.
 reqwest-for-configuration-only = { package = "reqwest", version = "0.11.13", default-features = false, optional = true }
 
+# for `interrupt` module
+parking_lot = "0.12.1"
+
 document-features = { version = "0.2.0", optional = true }
 
 [target.'cfg(target_vendor = "apple")'.dependencies]

Diff for: gix/Cargo.toml

@@ -11,7 +11,7 @@ fn main() -> anyhow::Result<()> {
         .nth(2)
         .context("The second argument is the directory to clone the repository into")?;
 
-    gix::interrupt::init_handler(|| {})?;
+    gix::interrupt::init_handler(1, || {})?;
     std::fs::create_dir_all(&dst)?;
     let url = gix::url::parse(repo_url.to_str().unwrap().into())?;
 

Diff for: gix/examples/clone.rs

@@ -3,7 +3,7 @@ use std::path::Path;
 use gix_tempfile::{AutoRemove, ContainingDirectory};
 
 fn main() -> anyhow::Result<()> {
-    gix::interrupt::init_handler(|| {})?;
+    gix::interrupt::init_handler(1, || {})?;
     eprintln!("About to emit the first term signal");
     let tempfile_path = Path::new("example-file.tmp");
     let _keep_tempfile = gix_tempfile::mark_at(tempfile_path, ContainingDirectory::Exists, AutoRemove::Tempfile)?;

Diff for: gix/examples/interrupt-handler-allows-graceful-shutdown.rs

@@ -1,6 +1,6 @@
 fn main() -> anyhow::Result<()> {
     {
-        let _deregister_on_drop = gix::interrupt::init_handler(|| {})?.auto_deregister();
+        let _deregister_on_drop = gix::interrupt::init_handler(1, || {})?.auto_deregister();
     }
     eprintln!("About to emit the first term signal, which acts just like a normal one");
     signal_hook::low_level::raise(signal_hook::consts::SIGTERM)?;

Diff for: gix/examples/reversible-interrupt-handlers.rs

@@ -8,46 +8,62 @@
 mod init {
     use std::{
         io,
-        sync::atomic::{AtomicBool, AtomicUsize, Ordering},
+        sync::atomic::{AtomicUsize, Ordering},
     };
 
-    static IS_INITIALIZED: AtomicBool = AtomicBool::new(false);
+    static DEREGISTER_COUNT: AtomicUsize = AtomicUsize::new(0);
+    static REGISTERED_HOOKS: once_cell::sync::Lazy<parking_lot::Mutex<Vec<(i32, signal_hook::SigId)>>> =
+        once_cell::sync::Lazy::new(Default::default);
+    static DEFAULT_BEHAVIOUR_HOOKS: once_cell::sync::Lazy<parking_lot::Mutex<Vec<signal_hook::SigId>>> =
+        once_cell::sync::Lazy::new(Default::default);
 
+    /// A type to help deregistering hooks registered with [`init_handler`](super::init_handler());
     #[derive(Default)]
-    pub struct Deregister(Vec<(i32, signal_hook::SigId)>);
+    pub struct Deregister {
+        do_reset: bool,
+    }
     pub struct AutoDeregister(Deregister);
 
     impl Deregister {
-        /// Remove all previously registered handlers, and assure the default behaviour is reinstated.
+        /// Remove all previously registered handlers, and assure the default behaviour is reinstated, if this is the last available instance.
         ///
         /// Note that only the instantiation of the default behaviour can fail.
         pub fn deregister(self) -> std::io::Result<()> {
-            if self.0.is_empty() {
+            let mut hooks = REGISTERED_HOOKS.lock();
+            let count = DEREGISTER_COUNT.fetch_sub(1, Ordering::SeqCst);
+            if count > 1 || hooks.is_empty() {
                 return Ok(());
             }
-            static REINSTATE_DEFAULT_BEHAVIOUR: AtomicBool = AtomicBool::new(true);
-            for (_, hook_id) in &self.0 {
+            if self.do_reset {
+                super::reset();
+            }
+            for (_, hook_id) in hooks.iter() {
                 signal_hook::low_level::unregister(*hook_id);
             }
-            IS_INITIALIZED.store(false, Ordering::SeqCst);
-            if REINSTATE_DEFAULT_BEHAVIOUR
-                .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |_| Some(false))
-                .expect("always returns value")
-            {
-                for (sig, _) in self.0 {
-                    // # SAFETY
-                    // * we only call a handler that is specifically designed to run in this environment.
-                    #[allow(unsafe_code)]
-                    unsafe {
-                        signal_hook::low_level::register(sig, move || {
-                            signal_hook::low_level::emulate_default_handler(sig).ok();
-                        })?;
-                    }
+
+            let hooks = hooks.drain(..);
+            let mut default_hooks = DEFAULT_BEHAVIOUR_HOOKS.lock();
+            // Even if dropped, `drain(..)` clears the vec which is a must.
+            for (sig, _) in hooks {
+                // # SAFETY
+                // * we only register a handler that is specifically designed to run in this environment.
+                #[allow(unsafe_code)]
+                unsafe {
+                    default_hooks.push(signal_hook::low_level::register(sig, move || {
+                        signal_hook::low_level::emulate_default_handler(sig).ok();
+                    })?);
                 }
             }
             Ok(())
         }
 
+        /// If called with `toggle` being `true`, when actually deregistering, we will also reset the trigger by
+        /// calling [`reset()`](super::reset()).
+        pub fn with_reset(mut self, toggle: bool) -> Self {
+            self.do_reset = toggle;
+            self
+        }
+
         /// Return a type that deregisters all installed signal handlers on drop.
         pub fn auto_deregister(self) -> AutoDeregister {
             AutoDeregister(self)
@@ -60,20 +76,33 @@ mod init {
         }
     }
 
-    /// Initialize a signal handler to listen to SIGINT and SIGTERM and trigger our [`trigger()`][super::trigger()] that way.
-    /// Also trigger `interrupt()` which promises to never use a Mutex, allocate or deallocate.
+    /// Initialize a signal handler to listen to SIGINT and SIGTERM and trigger our [`trigger()`](super::trigger()) that way.
+    /// Also trigger `interrupt()` which promises to never use a Mutex, allocate or deallocate, or do anything else that's blocking.
+    /// Use `grace_count` to determine how often the termination signal can be received before it's terminal, e.g. 1 would only terminate
+    /// the application the second time the signal is received.
+    /// Note that only the `grace_count` and `interrupt` of the first call are effective, all others will be ignored.
+    ///
+    /// Use the returned `Deregister` type to explicitly deregister hooks, or to do so automatically.
     ///
     /// # Note
     ///
     /// It will abort the process on second press and won't inform the user about this behaviour either as we are unable to do so without
     /// deadlocking even when trying to write to stderr directly.
-    pub fn init_handler(interrupt: impl Fn() + Send + Sync + Clone + 'static) -> io::Result<Deregister> {
-        if IS_INITIALIZED
-            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |_| Some(true))
-            .expect("always returns value")
-        {
-            return Err(io::Error::new(io::ErrorKind::Other, "Already initialized"));
+    pub fn init_handler(
+        grace_count: usize,
+        interrupt: impl Fn() + Send + Sync + Clone + 'static,
+    ) -> io::Result<Deregister> {
+        let prev_count = DEREGISTER_COUNT.fetch_add(1, Ordering::SeqCst);
+        if prev_count != 0 {
+            // Try to obtain the lock before we return just to wait for the signals to actually be registered.
+            let _guard = REGISTERED_HOOKS.lock();
+            return Ok(Deregister::default());
+        }
+        let mut guard = REGISTERED_HOOKS.lock();
+        if !guard.is_empty() {
+            return Ok(Deregister::default());
         }
+
         let mut hooks = Vec::with_capacity(signal_hook::consts::TERM_SIGNALS.len());
         for sig in signal_hook::consts::TERM_SIGNALS {
             // # SAFETY
@@ -88,7 +117,7 @@ mod init {
                         INTERRUPT_COUNT.store(0, Ordering::SeqCst);
                     }
                     let msg_idx = INTERRUPT_COUNT.fetch_add(1, Ordering::SeqCst);
-                    if msg_idx == 1 {
+                    if msg_idx == grace_count {
                         gix_tempfile::registry::cleanup_tempfiles_signal_safe();
                         signal_hook::low_level::emulate_default_handler(*sig).ok();
                     }
@@ -98,13 +127,19 @@ mod init {
                 hooks.push((*sig, hook_id));
             }
         }
+        for hook_id in DEFAULT_BEHAVIOUR_HOOKS.lock().drain(..) {
+            signal_hook::low_level::unregister(hook_id);
+        }
 
         // This means that they won't setup a handler allowing us to call them right before we actually abort.
         gix_tempfile::signal::setup(gix_tempfile::signal::handler::Mode::None);
 
-        Ok(Deregister(hooks))
+        *guard = hooks;
+        Ok(Deregister::default())
     }
 }
+pub use init::Deregister;
+
 use std::{
     io,
     sync::atomic::{AtomicBool, Ordering},

Diff for: gix/src/interrupt.rs

@@ -0,0 +1,58 @@
+use signal_hook::consts::SIGTERM;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+#[test]
+fn multi_registration() -> gix_testtools::Result {
+    static V1: AtomicUsize = AtomicUsize::new(0);
+    static V2: AtomicBool = AtomicBool::new(false);
+
+    let reg1 = gix::interrupt::init_handler(3, || {
+        V1.fetch_add(1, Ordering::SeqCst);
+    })
+    .expect("succeeds");
+    assert!(!gix::interrupt::is_triggered());
+    assert_eq!(V1.load(Ordering::Relaxed), 0);
+    let reg2 =
+        gix::interrupt::init_handler(2, || V2.store(true, Ordering::SeqCst)).expect("multi-initialization is OK");
+    assert!(!V2.load(Ordering::Relaxed));
+
+    signal_hook::low_level::raise(SIGTERM).expect("signal can be raised");
+    assert!(gix::interrupt::is_triggered(), "this happens automatically");
+    assert_eq!(V1.load(Ordering::Relaxed), 1, "the first trigger is invoked");
+    assert!(!V2.load(Ordering::Relaxed), "the second trigger was ignored");
+
+    reg1.deregister()?;
+    signal_hook::low_level::raise(SIGTERM).expect("signal can be raised");
+    assert_eq!(V1.load(Ordering::Relaxed), 2, "the first trigger is still invoked");
+
+    assert!(gix::interrupt::is_triggered(), "this happens automatically");
+    // now the registration is actually removed.
+    reg2.with_reset(true).deregister()?;
+    assert!(
+        !gix::interrupt::is_triggered(),
+        "the deregistration succeeded and this is an optional side-effect"
+    );
+
+    let reg1 = gix::interrupt::init_handler(3, || {
+        V1.fetch_add(1, Ordering::SeqCst);
+    })
+    .expect("succeeds");
+    assert_eq!(V1.load(Ordering::Relaxed), 2, "nothing changed yet");
+    let reg2 =
+        gix::interrupt::init_handler(2, || V2.store(true, Ordering::SeqCst)).expect("multi-initialization is OK");
+    assert!(!V2.load(Ordering::Relaxed));
+
+    signal_hook::low_level::raise(SIGTERM).expect("signal can be raised");
+    assert_eq!(V1.load(Ordering::Relaxed), 3, "the first trigger is invoked");
+    assert!(!V2.load(Ordering::Relaxed), "the second trigger was ignored");
+
+    reg2.auto_deregister();
+    reg1.with_reset(true).auto_deregister();
+
+    assert!(
+        !gix::interrupt::is_triggered(),
+        "the deregistration succeeded and this is an optional side-effect"
+    );
+
+    Ok(())
+}

Diff for: gix/tests/interrupt.rs

@@ -129,7 +129,7 @@ pub fn main() -> Result<()> {
     let auto_verbose = !progress && !args.no_verbose;
 
     let should_interrupt = Arc::new(AtomicBool::new(false));
-    gix::interrupt::init_handler({
+    gix::interrupt::init_handler(1, {
         let should_interrupt = Arc::clone(&should_interrupt);
         move || should_interrupt.store(true, Ordering::SeqCst)
     })?;

Diff for: src/plumbing/main.rs

@@ -20,7 +20,7 @@ pub fn main() -> Result<()> {
         time::util::local_offset::set_soundness(time::util::local_offset::Soundness::Unsound);
     }
     let should_interrupt = Arc::new(AtomicBool::new(false));
-    gix::interrupt::init_handler({
+    gix::interrupt::init_handler(1, {
         let should_interrupt = Arc::clone(&should_interrupt);
         move || should_interrupt.store(true, Ordering::SeqCst)
     })?;