atrg_core/
shutdown.rs

1//! Graceful shutdown utilities for the atrg server.
2//!
3//! Provides a signal-awaiting future suitable for
4//! [`axum::serve::Serve::with_graceful_shutdown`] and a post-shutdown cleanup
5//! helper that drains the database pool with a configurable timeout.
6
7use std::time::Duration;
8
9use tokio::signal;
10
11/// Default timeout (in seconds) for shutdown cleanup operations.
12const DEFAULT_SHUTDOWN_TIMEOUT_SECS: u64 = 30;
13
14/// Future that resolves when the process receives a shutdown signal.
15///
16/// On Unix this listens for both `SIGINT` (Ctrl+C) and `SIGTERM`.
17/// On other platforms only `SIGINT` is supported.
18///
19/// # Usage
20///
21/// ```rust,ignore
22/// axum::serve(listener, router)
23///     .with_graceful_shutdown(atrg_core::shutdown_signal())
24///     .await?;
25/// ```
26pub async fn shutdown_signal() {
27    let ctrl_c = async {
28        signal::ctrl_c()
29            .await
30            .expect("failed to install Ctrl+C handler");
31    };
32
33    #[cfg(unix)]
34    let terminate = async {
35        signal::unix::signal(signal::unix::SignalKind::terminate())
36            .expect("failed to install SIGTERM handler")
37            .recv()
38            .await;
39    };
40
41    #[cfg(not(unix))]
42    let terminate = std::future::pending::<()>();
43
44    tokio::select! {
45        _ = ctrl_c => {
46            tracing::info!("received Ctrl+C, initiating graceful shutdown");
47        }
48        _ = terminate => {
49            tracing::info!("received SIGTERM, initiating graceful shutdown");
50        }
51    }
52}
53
54/// Run cleanup tasks after the server stops accepting connections.
55///
56/// Currently this closes the SQLite connection pool, waiting up to `timeout`
57/// for in-flight queries to complete. If no timeout is provided the default
58/// of 30 seconds is used.
59pub async fn shutdown_cleanup(db: &sqlx::SqlitePool, timeout: Option<Duration>) {
60    let timeout = timeout.unwrap_or(Duration::from_secs(DEFAULT_SHUTDOWN_TIMEOUT_SECS));
61    tracing::info!(timeout_secs = timeout.as_secs(), "running shutdown cleanup");
62
63    let cleanup = async {
64        tracing::debug!("closing database connection pool");
65        db.close().await;
66        tracing::debug!("database pool closed");
67    };
68
69    match tokio::time::timeout(timeout, cleanup).await {
70        Ok(()) => tracing::info!("shutdown cleanup completed"),
71        Err(_) => tracing::warn!("shutdown cleanup timed out"),
72    }
73}
74
75// ---------------------------------------------------------------------------
76// Tests
77// ---------------------------------------------------------------------------
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn default_timeout_is_30_seconds() {
85        assert_eq!(DEFAULT_SHUTDOWN_TIMEOUT_SECS, 30);
86    }
87
88    #[tokio::test]
89    async fn shutdown_cleanup_completes_with_fresh_pool() {
90        let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
91        shutdown_cleanup(&pool, Some(Duration::from_secs(5))).await;
92        assert!(pool.is_closed());
93    }
94
95    #[tokio::test]
96    async fn shutdown_cleanup_handles_already_closed_pool() {
97        let pool = sqlx::SqlitePool::connect("sqlite::memory:").await.unwrap();
98        pool.close().await;
99        // Should not panic
100        shutdown_cleanup(&pool, Some(Duration::from_secs(1))).await;
101    }
102}
atrg_core/shutdown.rs

atrg_core/
shutdown.rs
