[dep-upgrade-202402] fix shutdown issue introduced when upgrading hyperdep-upgrade-202402

6 files changed, 71 insertions, 60 deletions

diff --git a/src/api/admin/api_server.rs b/src/api/admin/api_server.rs
index d5e1c777..50813d11 100644
--- a/src/api/admin/api_server.rs
+++ b/src/api/admin/api_server.rs
@@ -3,9 +3,9 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 
-use futures::future::Future;
 use http::header::{ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN, ALLOW};
 use hyper::{body::Incoming as IncomingBody, Request, Response, StatusCode};
+use tokio::sync::watch;
 
 use opentelemetry::trace::SpanRef;
 
@@ -65,11 +65,11 @@ impl AdminApiServer {
 	pub async fn run(
 		self,
 		bind_addr: UnixOrTCPSocketAddress,
-		shutdown_signal: impl Future<Output = ()>,
+		must_exit: watch::Receiver<bool>,
 	) -> Result<(), GarageError> {
 		let region = self.garage.config.s3_api.s3_region.clone();
 		ApiServer::new(region, self)
-			.run_server(bind_addr, Some(0o220), shutdown_signal)
+			.run_server(bind_addr, Some(0o220), must_exit)
 			.await
 	}
 
diff --git a/src/api/generic_server.rs b/src/api/generic_server.rs
index 7b37417e..9c49fdf3 100644
--- a/src/api/generic_server.rs
+++ b/src/api/generic_server.rs
@@ -18,6 +18,7 @@ use hyper_util::rt::TokioIo;
 
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::net::{TcpListener, TcpStream, UnixListener, UnixStream};
+use tokio::sync::watch;
 
 use opentelemetry::{
 	global,
@@ -104,20 +105,17 @@ impl<A: ApiHandler> ApiServer<A> {
 		self: Arc<Self>,
 		bind_addr: UnixOrTCPSocketAddress,
 		unix_bind_addr_mode: Option<u32>,
-		shutdown_signal: impl Future<Output = ()>,
+		must_exit: watch::Receiver<bool>,
 	) -> Result<(), GarageError> {
-		info!(
-			"{} API server listening on {}",
-			A::API_NAME_DISPLAY,
-			bind_addr
-		);
+		let server_name = format!("{} API", A::API_NAME_DISPLAY);
+		info!("{} server listening on {}", server_name, bind_addr);
 
 		match bind_addr {
 			UnixOrTCPSocketAddress::TCPSocket(addr) => {
 				let listener = TcpListener::bind(addr).await?;
 
 				let handler = move |request, socketaddr| self.clone().handler(request, socketaddr);
-				server_loop(listener, handler, shutdown_signal).await
+				server_loop(server_name, listener, handler, must_exit).await
 			}
 			UnixOrTCPSocketAddress::UnixSocket(ref path) => {
 				if path.exists() {
@@ -133,7 +131,7 @@ impl<A: ApiHandler> ApiServer<A> {
 				)?;
 
 				let handler = move |request, socketaddr| self.clone().handler(request, socketaddr);
-				server_loop(listener, handler, shutdown_signal).await
+				server_loop(server_name, listener, handler, must_exit).await
 			}
 		}
 	}
@@ -278,9 +276,10 @@ impl Accept for UnixListenerOn {
 }
 
 pub async fn server_loop<A, H, F, E>(
+	server_name: String,
 	listener: A,
 	handler: H,
-	shutdown_signal: impl Future<Output = ()>,
+	mut must_exit: watch::Receiver<bool>,
 ) -> Result<(), GarageError>
 where
 	A: Accept,
@@ -288,42 +287,57 @@ where
 	F: Future<Output = Result<Response<BoxBody<E>>, http::Error>> + Send + 'static,
 	E: Send + Sync + std::error::Error + 'static,
 {
-	tokio::pin!(shutdown_signal);
-
 	let (conn_in, mut conn_out) = tokio::sync::mpsc::unbounded_channel();
-	let connection_collector = tokio::spawn(async move {
-		let mut collection = FuturesUnordered::new();
-		loop {
-			let collect_next = async {
-				if collection.is_empty() {
-					futures::future::pending().await
-				} else {
-					collection.next().await
-				}
-			};
-			tokio::select! {
-				result = collect_next => {
-					trace!("HTTP connection finished: {:?}", result);
-				}
-				new_fut = conn_out.recv() => {
-					match new_fut {
-						Some(f) => collection.push(f),
-						None => break,
+	let connection_collector = tokio::spawn({
+		let server_name = server_name.clone();
+		async move {
+			let mut connections = FuturesUnordered::new();
+			loop {
+				let collect_next = async {
+					if connections.is_empty() {
+						futures::future::pending().await
+					} else {
+						connections.next().await
+					}
+				};
+				tokio::select! {
+					result = collect_next => {
+						trace!("{} server: HTTP connection finished: {:?}", server_name, result);
+					}
+					new_fut = conn_out.recv() => {
+						match new_fut {
+							Some(f) => connections.push(f),
+							None => break,
+						}
 					}
 				}
 			}
+			if !connections.is_empty() {
+				info!(
+					"{} server: {} connections still open",
+					server_name,
+					connections.len()
+				);
+				while let Some(conn_res) = connections.next().await {
+					trace!(
+						"{} server: HTTP connection finished: {:?}",
+						server_name,
+						conn_res
+					);
+					info!(
+						"{} server: {} connections still open",
+						server_name,
+						connections.len()
+					);
+				}
+			}
 		}
-		debug!("Collecting last open HTTP connections.");
-		while let Some(conn_res) = collection.next().await {
-			trace!("HTTP connection finished: {:?}", conn_res);
-		}
-		debug!("No more HTTP connections to collect");
 	});
 
-	loop {
+	while !*must_exit.borrow() {
 		let (stream, client_addr) = tokio::select! {
 			acc = listener.accept() => acc?,
-			_ = &mut shutdown_signal => break,
+			_ = must_exit.changed() => continue,
 		};
 
 		let io = TokioIo::new(stream);
@@ -343,6 +357,8 @@ where
 		conn_in.send(fut)?;
 	}
 
+	info!("{} server exiting", server_name);
+	drop(conn_in);
 	connection_collector.await?;
 
 	Ok(())
diff --git a/src/api/k2v/api_server.rs b/src/api/k2v/api_server.rs
index 128742c4..e97da2af 100644
--- a/src/api/k2v/api_server.rs
+++ b/src/api/k2v/api_server.rs
@@ -2,8 +2,8 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 
-use futures::future::Future;
 use hyper::{body::Incoming as IncomingBody, Method, Request, Response};
+use tokio::sync::watch;
 
 use opentelemetry::{trace::SpanRef, KeyValue};
 
@@ -42,10 +42,10 @@ impl K2VApiServer {
 		garage: Arc<Garage>,
 		bind_addr: UnixOrTCPSocketAddress,
 		s3_region: String,
-		shutdown_signal: impl Future<Output = ()>,
+		must_exit: watch::Receiver<bool>,
 	) -> Result<(), GarageError> {
 		ApiServer::new(s3_region, K2VApiServer { garage })
-			.run_server(bind_addr, None, shutdown_signal)
+			.run_server(bind_addr, None, must_exit)
 			.await
 	}
 }
diff --git a/src/api/s3/api_server.rs b/src/api/s3/api_server.rs
index 495c5832..4b815f79 100644
--- a/src/api/s3/api_server.rs
+++ b/src/api/s3/api_server.rs
@@ -2,9 +2,9 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 
-use futures::future::Future;
 use hyper::header;
 use hyper::{body::Incoming as IncomingBody, Request, Response};
+use tokio::sync::watch;
 
 use opentelemetry::{trace::SpanRef, KeyValue};
 
@@ -51,10 +51,10 @@ impl S3ApiServer {
 		garage: Arc<Garage>,
 		addr: UnixOrTCPSocketAddress,
 		s3_region: String,
-		shutdown_signal: impl Future<Output = ()>,
+		must_exit: watch::Receiver<bool>,
 	) -> Result<(), GarageError> {
 		ApiServer::new(s3_region, S3ApiServer { garage })
-			.run_server(addr, None, shutdown_signal)
+			.run_server(addr, None, must_exit)
 			.await
 	}
 
diff --git a/src/garage/server.rs b/src/garage/server.rs
index de8ac9e2..51b06b8e 100644
--- a/src/garage/server.rs
+++ b/src/garage/server.rs
@@ -88,7 +88,7 @@ pub async fn run_server(config_file: PathBuf, secrets: Secrets) -> Result<(), Er
 				garage.clone(),
 				s3_bind_addr.clone(),
 				config.s3_api.s3_region.clone(),
-				wait_from(watch_cancel.clone()),
+				watch_cancel.clone(),
 			)),
 		));
 	}
@@ -103,7 +103,7 @@ pub async fn run_server(config_file: PathBuf, secrets: Secrets) -> Result<(), Er
 					garage.clone(),
 					config.k2v_api.as_ref().unwrap().api_bind_addr.clone(),
 					config.s3_api.s3_region.clone(),
-					wait_from(watch_cancel.clone()),
+					watch_cancel.clone(),
 				)),
 			));
 		}
