ptth/src/graceful_shutdown.rs

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)
}