better_graceful/
lib.rs

1//! Gracefully shutdown.
2//!
3//! (Unix) Async signals are tricky to handle properly. This utility let you
4//! block the main thread and execute arbitrary (thread-safe) code to shutdown
5//! the process gracefully.
6//!
7//! # Usage
8//!
9//! 1. Initialize a [SignalGuard](struct.SignalGuard.html) before creating any
10//!    additional threads.
11//! 2. The [SignalGuard](struct.SignalGuard.html) will block necessary signals
12//!    (`SIGINT`, `SIGQUIT` and `SIGTERM` on *nix, `Ctrl+C` and `Ctrl+Break` on
13//!    Windows) during initialization.
14//! 3. Spawn new threads to do the real work.
15//! 4. Register a handle to properly shutdown the application.
16//! 5. The main thread will be blocked until a signal is received.
17//! 6. The handler will run in the main thread.
18//! 7. On Windows the process will terminate after the handler returns (and
19//!    potentially any libc `atexit` handlers).
20//!
21//! # Example
22//!
23//! ```no_run
24//! extern crate graceful;
25//!
26//! use std::sync::atomic::{ATOMIC_BOOL_INIT, AtomicBool, Ordering};
27//! use std::time::Duration;
28//! use std::thread;
29//!
30//! use graceful::SignalGuard;
31//!
32//! static STOP: AtomicBool = ATOMIC_BOOL_INIT;
33//!
34//! fn main() {
35//!     let signal_guard = SignalGuard::new();
36//!
37//! 	let handle = thread::spawn(|| {
38//!         println!("Worker thread started. Type Ctrl+C to stop.");
39//!         while !STOP.load(Ordering::Acquire) {
40//!             println!("working...");
41//!             thread::sleep(Duration::from_millis(500));
42//!         }
43//!         println!("Bye.");
44//!     });
45//!
46//! 	signal_guard.at_exit(move |sig| {
47//!         println!("Signal {} received.", sig);
48//!         STOP.store(true, Ordering::Release);
49//!         handle.join().unwrap();
50//!     });
51//! }
52//! ```
53//!
54
55#[cfg(unix)]
56mod platform {
57    extern crate nix;
58    use self::nix::sys::signal::{SigSet, SIGINT, SIGQUIT, SIGTERM};
59
60    pub struct SignalGuard(SigSet);
61
62    impl SignalGuard {
63        /// Block necessary signals (`SIGINT`, `SIGQUIT` and `SIGTERM` on *nix,
64        /// `Ctrl+C` and `Ctrl+Break` on Windows).
65        ///
66        /// New threads should be spawned after this.
67        pub fn new() -> SignalGuard {
68            let mut mask = SigSet::empty();
69            SignalGuard::init(&mut mask).unwrap();
70            SignalGuard(mask)
71        }
72
73        fn init(mask: &mut SigSet) -> nix::Result<()> {
74            mask.add(SIGINT);
75            mask.add(SIGQUIT);
76            mask.add(SIGTERM);
77            mask.thread_block()
78        }
79
80        /// Block the running thread until a signal is received. Then the
81        /// `handler` will be called in the main thread.
82        ///
83        /// Do not put any code after this.
84        pub fn at_exit<F: FnOnce(usize)>(&self, handler: F) {
85            let sig = self.0.wait().unwrap();
86            handler(sig as usize);
87        }
88    }
89}
90
91#[cfg(windows)]
92#[macro_use]
93extern crate lazy_static;
94
95#[cfg(windows)]
96mod platform {
97    extern crate winapi;
98
99    use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
100    use std::sync::Mutex;
101
102    use self::winapi::shared::minwindef::{BOOL, DWORD, TRUE};
103    use self::winapi::um::consoleapi::SetConsoleCtrlHandler;
104
105    lazy_static! {
106        static ref CHAN: (SyncSender<DWORD>, Mutex<Receiver<DWORD>>) = {
107            let channel = sync_channel(0);
108            (channel.0, Mutex::new(channel.1))
109        };
110    }
111
112    unsafe extern "system" fn handler(event: DWORD) -> BOOL {
113        CHAN.0.send(event).unwrap();
114        CHAN.0.send(0).unwrap();
115        TRUE
116    }
117
118    pub struct SignalGuard;
119
120    impl SignalGuard {
121        pub fn new() -> SignalGuard {
122            unsafe { SetConsoleCtrlHandler(Some(handler), TRUE) };
123            SignalGuard
124        }
125
126        pub fn at_exit<F: FnOnce(usize)>(&self, handler: F) {
127            let event = {
128                let receiver = CHAN.1.lock().unwrap();
129                receiver.recv().unwrap()
130            };
131            handler(event as usize);
132            CHAN.1.lock().unwrap().recv().unwrap();
133        }
134    }
135}
136
137pub use platform::SignalGuard;
better_graceful/lib.rs

better_graceful/
lib.rs
