From d4f0ba4029d48bc52d355f77f0784d1e691fa1c1 Mon Sep 17 00:00:00 2001 From: Li Yi Date: Thu, 2 Jun 2022 09:44:41 +0800 Subject: [PATCH] Handle SIGTERM properly Signed-off-by: Li Yi --- src/wagi_server.rs | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/wagi_server.rs b/src/wagi_server.rs index 9b8be0e..c919730 100644 --- a/src/wagi_server.rs +++ b/src/wagi_server.rs @@ -11,6 +11,7 @@ use hyper::{ use hyper::{Body, Response, Server}; use tokio::net::TcpStream; use tokio_rustls::server::TlsStream; +use tokio::signal::unix::SignalKind; pub struct WagiServer { routing_table: RoutingTable, @@ -18,6 +19,11 @@ pub struct WagiServer { address: SocketAddr, } +async fn shutdown_signal() { + // Wait for the SIGTERM signal + tokio::signal::unix::signal(SignalKind::terminate()).unwrap().recv().await; +} + impl WagiServer { pub async fn new(configuration: &WagiConfiguration, routing_table: RoutingTable) -> anyhow::Result { Ok(Self { @@ -65,9 +71,10 @@ impl WagiServer { })) }) }); - Server::builder(tls::TlsHyperAcceptor::new(&self.address, &tls.cert_path, &tls.key_path).await?) - .serve(mk_svc) - .await?; + let server = Server::builder(tls::TlsHyperAcceptor::new(&self.address, &tls.cert_path, &tls.key_path).await?) + .serve(mk_svc); + let graceful = server.with_graceful_shutdown(shutdown_signal()); + graceful.await?; }, None => { let mk_svc = make_service_fn(move |conn: &AddrStream| { @@ -80,7 +87,9 @@ impl WagiServer { })) } }); - Server::bind(&self.address).serve(mk_svc).await?; + let server = Server::bind(&self.address).serve(mk_svc); + let graceful = server.with_graceful_shutdown(shutdown_signal()); + graceful.await?; }, }