Skip to main content

jj_cli/
cleanup_guard.rs

1use std::panic::AssertUnwindSafe;
2use std::sync::Mutex;
3
4use slab::Slab;
5use tracing::instrument;
6
7/// Contains the callbacks passed to currently-live [`CleanupGuard`]s
8static LIVE_GUARDS: Mutex<GuardTable> = Mutex::new(Slab::new());
9
10type GuardTable = Slab<Box<dyn FnOnce() + Send>>;
11
12/// Prepare to run [`CleanupGuard`]s on `SIGINT`/`SIGTERM`/`SIGHUP`
13pub fn init() {
14    if let Err(e) = ctrlc::set_handler(|| {
15        // We must hold the lock for the remainder of the process's lifetime to avoid a
16        // race where a guard is created after we unlock but before we exit.
17        let guards = &mut *LIVE_GUARDS.lock().unwrap();
18        if let Err(e) = std::panic::catch_unwind(AssertUnwindSafe(|| {
19            for guard in guards.drain() {
20                guard();
21            }
22        })) {
23            match e.downcast::<String>() {
24                Ok(s) => eprintln!("ctrlc handler panicked: {s}"),
25                Err(_) => eprintln!("ctrlc handler panicked"),
26            }
27        }
28
29        #[cfg(feature = "git")]
30        gix::tempfile::registry::cleanup_tempfiles();
31
32        std::process::exit(1);
33    }) {
34        eprintln!("couldn't register signal handler: {e}");
35    }
36}
37
38/// A drop guard that also runs on `SIGINT`/`SIGTERM`/`SIGHUP`
39pub struct CleanupGuard {
40    slot: usize,
41}
42
43impl CleanupGuard {
44    /// Invoke `f` when dropped or killed by `SIGINT`/`SIGTERM`/`SIGHUP`
45    pub fn new<F: FnOnce() + Send + 'static>(f: F) -> Self {
46        let guards = &mut *LIVE_GUARDS.lock().unwrap();
47        Self {
48            slot: guards.insert(Box::new(f)),
49        }
50    }
51}
52
53impl Drop for CleanupGuard {
54    #[instrument(skip_all)]
55    fn drop(&mut self) {
56        let guards = &mut *LIVE_GUARDS.lock().unwrap();
57        let f = guards.remove(self.slot);
58        f();
59    }
60}