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 > ( self, server: F ) -> Result { 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) }