Skip to main content

anvil_core/
shutdown.rs

1//! Graceful shutdown signal handling for `smith serve` and queue workers.
2
3use std::time::Duration;
4use tokio::signal;
5use tokio_util::sync::CancellationToken;
6
7/// Returns a future that completes when the process receives SIGINT or SIGTERM.
8pub async fn shutdown_signal() {
9    let ctrl_c = async {
10        signal::ctrl_c()
11            .await
12            .expect("failed to install Ctrl+C handler");
13    };
14
15    #[cfg(unix)]
16    let terminate = async {
17        signal::unix::signal(signal::unix::SignalKind::terminate())
18            .expect("failed to install SIGTERM handler")
19            .recv()
20            .await;
21    };
22
23    #[cfg(not(unix))]
24    let terminate = std::future::pending::<()>();
25
26    tokio::select! {
27        _ = ctrl_c => tracing::info!("received Ctrl+C, beginning graceful shutdown"),
28        _ = terminate => tracing::info!("received SIGTERM, beginning graceful shutdown"),
29    }
30}
31
32/// A cancellation handle that subsystems (queue workers, schedulers) can subscribe to.
33#[derive(Debug, Clone)]
34pub struct ShutdownHandle {
35    token: CancellationToken,
36}
37
38impl Default for ShutdownHandle {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44impl ShutdownHandle {
45    pub fn new() -> Self {
46        Self {
47            token: CancellationToken::new(),
48        }
49    }
50
51    pub fn token(&self) -> CancellationToken {
52        self.token.clone()
53    }
54
55    pub fn trigger(&self) {
56        self.token.cancel();
57    }
58
59    pub fn is_shutdown(&self) -> bool {
60        self.token.is_cancelled()
61    }
62
63    pub async fn wait(&self) {
64        self.token.cancelled().await
65    }
66
67    /// Spawn the OS-signal listener; trigger this handle on SIGINT/SIGTERM.
68    pub fn install(self) -> Self {
69        let trigger = self.clone();
70        tokio::spawn(async move {
71            shutdown_signal().await;
72            trigger.trigger();
73        });
74        self
75    }
76}
77
78pub const DEFAULT_DRAIN_TIMEOUT: Duration = Duration::from_secs(30);