1 files changed, 32 insertions, 2 deletions
diff --git a/crates/auth/src/main.rs b/crates/auth/src/main.rs
index 4a71d69..ef8a358 100644
--- a/crates/auth/src/main.rs
+++ b/crates/auth/src/main.rs
@@ -8,6 +8,7 @@ use std::net::{Ipv6Addr, SocketAddr};
 
 use clap::Parser;
 use stack_up::{Configuration, Services, tracing::Tracing};
+use tokio::{signal, task::AbortHandle};
 use tracing::{info, trace};
 
 use crate::{error::AppError, state::AppState};
@@ -56,13 +57,42 @@ async fn main() -> Result<(), AppError> {
         .run(&services.postgres)
         .await?;
 
-    let state = AppState::create(services, &config).await?;
+    let (state, deletion_task) = AppState::create(services, &config).await?;
 
     let addr = SocketAddr::from((Ipv6Addr::UNSPECIFIED, config.application.port));
 
     let listener = tokio::net::TcpListener::bind(addr).await?;
     info!(port = addr.port(), "serving api");
 
-    axum::serve(listener, server::router(state)).await?;
+    axum::serve(listener, server::router(state))
+        .with_graceful_shutdown(shutdown_signal(deletion_task.abort_handle()))
+        .await?;
+
+    deletion_task.await??;
+
     Ok(())
 }
+
+async fn shutdown_signal(deletion_task_abort_handle: AbortHandle) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { deletion_task_abort_handle.abort() },
+        _ = terminate => { deletion_task_abort_handle.abort() },
+    }
+}