cli_batteries/
shutdown.rs

1use once_cell::sync::Lazy;
2use tokio::sync::watch::{self, Receiver, Sender};
3
4#[cfg(feature = "signals")]
5use eyre::Result as EyreResult;
6#[cfg(feature = "signals")]
7use tracing::{error, info};
8
9static NOTIFY: Lazy<(Sender<bool>, Receiver<bool>)> = Lazy::new(|| watch::channel(false));
10
11/// Send the signal to shutdown the program.
12#[allow(clippy::missing_panics_doc)] // Never panics
13pub fn shutdown() {
14    // Does not fail because the channel never closes.
15    NOTIFY.0.send(true).unwrap();
16}
17
18/// Reset the shutdown signal so it can be triggered again.
19///
20/// This is only useful for testing. Strange things can happen to any existing
21/// `await_shutdown()` futures.
22#[cfg(feature = "mock-shutdown")]
23#[allow(clippy::missing_panics_doc)] // Never panics
24#[allow(clippy::module_name_repetitions)] // Never panics
25pub fn reset_shutdown() {
26    // Does not fail because the channel never closes.
27    NOTIFY.0.send(false).unwrap();
28}
29
30/// Are we currently shutting down?
31#[must_use]
32pub fn is_shutting_down() -> bool {
33    *NOTIFY.1.borrow()
34}
35
36/// Wait for the program to shutdown.
37///
38/// Resolves immediately if the program is already shutting down.
39/// The resulting future is safe to cancel by dropping.
40#[allow(clippy::module_name_repetitions)]
41#[allow(clippy::missing_panics_doc)]
42pub async fn await_shutdown() {
43    let mut watch = NOTIFY.1.clone();
44    if *watch.borrow_and_update() {
45        return;
46    }
47    // Does not fail because the channel never closes.
48    watch.changed().await.unwrap();
49}
50
51#[cfg(feature = "signals")]
52pub fn watch_signals() {
53    tokio::spawn({
54        async move {
55            signal_shutdown()
56                .await
57                .map_err(|err| error!("Error handling Ctrl-C: {}", err))
58                .unwrap();
59            shutdown();
60        }
61    });
62}
63
64#[cfg(all(unix, feature = "signals"))]
65#[allow(clippy::module_name_repetitions)]
66async fn signal_shutdown() -> EyreResult<()> {
67    use tokio::signal::unix::{signal, SignalKind};
68
69    let sigint = signal(SignalKind::interrupt())?;
70    let sigterm = signal(SignalKind::terminate())?;
71    tokio::pin!(sigint);
72    tokio::pin!(sigterm);
73    tokio::select! {
74        _ = sigint.recv() => { info!("SIGINT received, shutting down"); }
75        _ = sigterm.recv() => { info!("SIGTERM received, shutting down"); }
76    };
77    Ok(())
78}
79
80#[cfg(all(not(unix), feature = "signals"))]
81#[allow(clippy::module_name_repetitions)]
82async fn signal_shutdown() -> EyreResult<()> {
83    use tokio::signal::ctrl_c;
84
85    ctrl_c().await?;
86    info!("Ctrl-C received, shutting down");
87    Ok(())
88}
cli_batteries/shutdown.rs

cli_batteries/
shutdown.rs
