tracel_xtask/utils/
cleanup.rs

1use std::sync::{Arc, LazyLock, Mutex};
2
3pub static CLEANUP_HANDLER: LazyLock<CleanupHandler> = LazyLock::new(CleanupHandler::new);
4
5pub struct RegisteredCleanupFunction {
6    pub handler: Box<dyn FnOnce() + Send + Sync>,
7    pub name: String,
8}
9
10#[derive(Clone)]
11pub struct CleanupHandler {
12    registered: Arc<Mutex<Vec<RegisteredCleanupFunction>>>,
13    terminated: Arc<Mutex<bool>>,
14}
15
16impl CleanupHandler {
17    fn new() -> Self {
18        let handler = CleanupHandler {
19            registered: Arc::new(Mutex::new(Vec::new())),
20            terminated: Arc::new(Mutex::new(false)),
21        };
22
23        let mut handler_ = handler.clone();
24
25        ctrlc::set_handler(move || {
26            if !handler_.registered.lock().unwrap().is_empty() {
27                println!();
28                warn!("Termination signal received, executing registered functions.");
29                handler_.terminate();
30            }
31            std::process::exit(1);
32        })
33        .expect("Should be able to set termination handler");
34
35        handler
36    }
37
38    pub fn register(
39        &self,
40        name: impl Into<String> + Clone,
41        handler: impl FnOnce() + Send + Sync + 'static,
42    ) {
43        trace!("Registering cleanup function for {}", name.clone().into());
44        self.registered
45            .lock()
46            .unwrap()
47            .push(RegisteredCleanupFunction {
48                handler: Box::new(handler),
49                name: name.into(),
50            });
51    }
52
53    fn terminate(&mut self) {
54        let mut terminated = self.terminated.lock().unwrap();
55        if *terminated {
56            return;
57        }
58        *terminated = true;
59        for f in (*self.registered.lock().unwrap()).drain(..) {
60            info!("Executing cleanup function: {}", f.name);
61            (f.handler)();
62        }
63    }
64}
65
66/// Since the CLEANUP_HANDLER is a static variable, the drop function is not automatically called when the program exits. It needs to be manually called with the `handle_cleanup!` macro.
67impl Drop for CleanupHandler {
68    fn drop(&mut self) {
69        if !self.registered.lock().unwrap().is_empty() {
70            println!();
71            warn!("Cleanup handler dropped, executing registered functions.");
72            self.terminate();
73        }
74    }
75}
76
77#[macro_export]
78macro_rules! register_cleanup {
79    ($name:expr, $handler:expr) => {
80        CLEANUP_HANDLER.register($name, $handler);
81    };
82}
83
84#[macro_export]
85macro_rules! handle_cleanup {
86    () => {
87        drop(CLEANUP_HANDLER.clone());
88    };
89}
tracel_xtask/utils/cleanup.rs

tracel_xtask/utils/
cleanup.rs
