sigterm 0.3.10

/* src/signal.rs */

//! Signal handling primitives.

use tokio::signal;

/// Represents a system signal.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Signal {
	/// Ctrl+C (SIGINT)
	Int,
	/// SIGTERM
	Term,
	/// SIGQUIT
	Quit,
	/// SIGHUP
	Hup,
	/// SIGUSR1
	User1,
	/// SIGUSR2
	User2,
}

/// Attempts to wait for a shutdown signal (Ctrl+C or SIGTERM).
///
/// This function returns a future that resolves when either:
/// - `Ctrl+C` is received (all platforms)
/// - `SIGTERM` is received (Unix-like systems)
///
/// Returns `std::io::Result::Err` if the signal handlers cannot be installed.
pub async fn try_wait() -> std::io::Result<()> {
	#[cfg(feature = "tracing")]
	tracing::debug!("waiting for shutdown signal (Ctrl+C or SIGTERM)");

	let ctrl_c = signal::ctrl_c();

	#[cfg(unix)]
	let mut terminate = signal::unix::signal(signal::unix::SignalKind::terminate())?;

	#[cfg(not(unix))]
	let terminate = std::future::pending::<()>();

	tokio::select! {
			res = ctrl_c => {
					#[cfg(feature = "tracing")]
					tracing::info!("received Ctrl+C signal");
					res
			},
			_ = async {
					#[cfg(unix)]
					{
							terminate.recv().await;
					}
					#[cfg(not(unix))]
					{
							let () = terminate.await;
					}
			} => {
					#[cfg(feature = "tracing")]
					tracing::info!("received SIGTERM signal");
					Ok(())
			},
	}
}

/// Waits for a shutdown signal (Ctrl+C or SIGTERM).
///
/// # Panics
///
/// This function panics if the signal handlers cannot be installed.
/// Use [`try_wait`] for a non-panicking version.
pub async fn wait() {
	try_wait().await.expect("failed to install signal handler");
}

/// Waits for any of the specified signals.
///
/// This function listens for a list of signals and returns the signal that was received.
///
/// **Note:** This function requires the `rt` and `sync` features to be enabled on `tokio`.
#[cfg(all(unix, feature = "rt", feature = "sync"))]
pub async fn wait_for(signals: &[Signal]) -> std::io::Result<Signal> {
	use signal::unix::{SignalKind, signal};
	use tokio::sync::mpsc;

	#[cfg(feature = "tracing")]
	tracing::debug!(?signals, "waiting for specified signals");

	let (tx, mut rx) = mpsc::channel(1);

	// Pre-register all signals to ensure we can listen to them.
	// If any fail, we return error immediately.
	let mut listeners = Vec::with_capacity(signals.len());

	for &sig in signals {
		let kind = match sig {
			Signal::Int => SignalKind::interrupt(),
			Signal::Term => SignalKind::terminate(),
			Signal::Quit => SignalKind::quit(),
			Signal::Hup => SignalKind::hangup(),
			Signal::User1 => SignalKind::user_defined1(),
			Signal::User2 => SignalKind::user_defined2(),
		};

		// This fails fast if we can't register the signal
		let stream = signal(kind)?;
		listeners.push((sig, stream));
	}

	for (sig, mut stream) in listeners {
		let tx = tx.clone();
		tokio::spawn(async move {
			stream.recv().await;
			let _ = tx.send(sig).await;
		});
	}

	// Drop original sender so channel doesn't hang
	drop(tx);

	// This unwrap is safe because we hold the receiver and have ensured listeners are registered.
	// If all tasks die without sending (impossible unless runtime dies?), it returns None.
	// But since we are waiting for external signals, valid execution implies we wait here.
	// If rx returns None, it means all senders dropped, which implies all spawned tasks finished?
	// But they only finish if they receive a signal.
	// So this await effectively waits for the first signal.
	let received = rx
		.recv()
		.await
		.expect("signal listeners task failed unexpectedly");

	#[cfg(feature = "tracing")]
	tracing::info!(?received, "received signal");

	Ok(received)
}

/// Returns a future that resolves when a system shutdown signal is received or
/// the provided cancellation token is cancelled.
///
/// This is a helper designed to be passed directly to server frameworks like Axum:
/// `.with_graceful_shutdown(sigterm::shutdown_signal(token))`
#[cfg(feature = "cancel")]
pub async fn shutdown_signal(token: crate::CancellationToken) {
	#[cfg(feature = "tracing")]
	tracing::debug!("waiting for shutdown signal or cancellation token");

	tokio::select! {
			_ = wait() => {
					#[cfg(feature = "tracing")]
					tracing::info!("shutdown triggered by system signal");
			},
			_ = token.cancelled() => {
					#[cfg(feature = "tracing")]
					tracing::info!("shutdown triggered by cancellation token");
			},
	}
}

/// Waits for a shutdown signal with a timeout.
///
/// Returns `Ok(())` if a signal was received, or `Err(Elapsed)` if the timeout expired.
pub async fn wait_timeout(
	duration: std::time::Duration,
) -> Result<(), tokio::time::error::Elapsed> {
	#[cfg(feature = "tracing")]
	tracing::debug!(?duration, "waiting for shutdown signal with timeout");

	let result = tokio::time::timeout(duration, wait()).await;

	#[cfg(feature = "tracing")]
	if result.is_err() {
		tracing::debug!("shutdown wait timed out");
	}

	result
}