From 8ab6dd9803d9b74c306d2fce629e59779391fac6 Mon Sep 17 00:00:00 2001
From: Natty <natty.sh.git@gmail.com>
Date: Wed, 2 Aug 2023 03:18:16 +0200
Subject: [PATCH] Graceful shutdown on SIGTERM

---
 src/main.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/src/main.rs b/src/main.rs
index c7681fd..40dcc28 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,9 +15,11 @@ use magnetar_calckey_model::{CacheConnectorConfig, CalckeyCache, CalckeyModel, C
 use miette::{miette, IntoDiagnostic};
 use std::net::SocketAddr;
 use std::sync::Arc;
+use tokio::signal;
 use tower_http::cors::{Any, CorsLayer};
 use tower_http::trace::TraceLayer;
 use tracing::info;
+use tracing::log::error;
 use tracing_subscriber::EnvFilter;
 
 #[tokio::main]
@@ -85,6 +87,33 @@ async fn main() -> miette::Result<()> {
     info!("Serving on: {addr}");
     axum::Server::bind(&addr)
         .serve(app.into_make_service())
+        .with_graceful_shutdown(shutdown_signal())
         .await
         .map_err(|e| miette!("Error running server: {}", e))
 }
+
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        if let Err(e) = signal::ctrl_c().await {
+            error!("Ctrl+C signal handler error: {}", e);
+        }
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("SIGTERM handler error")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
+    }
+
+    info!("Shutting down...");
+}