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