From 5dc189cf832917486ee452b092eb3c9174abf50f Mon Sep 17 00:00:00 2001
From: Marijan Smetko <msmetko@msmetko.xyz>
Date: Tue, 30 May 2023 06:51:16 +0200
Subject: [PATCH] Add graceful shutdown on SIGTERM (#1014)

* Add graceful shutdown on SIGTERM

* Fix linter
---
 atuin-server/src/lib.rs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/atuin-server/src/lib.rs b/atuin-server/src/lib.rs
index 2a77994..01873af 100644
--- a/atuin-server/src/lib.rs
+++ b/atuin-server/src/lib.rs
@@ -8,6 +8,8 @@ use eyre::{Context, Result};
 
 use crate::settings::Settings;
 
+use tokio::signal;
+
 pub mod auth;
 pub mod calendar;
 pub mod database;
@@ -17,6 +19,20 @@ pub mod router;
 pub mod settings;
 pub mod utils;
 
+async fn shutdown_signal() {
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to register signal handler")
+            .recv()
+            .await;
+    };
+
+    tokio::select! {
+        _ = terminate => (),
+    }
+    eprintln!("Shutting down gracefully...");
+}
+
 pub async fn launch(settings: Settings, host: String, port: u16) -> Result<()> {
     let host = host.parse::<IpAddr>()?;
 
@@ -28,6 +44,7 @@ pub async fn launch(settings: Settings, host: String, port: u16) -> Result<()> {
 
     Server::bind(&SocketAddr::new(host, port))
         .serve(r.into_make_service())
+        .with_graceful_shutdown(shutdown_signal())
         .await?;
 
     Ok(())