cdp_html_shot/
exit_hook.rs

1use std::panic;
2use std::sync::{Arc, Once};
3
4/// A struct that manages cleanup functions to be executed on process termination.
5pub struct ExitHook {
6    cleanup_fn: Arc<dyn Fn() + Send + Sync + 'static>,
7}
8
9impl ExitHook {
10    /**
11    Creates a new ExitHook with the given cleanup function.
12
13    # Arguments
14
15    * `f` - A function to be executed when the process exits
16
17    # Example
18
19    ```rust
20    use cdp_html_shot::ExitHook;
21    let hook = ExitHook::new(|| println!("Cleaning up..."));
22
23    hook.register().unwrap();
24    ```
25    */
26    pub fn new<F>(f: F) -> Self
27    where
28        F: Fn() + Send + Sync + 'static
29    {
30        ExitHook {
31            cleanup_fn: Arc::new(f),
32        }
33    }
34
35    /// Registers all necessary hooks for process termination.
36    pub fn register(&self) -> Result<(), Box<dyn std::error::Error>> {
37        static INIT: Once = Once::new();
38        let cleanup_fn = Arc::clone(&self.cleanup_fn);
39
40        // Set up panic hook
41        let original_hook = panic::take_hook();
42        let panic_cleanup_fn = Arc::clone(&cleanup_fn);
43        panic::set_hook(Box::new(move |panic_info| {
44            panic_cleanup_fn();
45            original_hook(panic_info);
46        }));
47
48        // Set up Ctrl+C handler
49        INIT.call_once(|| {
50            let ctrl_c_cleanup_fn = Arc::clone(&cleanup_fn);
51            if let Err(e) = ctrlc::set_handler(move || {
52                ctrl_c_cleanup_fn();
53                std::process::exit(0);
54            }) {
55                eprintln!("Error setting Ctrl-C handler: {}", e);
56            }
57        });
58
59        Ok(())
60    }
61}
62
63impl Drop for ExitHook {
64    fn drop(&mut self) {
65        (self.cleanup_fn)();
66    }
67}
cdp_html_shot/exit_hook.rs

cdp_html_shot/
exit_hook.rs