@@ -116,10 +116,7 @@ pub async fn run_server(config_file: PathBuf, secrets: Secrets) -> Result<(), Er
 		let web_server = WebServer::new(garage.clone(), web_config.root_domain.clone());
 		servers.push((
 			"Web",
-			tokio::spawn(web_server.run(
-				web_config.bind_addr.clone(),
-				wait_from(watch_cancel.clone()),
-			)),
+			tokio::spawn(web_server.run(web_config.bind_addr.clone(), watch_cancel.clone())),
 		));
 	}
 
@@ -127,9 +124,7 @@ pub async fn run_server(config_file: PathBuf, secrets: Secrets) -> Result<(), Er
 		info!("Launching Admin API server...");
 		servers.push((
 			"Admin",
-			tokio::spawn(
-				admin_server.run(admin_bind_addr.clone(), wait_from(watch_cancel.clone())),
-			),
+			tokio::spawn(admin_server.run(admin_bind_addr.clone(), watch_cancel.clone())),
 		));
 	}
 
diff --git a/src/web/web_server.rs b/src/web/web_server.rs
index 766e3829..269f37f2 100644
--- a/src/web/web_server.rs
+++ b/src/web/web_server.rs
@@ -2,7 +2,8 @@ use std::fs::{self, Permissions};
 use std::os::unix::prelude::PermissionsExt;
 use std::{convert::Infallible, sync::Arc};
 
