tokio_shutdown/
lib.rs

1//! # Tokio Shutdown
2//!
3//! Tiny crate that allows to wait for a stop signal across multiple threads. Helpful mostly in
4//! server applications that run indefinitely and need a signal for graceful shutdowns.
5//!
6//! # Examples
7//!
8//! This example installs the global shutdown handler and will print a message once a single is
9//! received. For demonstration purposes it creates other tasks that wait for shutdown as well.
10//!
11//! ```no_run
12//! use tokio_shutdown::Shutdown;
13//!
14//! #[tokio::main(flavor = "current_thread")]
15//! async fn main() {
16//!     let shutdown = Shutdown::new().expect("shutdown creation works on first call");
17//!
18//!     // Pass a copy of the shutdown handler to another task.
19//!     // Clones of `Shutdown` are cheap.
20//!     let clone = shutdown.clone();
21//!     tokio::spawn(async move {
22//!         clone.handle().await;
23//!         println!("task 1 shutting down");
24//!     });
25//!
26//!     // Alternatively, pass a new handle to the new task.
27//!     // Both this and the above way work, choose whatever works best for your use case.
28//!     let handle = shutdown.handle();
29//!     tokio::spawn(async move {
30//!         handle.await;
31//!         println!("task 2 shutting down");
32//!     });
33//!
34//!     shutdown.handle().await;
35//!     println!("application shutting down");
36//! }
37//! ```
38//!
39//! Please have a look at the examples directory of this project for further usage instructions.
40
41#![forbid(unsafe_code)]
42#![deny(rust_2018_idioms, clippy::all, clippy::pedantic)]
43
44use std::{
45    error::Error,
46    fmt::{self, Display},
47    future::Future,
48    sync::atomic::{AtomicBool, Ordering},
49};
50
51use tokio::{signal, sync::watch};
52
53/// Error that occurs when the [`Shutdown::new`] function is called more than once in a process
54/// lifetime.
55#[derive(Debug, PartialEq, Eq)]
56pub struct AlreadyCreatedError;
57
58impl Error for AlreadyCreatedError {}
59
60impl Display for AlreadyCreatedError {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        f.write_str("shutdown handler already created")
63    }
64}
65
66static CREATED: AtomicBool = AtomicBool::new(false);
67
68/// The global shutdown handler for an application. It can be cloned cheaply wherever needed.
69///
70/// New handles can be created with the [`handle`](Self::handle) function, which creates futures
71/// that will complete once a shutdown signal is received.
72#[derive(Clone)]
73pub struct Shutdown {
74    receiver: watch::Receiver<()>,
75}
76
77impl Shutdown {
78    /// Create a new shutdown handle. This can only be called once per application instance.
79    ///
80    /// Signal handles can only be registered once for the duration of the entire process and
81    /// creating another shutdown handler would break the previous one without notice.
82    ///
83    /// # Errors
84    ///
85    /// If this function is called more than once during the lifetime of a process, an error will be
86    /// returned.
87    pub fn new() -> Result<Shutdown, AlreadyCreatedError> {
88        if (CREATED).swap(true, Ordering::SeqCst) {
89            return Err(AlreadyCreatedError);
90        }
91
92        let (tx, rx) = watch::channel(());
93        let handle = register_handlers();
94
95        tokio::spawn(async move {
96            handle.await;
97            tx.send(()).ok();
98        });
99
100        Ok(Self { receiver: rx })
101    }
102
103    /// Create a new handle that can be awaited on. The future will complete once a shutdown signal
104    /// is received.
105    pub fn handle(&self) -> impl Future<Output = ()> {
106        let mut rx = self.receiver.clone();
107
108        async move {
109            rx.changed().await.ok();
110        }
111    }
112}
113
114fn register_handlers() -> impl Future<Output = ()> {
115    let ctrl_c = async {
116        signal::ctrl_c()
117            .await
118            .expect("failed to install Ctrl+C handler");
119    };
120
121    #[cfg(unix)]
122    let terminate = async {
123        signal::unix::signal(signal::unix::SignalKind::terminate())
124            .expect("failed to install signal handler")
125            .recv()
126            .await;
127    };
128
129    #[cfg(not(unix))]
130    let terminate = std::future::pending::<()>();
131
132    async {
133        tokio::select! {
134            () = ctrl_c => {},
135            () = terminate => {},
136        }
137
138        #[cfg(feature = "tracing")]
139        tracing::info!("shutdown signal received");
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[tokio::test]
148    async fn fail_create_two_instances() {
149        assert!(Shutdown::new().is_ok());
150        assert_eq!(Some(AlreadyCreatedError), Shutdown::new().err());
151    }
152}
tokio_shutdown/lib.rs

tokio_shutdown/
lib.rs
