fix: gracefully shutdown on CTRL+C and SIGTERM

Fixes redlib-org#205
internationalcrisis · Oct 1, 2024 · 64c6b5a · 64c6b5a
1 parent e57eaa0
commit 64c6b5a
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/server.rs b/src/server.rs
@@ -26,6 +26,8 @@ use std::{
 };
 use time::Duration;
 
+use tokio::signal::unix::{signal, SignalKind};
+
 use crate::dbg_msg;
 
 type BoxResponse = Pin<Box<dyn Future<Output = Result<Response<Body>, String>> + Send>>;
@@ -267,8 +269,13 @@ impl Server {
 
 		// Bind server to address specified above. Gracefully shut down if CTRL+C is pressed
 		let server = HyperServer::bind(address).serve(make_svc).with_graceful_shutdown(async {
-			// Wait for the CTRL+C signal
-			tokio::signal::ctrl_c().await.expect("Failed to install CTRL+C signal handler");
+			// Wait for the CTRL+C or SIGTERM signals
+			let mut signal_terminate = signal(SignalKind::terminate()).expect("Failed to install SIGTERM signal handler");
+
+			tokio::select! {
+				_ = signal_terminate.recv() => (),
+				_ = tokio::signal::ctrl_c() => ()
+			}
 		});
 
 		server.boxed()