Skip to main content

kerniq_runtime/
shutdown.rs

1//! Shutdown coordination.
2
3use std::future::Future;
4
5use tokio_util::sync::CancellationToken;
6
7/// A cloneable handle used to request and observe server shutdown.
8///
9/// Every clone refers to the same underlying signal. The server stops
10/// accepting connections once it is triggered, background tasks receive it,
11/// and long-lived handlers (WebSockets, streaming responses) should watch it
12/// to finish cleanly:
13///
14/// ```
15/// use kerniq_runtime::Shutdown;
16///
17/// # async fn demo() {
18/// let shutdown = Shutdown::new();
19/// let watcher = shutdown.clone();
20/// let task = tokio::spawn(async move {
21///     watcher.wait().await;
22///     "stopped"
23/// });
24/// shutdown.trigger();
25/// assert_eq!(task.await.unwrap(), "stopped");
26/// # }
27/// ```
28#[derive(Debug, Clone, Default)]
29pub struct Shutdown {
30    token: CancellationToken,
31}
32
33impl Shutdown {
34    /// Creates a new, untriggered handle.
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Requests shutdown. Idempotent.
40    pub fn trigger(&self) {
41        self.token.cancel();
42    }
43
44    /// Whether shutdown has been requested.
45    pub fn is_triggered(&self) -> bool {
46        self.token.is_cancelled()
47    }
48
49    /// Completes once shutdown has been requested. Cancel-safe.
50    pub async fn wait(&self) {
51        self.token.cancelled().await;
52    }
53
54    /// An owned future that completes on shutdown (useful for `select!` in
55    /// `'static` contexts).
56    pub fn wait_owned(&self) -> impl Future<Output = ()> + Send + 'static {
57        self.token.clone().cancelled_owned()
58    }
59
60    /// The underlying Tokio cancellation token, for interoperability.
61    pub fn token(&self) -> &CancellationToken {
62        &self.token
63    }
64}
65
66impl From<CancellationToken> for Shutdown {
67    fn from(token: CancellationToken) -> Self {
68        Self { token }
69    }
70}
71
72/// Completes when the process receives Ctrl+C (SIGINT) or, on Unix, SIGTERM.
73///
74/// If a signal handler cannot be installed the failure is logged and that
75/// signal is ignored (the future then waits on the remaining signals).
76pub async fn os_signal() {
77    let ctrl_c = async {
78        if let Err(err) = tokio::signal::ctrl_c().await {
79            tracing::warn!(error = %err, "failed to listen for Ctrl+C");
80            std::future::pending::<()>().await;
81        }
82    };
83
84    #[cfg(unix)]
85    let terminate = async {
86        use tokio::signal::unix::{SignalKind, signal};
87        match signal(SignalKind::terminate()) {
88            Ok(mut stream) => {
89                stream.recv().await;
90            }
91            Err(err) => {
92                tracing::warn!(error = %err, "failed to listen for SIGTERM");
93                std::future::pending::<()>().await;
94            }
95        }
96    };
97    #[cfg(not(unix))]
98    let terminate = std::future::pending::<()>();
99
100    tokio::select! {
101        () = ctrl_c => tracing::info!("received Ctrl+C, shutting down"),
102        () = terminate => tracing::info!("received SIGTERM, shutting down"),
103    }
104}