fuel_web_utils/
shutdown.rs1use std::{
2 sync::Arc,
3 time::Duration,
4};
5use tokio_util::sync::CancellationToken;
6
7pub const GRACEFUL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(90);
8
9#[derive(Clone)]
10pub struct ShutdownController {
11 token: CancellationToken,
12}
13
14impl Default for ShutdownController {
15 fn default() -> Self {
16 Self::new()
17 }
18}
19
20impl ShutdownController {
21 pub fn new() -> Self {
22 Self {
23 token: CancellationToken::new(),
24 }
25 }
26
27 pub fn token(&self) -> &CancellationToken {
28 &self.token
29 }
30
31 pub fn spawn_signal_handler(self: Arc<Self>) -> Arc<Self> {
32 tokio::spawn({
33 let shutdown = self.clone();
34 async move {
35 tokio::signal::ctrl_c()
36 .await
37 .expect("Failed to listen for ctrl+c");
38 tracing::info!("Received shutdown signal");
39 shutdown.initiate_shutdown();
40 }
41 });
42 self
43 }
44
45 pub fn initiate_shutdown(&self) {
46 tracing::info!("Initiating graceful shutdown...");
47 self.token.cancel();
48 }
49
50 pub fn is_shutdown_initiated(&self) -> bool {
51 self.token.is_cancelled()
52 }
53
54 pub async fn wait_for_shutdown(&self) {
55 self.token.cancelled().await;
56 }
57}
58
59#[cfg(test)]
60mod tests {
61 use std::time::Duration;
62
63 use super::*;
64
65 #[tokio::test]
66 async fn test_manual_shutdown() {
67 let controller = ShutdownController::new();
68 assert!(
69 !controller.is_shutdown_initiated(),
70 "Controller should not be shutdown initially"
71 );
72
73 controller.initiate_shutdown();
74 assert!(
75 controller.is_shutdown_initiated(),
76 "Controller should be shutdown after initiation"
77 );
78 }
79
80 #[tokio::test]
81 async fn test_wait_for_shutdown_timeout() {
82 let controller = ShutdownController::new();
83
84 let timeout = Duration::from_millis(50);
85 let result = tokio::time::timeout(timeout, controller.wait_for_shutdown()).await;
86
87 assert!(
88 result.is_err(),
89 "wait_for_shutdown should not complete without initiation"
90 );
91 }
92
93 #[tokio::test]
94 async fn test_clone_behavior() {
95 let controller = ShutdownController::new();
96 let cloned = controller.clone();
97
98 cloned.initiate_shutdown();
100
101 assert!(
102 controller.is_shutdown_initiated(),
103 "Original should be shutdown"
104 );
105 assert!(cloned.is_shutdown_initiated(), "Clone should be shutdown");
106 }
107}