fix: graceful shutdown

4 files changed, 77 insertions, 4 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");
+}
diff --git a/crates/rule-executor/src/main.rs b/crates/rule-executor/src/main.rs
index 18c9222..ed284c6 100644
--- a/crates/rule-executor/src/main.rs
+++ b/crates/rule-executor/src/main.rs
@@ -1,4 +1,6 @@
+#[allow(dead_code)]
 mod cnfg;
+
 mod processor;
 mod state;
 
diff --git a/crates/warden/src/main.rs b/crates/warden/src/main.rs
index 90fedfd..ecc9cbc 100644
--- a/crates/warden/src/main.rs
+++ b/crates/warden/src/main.rs
@@ -5,10 +5,14 @@ mod state;
 mod version;
 
 use std::net::{Ipv6Addr, SocketAddr};
+use tokio::signal;
 
 use clap::{Parser, command};
 use tracing::{error, info, trace};
-use warden_stack::{Configuration, Services, tracing::Tracing};
+use warden_stack::{
+    Configuration, Services,
+    tracing::{SdkTracerProvider, Tracing},
+};
 
 use crate::state::AppState;
 
@@ -42,6 +46,8 @@ async fn main() -> Result<(), error::AppError> {
         .loki(&config.application, &config.monitoring)?
         .build(&config.monitoring);
 
+    let provider = tracing.otel_provider;
+
     tokio::spawn(tracing.loki_task);
 
     let mut services = Services::builder()
@@ -91,7 +97,37 @@ async fn main() -> Result<(), error::AppError> {
     info!(port = addr.port(), "starting warden");
 
     let router = server::router(state).merge(server::metrics_app());
-    axum::serve(listener, router).await?;
+    axum::serve(listener, router)
+        .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");
+}
diff --git a/lib/warden-core/src/lib.rs b/lib/warden-core/src/lib.rs
index c97bef3..80fb34d 100644
--- a/lib/warden-core/src/lib.rs
+++ b/lib/warden-core/src/lib.rs
@@ -22,6 +22,7 @@ pub mod iso20022;
 
 /// Message in transit
 #[allow(missing_docs)]
+#[allow(clippy::large_enum_variant)]
 #[cfg(feature = "message")]
 pub mod message;