1 files changed, 36 insertions, 2 deletions
diff --git a/crates/configuration/src/main.rs b/crates/configuration/src/main.rs
index 6f90b69..a7843a1 100644
--- a/crates/configuration/src/main.rs
+++ b/crates/configuration/src/main.rs
@@ -7,9 +7,13 @@ use std::net::{Ipv6Addr, SocketAddr};
 use crate::{server::error::AppError, state::AppState};
 use axum::http::header::CONTENT_TYPE;
 use clap::Parser;
+use tokio::signal;
 use tower::{make::Shared, steer::Steer};
 use tracing::{error, info, trace};
-use warden_stack::{Configuration, Services, tracing::Tracing};
+use warden_stack::{
+    Configuration, Services,
+    tracing::{SdkTracerProvider, Tracing},
+};
 
 /// warden-config
 #[derive(Parser, Debug)]
@@ -114,7 +118,37 @@ async fn main() -> Result<(), AppError> {
     let listener = tokio::net::TcpListener::bind(addr).await?;
     info!(port = addr.port(), "starting config-api");
 
-    axum::serve(listener, Shared::new(service)).await?;
+    axum::serve(listener, Shared::new(service))
+        .with_graceful_shutdown(shutdown_signal(provider))
+        .await?;
 
     Ok(())
 }
+
+async fn shutdown_signal(provider: SdkTracerProvider) {
+    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 => {},
+        _ = terminate => {},
+    }
+
+    provider
+        .shutdown()
+        .expect("failed to shutdown trace provider");
+}