111 lines
2.1 KiB
Rust
111 lines
2.1 KiB
Rust
// False positive on futures::select! macro
|
|
#![allow (clippy::mut_mut)]
|
|
|
|
use std::{
|
|
cell::Cell,
|
|
time::Duration,
|
|
};
|
|
|
|
use futures::prelude::*;
|
|
use tokio::{
|
|
sync::oneshot,
|
|
time::sleep,
|
|
};
|
|
|
|
use crate::prelude::*;
|
|
|
|
#[must_use]
|
|
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 (()).expect ("Error sending Ctrl-C to program");
|
|
}
|
|
}).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 {
|
|
let desc = match self {
|
|
ShutdownError::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 {
|
|
/// # Errors
|
|
///
|
|
/// `ForcedShutdown` if the graceful shutdown doesn't complete in time
|
|
|
|
pub async fn wrap_server <
|
|
T,
|
|
F: Future <Output = T>
|
|
> (
|
|
self,
|
|
server: F
|
|
) -> Result <T, ShutdownError> {
|
|
let fut = async move {
|
|
self.rx.await.expect ("Error awaiting graceful shutdown signal");
|
|
self.tx.send (()).expect ("Error forwarding graceful shutdown signal");
|
|
let timeout = 5;
|
|
debug! ("Starting graceful shutdown. Forcing shutdown in {} seconds", timeout);
|
|
sleep (Duration::from_secs (timeout)).await;
|
|
|
|
error! ("Forcing shutdown");
|
|
};
|
|
|
|
futures::select! {
|
|
x = server.fuse () => {
|
|
info! ("Shut down gracefully");
|
|
Ok (x)
|
|
},
|
|
_ = fut.fuse () => Err (ShutdownError::ForcedShutdown),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn init_with_force () -> (oneshot::Receiver <()>, ForcedShutdown) {
|
|
let (tx, rx) = oneshot::channel ();
|
|
|
|
let f = ForcedShutdown {
|
|
rx: init (),
|
|
tx,
|
|
};
|
|
|
|
(rx, f)
|
|
}
|