hyperi_rustlib/
shutdown.rs

1// Project:   hyperi-rustlib
2// File:      src/shutdown.rs
3// Purpose:   Unified graceful shutdown with global CancellationToken
4// Language:  Rust
5//
6// License:   FSL-1.1-ALv2
7// Copyright: (c) 2026 HYPERI PTY LIMITED
8
9//! Unified graceful shutdown manager.
10//!
11//! Provides a global [`CancellationToken`] that all modules can listen on
12//! for coordinated graceful shutdown. One place handles SIGTERM/SIGINT,
13//! all modules drain gracefully.
14//!
15//! ## Usage
16//!
17//! ```rust,no_run
18//! use hyperi_rustlib::shutdown;
19//!
20//! #[tokio::main]
21//! async fn main() {
22//!     // Install the signal handler once at startup
23//!     let token = shutdown::install_signal_handler();
24//!
25//!     // Pass token to workers, or they can call shutdown::token() directly
26//!     tokio::spawn(async move {
27//!         loop {
28//!             tokio::select! {
29//!                 _ = token.cancelled() => {
30//!                     // drain and exit
31//!                     break;
32//!                 }
33//!                 _ = do_work() => {}
34//!             }
35//!         }
36//!     });
37//! }
38//!
39//! async fn do_work() {
40//!     tokio::time::sleep(std::time::Duration::from_secs(1)).await;
41//! }
42//! ```
43
44use std::sync::OnceLock;
45use tokio_util::sync::CancellationToken;
46
47static TOKEN: OnceLock<CancellationToken> = OnceLock::new();
48
49/// Get the global shutdown token.
50///
51/// All modules should clone this token and listen for cancellation
52/// in their main loops via `token.cancelled().await`.
53///
54/// The token is created lazily on first access.
55pub fn token() -> CancellationToken {
56    TOKEN.get_or_init(CancellationToken::new).clone()
57}
58
59/// Check if shutdown has been requested.
60pub fn is_shutdown() -> bool {
61    TOKEN.get().is_some_and(CancellationToken::is_cancelled)
62}
63
64/// Trigger shutdown programmatically.
65///
66/// Cancels the global token. All modules listening on it will
67/// begin their drain/cleanup sequence.
68pub fn trigger() {
69    if let Some(t) = TOKEN.get() {
70        t.cancel();
71    }
72}
73
74/// Wait for SIGTERM or SIGINT, then trigger shutdown.
75///
76/// Call this once at application startup. It spawns a background
77/// task that waits for the OS signal, then cancels the global token.
78///
79/// Returns the token for use in `tokio::select!` or other async
80/// shutdown coordination.
81#[must_use]
82pub fn install_signal_handler() -> CancellationToken {
83    let t = token();
84    let cancel = t.clone();
85
86    tokio::spawn(async move {
87        wait_for_signal().await;
88        cancel.cancel();
89
90        #[cfg(feature = "logger")]
91        tracing::info!("Shutdown signal received, cancelling all tasks");
92    });
93
94    t
95}
96
97/// Wait for SIGTERM or SIGINT.
98async fn wait_for_signal() {
99    let ctrl_c = async {
100        if let Err(e) = tokio::signal::ctrl_c().await {
101            tracing::error!(error = %e, "Failed to install Ctrl+C handler");
102            std::future::pending::<()>().await;
103        }
104    };
105
106    #[cfg(unix)]
107    let terminate = async {
108        match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) {
109            Ok(mut sig) => {
110                sig.recv().await;
111            }
112            Err(e) => {
113                tracing::error!(
114                    error = %e,
115                    "Failed to install SIGTERM handler, only Ctrl+C will trigger shutdown"
116                );
117                std::future::pending::<()>().await;
118            }
119        }
120    };
121
122    #[cfg(unix)]
123    tokio::select! {
124        () = ctrl_c => {},
125        () = terminate => {},
126    }
127
128    #[cfg(not(unix))]
129    ctrl_c.await;
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn token_is_not_cancelled_initially() {
138        // Use a fresh token (not the global) to avoid test pollution
139        let t = CancellationToken::new();
140        assert!(!t.is_cancelled());
141    }
142
143    #[test]
144    fn trigger_cancels_token() {
145        let t = CancellationToken::new();
146        assert!(!t.is_cancelled());
147        t.cancel();
148        assert!(t.is_cancelled());
149    }
150
151    #[test]
152    fn token_is_cloneable_and_shared() {
153        let t = CancellationToken::new();
154        let c1 = t.clone();
155        let c2 = t.clone();
156
157        assert!(!c1.is_cancelled());
158        assert!(!c2.is_cancelled());
159
160        t.cancel();
161
162        assert!(c1.is_cancelled());
163        assert!(c2.is_cancelled());
164    }
165
166    #[test]
167    fn multiple_triggers_are_idempotent() {
168        let t = CancellationToken::new();
169        t.cancel();
170        t.cancel(); // second cancel should not panic
171        assert!(t.is_cancelled());
172    }
173
174    #[tokio::test]
175    async fn cancelled_future_resolves_after_cancel() {
176        let t = CancellationToken::new();
177        let c = t.clone();
178
179        tokio::spawn(async move {
180            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
181            c.cancel();
182        });
183
184        // This should resolve once the token is cancelled
185        t.cancelled().await;
186        assert!(t.is_cancelled());
187    }
188
189    #[tokio::test]
190    async fn child_token_cancelled_by_parent() {
191        let parent = CancellationToken::new();
192        let child = parent.child_token();
193
194        assert!(!child.is_cancelled());
195        parent.cancel();
196        assert!(child.is_cancelled());
197    }
198}
hyperi_rustlib/shutdown.rs

hyperi_rustlib/
shutdown.rs
