feat(serve): graceful shutdown on Ctrl+C and SIGTERM

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| {

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<F, Fut, A>(
   addr: A,
+  shutdown: impl Future<Output = ()>,
   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<F, Fut>(
+  accepted: std::io::Result<(TcpStream, SocketAddr)>,
+  handler_func_ref: &Arc<F>,
+  graceful: &GracefulShutdown,
+) where
+  F: 'static + Send + Sync + Fn(Request<Incoming>, ResponseBuilder) -> Fut,
+  Fut: Send + Future<Output = anyhow::Result<Response<BoxBody<HyperBytes, std::io::Error>>>>,
+{
+  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::<Response<BoxBody<HyperBytes, std::io::Error>>, anyhow::Error>(handler_response)
-        }
-      });
+      Ok::<Response<BoxBody<HyperBytes, std::io::Error>>, 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<BoxBody<HyperBytes, std::io::Error>> {

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;

src/platform/shutdown/mod.rs

@@ -0,0 +1,3 @@
+mod shutdown_signal;
+
+pub use self::shutdown_signal::*;

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 => {},
+  }
+}