From 717e7cf5a417cbae317241219f04871a2fe14fc2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=A3o=20Ant=C3=B4nio=20Cardoso?= <joao.maker@gmail.com>
Date: Thu, 3 Oct 2024 15:05:25 -0300
Subject: [PATCH] src: web: Add graceful shutdown

---
 src/web/mod.rs | 40 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/web/mod.rs b/src/web/mod.rs
index e4e2c72e..fd449db0 100644
--- a/src/web/mod.rs
+++ b/src/web/mod.rs
@@ -18,7 +18,10 @@ use axum::{
 };
 use futures::{sink::SinkExt, stream::StreamExt};
 use serde::Deserialize;
-use tokio::sync::{broadcast, mpsc, RwLock};
+use tokio::{
+    signal,
+    sync::{broadcast, mpsc, RwLock},
+};
 use tracing::*;
 use uuid::Uuid;
 
@@ -227,9 +230,16 @@ pub async fn run(address: String) {
                 continue;
             }
         };
-        if let Err(error) = axum::serve(listener, router.clone()).into_future().await {
-            error!("WebServer error: {}", error);
+
+        if let Err(error) = axum::serve(listener, router.clone())
+            .with_graceful_shutdown(shutdown_signal())
+            .await
+        {
+            error!("WebServer error: {error}");
+            continue;
         }
+
+        break;
     }
 }
 
@@ -240,3 +250,27 @@ where
     let mut router = SERVER.router.lock().unwrap();
     modifier(&mut router);
 }
+
+async fn shutdown_signal() {
+    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 => {},
+    }
+}