ptth/src/graceful_shutdown.rs

104 lines
1.8 KiB
Rust
Raw Normal View History

use std::{
cell::Cell,
time::Duration,
};
use futures::prelude::*;
use tokio::{
sync::oneshot,
time::delay_for,
};
use crate::prelude::*;
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)
}