-use futures::future::Future;
+use tokio::net::{TcpListener, UnixListener};
+use tokio::sync::watch;
 
 use hyper::{
 	body::Incoming as IncomingBody,
@@ -10,8 +11,6 @@ use hyper::{
 	Method, Request, Response, StatusCode,
 };
 
-use tokio::net::{TcpListener, UnixListener};
-
 use opentelemetry::{
 	global,
 	metrics::{Counter, ValueRecorder},
@@ -84,8 +83,9 @@ impl WebServer {
 	pub async fn run(
 		self: Arc<Self>,
 		bind_addr: UnixOrTCPSocketAddress,
-		shutdown_signal: impl Future<Output = ()>,
+		must_exit: watch::Receiver<bool>,
 	) -> Result<(), GarageError> {
+		let server_name = "Web".into();
 		info!("Web server listening on {}", bind_addr);
 
 		match bind_addr {
@@ -94,7 +94,7 @@ impl WebServer {
 
 				let handler =
 					move |stream, socketaddr| self.clone().handle_request(stream, socketaddr);
-				server_loop(listener, handler, shutdown_signal).await
+				server_loop(server_name, listener, handler, must_exit).await
 			}
 			UnixOrTCPSocketAddress::UnixSocket(ref path) => {
 				if path.exists() {
@@ -108,7 +108,7 @@ impl WebServer {
 
 				let handler =
 					move |stream, socketaddr| self.clone().handle_request(stream, socketaddr);
-				server_loop(listener, handler, shutdown_signal).await
+				server_loop(server_name, listener, handler, must_exit).await
 			}
 		}
 	}