Skip to main content

fuel_web_utils/
shutdown.rs

1use 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        // Initiate shutdown from clone
99        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}