make shutdown grace periods configurable

Signed-off-by: Jason Volk <jason@zemos.net>

conduwuit-example.toml

@@ -389,6 +389,14 @@
 #
 #client_response_timeout = 120
 
+# Grace period for clean shutdown of client requests (seconds).
+#
+#client_shutdown_timeout = 10
+
+# Grace period for clean shutdown of federation requests (seconds).
+#
+#sender_shutdown_timeout = 5
+
 # Enables registration. If set to false, no users can register on this
 # server.
 #

src/core/config/mod.rs

@@ -498,6 +498,18 @@ pub struct Config {
 	#[serde(default = "default_client_response_timeout")]
 	pub client_response_timeout: u64,
 
+	/// Grace period for clean shutdown of client requests (seconds).
+	///
+	/// default: 10
+	#[serde(default = "default_client_shutdown_timeout")]
+	pub client_shutdown_timeout: u64,
+
+	/// Grace period for clean shutdown of federation requests (seconds).
+	///
+	/// default: 5
+	#[serde(default = "default_sender_shutdown_timeout")]
+	pub sender_shutdown_timeout: u64,
+
 	/// Enables registration. If set to false, no users can register on this
 	/// server.
 	///
@@ -2194,3 +2206,7 @@ fn default_client_receive_timeout() -> u64 { 75 }
 fn default_client_request_timeout() -> u64 { 180 }
 
 fn default_client_response_timeout() -> u64 { 120 }
+
+fn default_client_shutdown_timeout() -> u64 { 15 }
+
+fn default_sender_shutdown_timeout() -> u64 { 5 }

src/router/request.rs

@@ -1,6 +1,7 @@
 use std::{
 	fmt::Debug,
 	sync::{atomic::Ordering, Arc},
+	time::Duration,
 };
 
 use axum::{
@@ -9,7 +10,9 @@ use axum::{
 };
 use conduwuit::{debug, debug_error, debug_warn, err, error, trace, Result};
 use conduwuit_service::Services;
+use futures::FutureExt;
 use http::{Method, StatusCode, Uri};
+use tokio::time::sleep;
 use tracing::Span;
 
 #[tracing::instrument(
@@ -63,8 +66,14 @@ pub(crate) async fn handle(
 	let task = services.server.runtime().spawn(async move {
 		tokio::select! {
 			response = execute(&services_, req, next, parent) => response,
-			() = services_.server.until_shutdown() =>
+			response = services_.server.until_shutdown()
-				StatusCode::SERVICE_UNAVAILABLE.into_response(),
+				.then(|()| {
+					let timeout = services_.server.config.client_shutdown_timeout;
+					let timeout = Duration::from_secs(timeout);
+					sleep(timeout)
+				})
+				.map(|()| StatusCode::SERVICE_UNAVAILABLE)
+				.map(IntoResponse::into_response) => response,
 		}
 	});
 

src/router/run.rs

@@ -122,7 +122,8 @@ async fn handle_shutdown(server: Arc<Server>, tx: Sender<()>, handle: axum_serve
 		error!("failed sending shutdown transaction to channel: {e}");
 	}
 
-	let timeout = Duration::from_secs(36);
+	let timeout = server.config.client_shutdown_timeout;
+	let timeout = Duration::from_secs(timeout);
 	debug!(
 		?timeout,
 		handle_active = ?server.metrics.requests_handle_active.load(Ordering::Relaxed),

src/service/sending/sender.rs

@@ -67,8 +67,6 @@ type SendingFuture<'a> = BoxFuture<'a, SendingResult>;
 type SendingFutures<'a> = FuturesUnordered<SendingFuture<'a>>;
 type CurTransactionStatus = HashMap<Destination, TransactionStatus>;
 
-const CLEANUP_TIMEOUT_MS: u64 = 3500;
-
 const SELECT_PRESENCE_LIMIT: usize = 256;
 const SELECT_RECEIPT_LIMIT: usize = 256;
 const SELECT_EDU_LIMIT: usize = EDU_LIMIT - 2;
@@ -216,8 +214,9 @@ impl Service {
 			time::{sleep_until, Instant},
 		};
 
+		let timeout = self.server.config.sender_shutdown_timeout;
+		let timeout = Duration::from_secs(timeout);
 		let now = Instant::now();
-		let timeout = Duration::from_millis(CLEANUP_TIMEOUT_MS);
 		let deadline = now.checked_add(timeout).unwrap_or(now);
 		loop {
 			trace!("Waiting for {} requests to complete...", futures.len());