ptth_core 2.0.0

// 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::*;

/// Initializes a graceful shutdown using the `ctrlc` crate

#[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 {
	
}

/// The type returned by `init_with_force`

pub struct ForcedShutdown {
	rx: oneshot::Receiver <()>,
	tx: oneshot::Sender <()>,
}

impl ForcedShutdown {
	/// Wraps a future in a graceful shutdown that times out into a 
	/// forced shutdown.
	/// 
	/// # 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),
		}
	}
}

/// Initializes a graceful shutdown that times out into a forced shutdown

#[must_use]
pub fn init_with_force () -> (oneshot::Receiver <()>, ForcedShutdown) {
	let (tx, rx) = oneshot::channel ();
	
	let f = ForcedShutdown {
		rx: init (),
		tx,
	};
	
	(rx, f)
}