diff --git a/src/channels/webhook_server.rs b/src/channels/webhook_server.rs index 2425ab32..228abf0a 100644 --- a/src/channels/webhook_server.rs +++ b/src/channels/webhook_server.rs @@ -139,12 +139,19 @@ impl WebhookServer { self.config.addr } + /// Take ownership of shutdown primitives so callers can perform async + /// shutdown work without holding external locks around this server. + pub fn begin_shutdown(&mut self) -> (Option>, Option>) { + (self.shutdown_tx.take(), self.handle.take()) + } + /// Signal graceful shutdown and wait for the server task to finish. pub async fn shutdown(&mut self) { - if let Some(tx) = self.shutdown_tx.take() { + let (shutdown_tx, handle) = self.begin_shutdown(); + if let Some(tx) = shutdown_tx { let _ = tx.send(()); } - if let Some(handle) = self.handle.take() { + if let Some(handle) = handle { let _ = handle.await; } } @@ -269,6 +276,35 @@ mod tests { server.shutdown().await; } + #[tokio::test] + async fn test_begin_shutdown_takes_handles_for_lock_free_shutdown() { + let addr = SocketAddr::from((std::net::Ipv4Addr::LOCALHOST, 0)); + let mut server = WebhookServer::new(WebhookServerConfig { addr }); + + let test_router = axum::Router::new().route( + "/health", + axum::routing::get(|| async { Json(json!({"status": "ok"})) }), + ); + server.add_routes(test_router); + server.start().await.expect("Failed to start server"); // safety: test assertion for setup precondition + + let (shutdown_tx, handle) = server.begin_shutdown(); + assert!(shutdown_tx.is_some(), "shutdown sender should be available"); // safety: test assertion for expected server state + assert!(handle.is_some(), "server handle should be available"); // safety: test assertion for expected server state + + // begin_shutdown() should leave no handles behind on the server. + let (shutdown_tx2, handle2) = server.begin_shutdown(); + assert!(shutdown_tx2.is_none(), "shutdown sender should be consumed"); // safety: test assertion for postcondition + assert!(handle2.is_none(), "server handle should be consumed"); // safety: test assertion for postcondition + + if let Some(tx) = shutdown_tx { + let _ = tx.send(()); + } + if let Some(handle) = handle { + let _ = handle.await; + } + } + #[tokio::test] async fn test_restart_with_addr_rollback_on_bind_failure() { use std::net::TcpListener as StdTcpListener; diff --git a/src/main.rs b/src/main.rs index 12a8caf6..a7d95bec 100644 --- a/src/main.rs +++ b/src/main.rs @@ -920,7 +920,16 @@ async fn async_main() -> anyhow::Result<()> { } if let Some(ref ws_arc) = webhook_server { - ws_arc.lock().await.shutdown().await; + let (shutdown_tx, handle) = { + let mut ws = ws_arc.lock().await; + ws.begin_shutdown() + }; + if let Some(tx) = shutdown_tx { + let _ = tx.send(()); + } + if let Some(handle) = handle { + let _ = handle.await; + } } if let Some(tunnel) = active_tunnel {