sigterm 0.3.10

/* src/shutdown.rs */

//!
//! This module provides a simple mechanism for notifying a single task or
//! listener to shut down. It is built on top of `tokio::sync::oneshot`.

use tokio::sync::oneshot;

/// A receiver that waits for a shutdown signal.
///
/// This is typically moved into a background task's loop.
pub struct Shutdown {
	receiver: oneshot::Receiver<()>,
}

impl Shutdown {
	/// Creates a new shutdown token pair.
	pub fn new() -> (Self, ShutdownHandle) {
		#[cfg(feature = "tracing")]
		tracing::trace!("creating new shutdown channel");

		let (tx, rx) = oneshot::channel();
		(Self { receiver: rx }, ShutdownHandle { sender: tx })
	}

	/// Waits for the shutdown signal.
	///
	/// This resolves when [`ShutdownHandle::shutdown`] is called or the handle is dropped.
	pub async fn recv(self) {
		#[cfg(feature = "tracing")]
		tracing::debug!("waiting for shutdown signal");

		let _ = self.receiver.await;

		#[cfg(feature = "tracing")]
		tracing::debug!("shutdown signal received");
	}
}

/// A handle to trigger the shutdown of a task.
pub struct ShutdownHandle {
	sender: oneshot::Sender<()>,
}

impl ShutdownHandle {
	/// Triggers the shutdown signal.
	pub fn shutdown(self) {
		#[cfg(feature = "tracing")]
		tracing::info!("triggering shutdown");

		let _ = self.sender.send(());
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use std::time::Duration;
	use tokio::time::timeout;

	#[tokio::test]
	async fn test_shutdown_trigger() {
		let (shutdown, handle) = Shutdown::new();

		let task = tokio::spawn(async move {
			shutdown.recv().await;
		});

		handle.shutdown();

		// Should complete quickly
		assert!(timeout(Duration::from_millis(100), task).await.is_ok());
	}

	#[tokio::test]
	async fn test_shutdown_handle_drop() {
		let (shutdown, handle) = Shutdown::new();

		let task = tokio::spawn(async move {
			shutdown.recv().await;
		});

		// Dropping the handle should also unblock recv (channel closed)
		drop(handle);

		assert!(timeout(Duration::from_millis(100), task).await.is_ok());
	}
}