Add graceful shutdown on SIGTERM (#1014)

* Add graceful shutdown on SIGTERM

* Fix linter

1 files changed, 17 insertions, 0 deletions

diff --git a/atuin-server/src/lib.rs b/atuin-server/src/lib.rs
index 2a77994c7..01873af96 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(())