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)]
92mod platform {
93    extern crate winapi;
94    extern crate kernel32;
95
96    use std::sync::mpsc::{sync_channel, SyncSender, Receiver};
97
98    use kernel32::SetConsoleCtrlHandler;
99    use winapi::{BOOL, DWORD, TRUE};
100
101    static CHAN: (SyncSender<DWORD>, Receiver<DWORD>) = sync_channel(0);
102
103    unsafe extern "system" fn handler(event: DWORD) -> BOOL {
104        CHAN.0.send(event);
105        CHAN.0.send(0);
106        FALSE
107    }
108
109    pub struct SignalGuard;
110
111    impl SignalGuard {
112        pub fn new() -> SignalGuard {
113            unsafe { SetConsoleCtrlHandler(Some(handler), TRUE) };
114            SignalGuard
115        }
116
117        pub fn at_exit<F: FnOnce(usize)>(&self, handler: F) {
118            let event = CHAN.1.recv().unwrap();
119            handler(event as usize);
120            CHAN.1.recv().unwrap();
121        }
122    }
123}
124
125pub use platform::SignalGuard;
graceful/lib.rs

graceful/
lib.rs
