103 lines
1.9 KiB
Rust
103 lines
1.9 KiB
Rust
use std::{
|
|
cell::Cell,
|
|
time::Duration,
|
|
};
|
|
|
|
use futures::prelude::*;
|
|
use tokio::{
|
|
sync::oneshot,
|
|
time::delay_for,
|
|
};
|
|
use tracing::{debug, error, info, trace, warn};
|
|
|
|
pub fn init () -> oneshot::Receiver <()> {
|
|
let (tx, rx) = oneshot::channel::<()> ();
|
|
|
|
// I have to put the tx into a Cell here so that if Ctrl-C gets
|
|
// called multiple times, we won't send multiple shutdowns to the
|
|
// oneshot channel. (Which would be a compile error)
|
|
|
|
let tx = Some (tx);
|
|
let tx = Cell::new (tx);
|
|
|
|
ctrlc::set_handler (move || {
|
|
let tx = tx.replace (None);
|
|
|
|
if let Some (tx) = tx {
|
|
tx.send (()).unwrap ();
|
|
}
|
|
}).expect ("Error setting Ctrl-C handler");
|
|
|
|
rx
|
|
}
|
|
|
|
#[derive (Debug)]
|
|
pub enum ShutdownError {
|
|
ForcedShutdown,
|
|
}
|
|
|
|
use std::{
|
|
error,
|
|
fmt,
|
|
};
|
|
|
|
impl fmt::Display for ShutdownError {
|
|
fn fmt (&self, f: &mut fmt::Formatter <'_>) -> fmt::Result {
|
|
use ShutdownError::*;
|
|
|
|
let desc = match self {
|
|
ForcedShutdown => "Shutdown was forced after a timeout",
|
|
};
|
|
|
|
write! (f, "{}", desc)
|
|
}
|
|
}
|
|
|
|
impl error::Error for ShutdownError {
|
|
|
|
}
|
|
|
|
pub struct ForcedShutdown {
|
|
rx: oneshot::Receiver <()>,
|
|
tx: oneshot::Sender <()>,
|
|
}
|
|
|
|
impl ForcedShutdown {
|
|
pub async fn wrap_server <
|
|
T,
|
|
F: Future <Output = T>
|
|
> (
|
|
self,
|
|
server: F
|
|
) -> Result <T, ShutdownError> {
|
|
let fut = async move {
|
|
self.rx.await.unwrap ();
|
|
self.tx.send (()).unwrap ();
|
|
let timeout = 5;
|
|
debug! ("Starting graceful shutdown. Forcing shutdown in {} seconds", timeout);
|
|
delay_for (Duration::from_secs (timeout)).await;
|
|
|
|
error! ("Forcing shutdown");
|
|
};
|
|
|
|
futures::select! {
|
|
x = server.fuse () => {
|
|
info! ("Shut down gracefully");
|
|
Ok (x)
|
|
},
|
|
_ = fut.fuse () => Err (ShutdownError::ForcedShutdown),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn init_with_force () -> (oneshot::Receiver <()>, ForcedShutdown) {
|
|
let (tx, rx) = oneshot::channel ();
|
|
|
|
let f = ForcedShutdown {
|
|
rx: init (),
|
|
tx,
|
|
};
|
|
|
|
(rx, f)
|
|
}
|