diff --git a/src/cmd/serve/mod.rs b/src/cmd/serve/mod.rs index 89f7ad5..35d9756 100644 --- a/src/cmd/serve/mod.rs +++ b/src/cmd/serve/mod.rs @@ -19,6 +19,7 @@ use crate::platform::proxy::build_client; use crate::platform::self_upgrade; use crate::platform::self_upgrade::TARGET_REPO; use crate::platform::self_upgrade::UpgradeOptions; +use crate::platform::shutdown::shutdown_signal; use crate::platform::transpile::scss::ScssTransformer; use crate::platform::transpile::typescript::TypeScriptTransformer; use crate::platform::watcher::Watcher; @@ -232,7 +233,7 @@ pub async fn main(args: ServeCommand) -> anyhow::Result<()> { )), }); - http1_server(&config.domain, { + http1_server(&config.domain, shutdown_signal(), { let state = state.clone(); move |req, res| { diff --git a/src/platform/http1/http1_server.rs b/src/platform/http1/http1_server.rs index e184d8e..eabddf5 100644 --- a/src/platform/http1/http1_server.rs +++ b/src/platform/http1/http1_server.rs @@ -1,4 +1,5 @@ use std::future::Future; +use std::net::SocketAddr; use std::sync::Arc; use http_body_util::BodyExt; @@ -12,12 +13,18 @@ use hyper::http::response::Builder as ResponseBuilder; use hyper::server::conn::http1; use hyper::service::service_fn; use hyper_util::rt::TokioIo; +use hyper_util::server::graceful::GracefulShutdown; use tokio::net::TcpListener; +use tokio::net::TcpStream; use tokio::net::ToSocketAddrs; -/// Simple wrapper around hyper to make it a little nicer to use +/// Simple wrapper around hyper to make it a little nicer to use. +/// +/// The returned future resolves once the listener stops and every in-flight +/// connection has drained. pub async fn http1_server( addr: A, + shutdown: impl Future, handle_func: F, ) -> anyhow::Result<()> where @@ -27,35 +34,51 @@ where { let listener = TcpListener::bind(&addr).await?; let handler_func_ref = Arc::new(handle_func); + let graceful = GracefulShutdown::new(); + tokio::pin!(shutdown); loop { - let Ok((stream, _)) = listener.accept().await else { - continue; - }; - let io = TokioIo::new(stream); - let handler_func_ref = handler_func_ref.clone(); + tokio::select! { + _ = &mut shutdown => break, + accepted = listener.accept() => handle_accepted(accepted, &handler_func_ref, &graceful), + } + } + + graceful.shutdown().await; + + Ok(()) +} + +fn handle_accepted( + accepted: std::io::Result<(TcpStream, SocketAddr)>, + handler_func_ref: &Arc, + graceful: &GracefulShutdown, +) where + F: 'static + Send + Sync + Fn(Request, ResponseBuilder) -> Fut, + Fut: Send + Future>>>, +{ + let Ok((stream, _)) = accepted else { + return; + }; + let io = TokioIo::new(stream); + let handler_func_ref = handler_func_ref.clone(); - tokio::task::spawn(async move { - let service_builder = http1::Builder::new(); - let service_handler = service_fn(move |req| { - let fut = handler_func_ref(req, Response::builder()); + let service_handler = service_fn(move |req| { + let fut = handler_func_ref(req, Response::builder()); - async move { - let handler_response = match fut.await { - Ok(handler_response) => handler_response, - Err(handler_error) => handle_error(handler_error), - }; + async move { + let handler_response = match fut.await { + Ok(handler_response) => handler_response, + Err(handler_error) => handle_error(handler_error), + }; - Ok::>, anyhow::Error>(handler_response) - } - }); + Ok::>, anyhow::Error>(handler_response) + } + }); - service_builder - .serve_connection(io, service_handler) - .await - .ok(); - }); - } + let connection = http1::Builder::new().serve_connection(io, service_handler); + + tokio::spawn(graceful.watch(connection)); } fn handle_error(error: anyhow::Error) -> Response> { diff --git a/src/platform/mod.rs b/src/platform/mod.rs index b4e11d4..77ce0c5 100644 --- a/src/platform/mod.rs +++ b/src/platform/mod.rs @@ -6,5 +6,6 @@ pub mod ip_address; pub mod logger; pub mod proxy; pub mod self_upgrade; +pub mod shutdown; pub mod transpile; pub mod watcher; diff --git a/src/platform/shutdown/mod.rs b/src/platform/shutdown/mod.rs new file mode 100644 index 0000000..e8d4682 --- /dev/null +++ b/src/platform/shutdown/mod.rs @@ -0,0 +1,3 @@ +mod shutdown_signal; + +pub use self::shutdown_signal::*; diff --git a/src/platform/shutdown/shutdown_signal.rs b/src/platform/shutdown/shutdown_signal.rs new file mode 100644 index 0000000..55e2454 --- /dev/null +++ b/src/platform/shutdown/shutdown_signal.rs @@ -0,0 +1,25 @@ +/// Resolves on Ctrl+C, or on SIGTERM where the platform provides it. +pub async fn shutdown_signal() { + let ctrl_c = async { + tokio::signal::ctrl_c().await.ok(); + }; + + #[cfg(unix)] + let terminate = async { + if let Ok(mut signal) = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + { + signal.recv().await; + } else { + std::future::pending::<()>().await; + } + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } +}