Skip to main content

chio_http_serve/
signal.rs

1use std::future::Future;
2
3use tokio::sync::watch;
4use tracing::error;
5
6/// Resolve on the first SIGTERM (unix) or Ctrl-C / SIGINT (all platforms).
7///
8/// The future is fail-closed against handler-install failure: a handler that
9/// cannot be installed logs loudly and yields to a pending branch rather than
10/// resolving, so a missing handler never triggers a spurious shutdown at
11/// startup. If both handlers fail to install the future simply never resolves
12/// and the process keeps running until the platform escalates to an
13/// unconditional kill, which is no worse than having no shutdown wiring at all.
14pub async fn shutdown_signal() {
15    let ctrl_c = async {
16        match tokio::signal::ctrl_c().await {
17            Ok(()) => {}
18            Err(source) => {
19                error!(%source, "cannot install Ctrl-C handler; the SIGTERM path still governs shutdown");
20                std::future::pending::<()>().await;
21            }
22        }
23    };
24
25    #[cfg(unix)]
26    let terminate = async {
27        use tokio::signal::unix::{signal, SignalKind};
28        match signal(SignalKind::terminate()) {
29            Ok(mut stream) => {
30                stream.recv().await;
31            }
32            Err(source) => {
33                error!(%source, "cannot install SIGTERM handler; the Ctrl-C path still governs shutdown");
34                std::future::pending::<()>().await;
35            }
36        }
37    };
38
39    #[cfg(not(unix))]
40    let terminate = std::future::pending::<()>();
41
42    tokio::select! {
43        _ = ctrl_c => {}
44        _ = terminate => {}
45    }
46}
47
48/// Owns the shutdown watch channel that fans one stop signal out to the serve
49/// loop and every cooperating background task. The receivers it hands out are
50/// cheap to clone.
51pub struct ShutdownController {
52    tx: watch::Sender<bool>,
53    rx: watch::Receiver<bool>,
54}
55
56impl ShutdownController {
57    /// Spawn the OS-signal task and return a controller that is live
58    /// immediately. Requires an active Tokio runtime; the spawned task
59    /// translates the first [`shutdown_signal`] into a `true` on the channel.
60    #[must_use]
61    pub fn install() -> Self {
62        let (tx, rx) = watch::channel(false);
63        let task_tx = tx.clone();
64        tokio::spawn(async move {
65            shutdown_signal().await;
66            // A closed channel (every receiver dropped) is not an error here: it
67            // just means there is nothing left to notify.
68            let _ = task_tx.send(true);
69        });
70        Self { tx, rx }
71    }
72
73    /// Build a controller with no OS-signal task, triggered only through
74    /// [`ShutdownController::trigger`]. Useful for tests and for embeddings that
75    /// drive shutdown from their own signal source.
76    #[must_use]
77    pub fn manual() -> Self {
78        let (tx, rx) = watch::channel(false);
79        Self { tx, rx }
80    }
81
82    /// A receiver for a cooperating background loop. The canonical stop check is
83    /// `while !*rx.borrow_and_update() { if rx.changed().await.is_err() { break } }`.
84    #[must_use]
85    pub fn subscribe(&self) -> watch::Receiver<bool> {
86        self.rx.clone()
87    }
88
89    /// Request a drain programmatically. This is the same path a signal takes,
90    /// so an admin endpoint or a fatal-error handler can ask for the same clean
91    /// teardown.
92    pub fn trigger(&self) {
93        // A closed channel means every receiver is gone; there is nothing to
94        // wake, so a send error is expected and ignored.
95        let _ = self.tx.send(true);
96    }
97
98    /// Whether a shutdown has already been requested.
99    #[must_use]
100    pub fn is_shutdown(&self) -> bool {
101        *self.rx.borrow()
102    }
103
104    /// The future to hand to `axum::serve(..).with_graceful_shutdown(..)`.
105    pub fn signalled(&self) -> impl Future<Output = ()> + Send + 'static {
106        wait_for_shutdown(self.rx.clone())
107    }
108}
109
110/// Resolve once `rx` observes a shutdown request. A dropped sender is treated as
111/// a request so a lost controller can never wedge a drain open.
112pub(crate) async fn wait_for_shutdown(mut rx: watch::Receiver<bool>) {
113    while !*rx.borrow_and_update() {
114        if rx.changed().await.is_err() {
115            break;
116        }
117    }
118}