69 lines
2.5 KiB
Rust
69 lines
2.5 KiB
Rust
use std::io;
|
||
use std::sync::atomic::AtomicBool;
|
||
use std::sync::{Arc, Mutex};
|
||
|
||
use signal_hook::consts::TERM_SIGNALS;
|
||
use signal_hook::flag;
|
||
use signal_hook::iterator::{Signals, SignalsInfo};
|
||
|
||
use crate::server;
|
||
|
||
/// Install the required signal handlers for terminating signals.
|
||
pub fn install_signal_handlers() -> io::Result<(Arc<AtomicBool>, SignalsInfo)> {
|
||
let term = Arc::new(AtomicBool::new(false));
|
||
|
||
// For each terminating signal, we register both a shutdown handler and a handler that sets an
|
||
// atomic bool. With this, the process will get killed immediately once it receives a second
|
||
// termination signal (e.g. a double ctrl-c).
|
||
// https://docs.rs/signal-hook/0.3.15/signal_hook/#a-complex-signal-handling-with-a-background-thread
|
||
for sig in TERM_SIGNALS {
|
||
// When terminated by a second term signal, exit with exit code 1.
|
||
// This will do nothing the first time (because term_now is false).
|
||
flag::register_conditional_shutdown(*sig, 1, Arc::clone(&term))?;
|
||
// But this will "arm" the above for the second time, by setting it to true.
|
||
// The order of registering these is important, if you put this one first, it will
|
||
// first arm and then terminate ‒ all in the first round.
|
||
flag::register(*sig, Arc::clone(&term))?;
|
||
}
|
||
|
||
let signals = TERM_SIGNALS;
|
||
|
||
Ok((term, Signals::new(signals)?))
|
||
}
|
||
|
||
/// Loop that handles terminating signals as they come in.
|
||
pub fn handle_signals(
|
||
signals: &mut SignalsInfo,
|
||
counter: Arc<Mutex<server::ServerProcess>>,
|
||
) -> io::Result<()> {
|
||
let mut force = false;
|
||
|
||
// We only register terminating signals, so we don't need to differentiate between what kind of
|
||
// signal came in
|
||
for _ in signals {
|
||
// If term is already true, this is the second signal, meaning we kill the process
|
||
// immediately.
|
||
// This will currently not work, as the initial stop command will block the kill from
|
||
// happening.
|
||
if force {
|
||
let mut server = counter.lock().unwrap();
|
||
return server.kill();
|
||
}
|
||
// The stop command runs in a separate thread to avoid blocking the signal handling loop.
|
||
// After stopping the server, the thread terminates the process.
|
||
else {
|
||
let clone = Arc::clone(&counter);
|
||
|
||
std::thread::spawn(move || {
|
||
let mut server = clone.lock().unwrap();
|
||
let _ = server.stop();
|
||
std::process::exit(0);
|
||
});
|
||
}
|
||
|
||
force = true;
|
||
}
|
||
|
||
Ok(())
|
||
}
|