diff --git a/server/src/transport/ws.rs b/server/src/transport/ws.rs index 87241cd6a3..ac0b8e82c4 100644 --- a/server/src/transport/ws.rs +++ b/server/src/transport/ws.rs @@ -71,6 +71,9 @@ where + Sync + 'static, { + let cancel = tokio_util::sync::CancellationToken::new(); + let _g = cancel.drop_guard_ref(); + let BackgroundTaskParams { server_cfg, conn, @@ -150,6 +153,7 @@ where let rpc_service = rpc_service.clone(); let sink = sink.clone(); let extensions = extensions.clone(); + let cancel2 = cancel.clone(); tokio::spawn(async move { let first_non_whitespace = data.iter().enumerate().take(128).find(|(_, byte)| !byte.is_ascii_whitespace()); @@ -163,7 +167,19 @@ where } }; - let rp = handle_rpc_call(&data[idx..], is_single, batch_requests_config, &*rpc_service, extensions).await; + let Some(rp) = cancel2 + .run_until_cancelled(handle_rpc_call( + &data[idx..], + is_single, + batch_requests_config, + &*rpc_service, + extensions, + )) + .await + else { + // Cancelled + return; + }; // Subscriptions are handled by the subscription callback and // "ordinary notifications" should not be sent back to the client. diff --git a/tests/tests/helpers.rs b/tests/tests/helpers.rs index 70ac0c2dd8..d3b1df8c57 100644 --- a/tests/tests/helpers.rs +++ b/tests/tests/helpers.rs @@ -143,6 +143,10 @@ pub async fn server_with_subscription() -> SocketAddr { } pub async fn server() -> SocketAddr { + server_with_context((None,)).await +} + +pub async fn server_with_context(context: (Option>,)) -> SocketAddr { #[derive(Debug, Clone)] struct ConnectionDetails { inner: S, @@ -177,7 +181,7 @@ pub async fn server() -> SocketAddr { } } - let mut module = RpcModule::new(()); + let mut module = RpcModule::new(context); module.register_method("say_hello", |_, _, _| "hello").unwrap(); module.register_method("get_connection_id", |_, _, ext| *ext.get::().unwrap()).unwrap(); module @@ -196,6 +200,20 @@ pub async fn server() -> SocketAddr { }) .unwrap(); + module + .register_async_method("blocking_hello_with_cleanup_notif", |_, notif, _| async move { + struct CleanupGuard(Arc); + impl Drop for CleanupGuard { + fn drop(&mut self) { + self.0.notify_waiters(); + } + } + let _g = notif.0.clone().map(|n| CleanupGuard(n)); + tokio::time::sleep(std::time::Duration::MAX).await; + "hello" + }) + .unwrap(); + struct CustomError; impl From for ErrorObjectOwned { @@ -248,7 +266,18 @@ pub async fn server() -> SocketAddr { .connection_id(connection_id) .build(methods2.clone(), stop_hdl2.clone()); - async move { tower_service.call(req).await } + let session_closed = tower_service.on_session_closed(); + tokio::spawn(async move { + println!("session started, connection_id: {connection_id}"); + session_closed.await; + println!("session closed, connection_id: {connection_id}"); + }); + + async move { + // `session_closed` won't be scheduled for http without this, maybe a bug? + tokio::task::yield_now().await; + tower_service.call(req).await + } }); // Spawn a new task to serve each respective (Hyper) connection. diff --git a/tests/tests/integration_tests.rs b/tests/tests/integration_tests.rs index fde43082ec..a3789c3644 100644 --- a/tests/tests/integration_tests.rs +++ b/tests/tests/integration_tests.rs @@ -58,7 +58,7 @@ use tokio::time::interval; use tokio_stream::wrappers::IntervalStream; use tower_http::cors::CorsLayer; -use crate::helpers::server_with_sleeping_subscription; +use crate::helpers::{server_with_context, server_with_sleeping_subscription}; type HttpBody = http_body_util::Full; @@ -289,6 +289,44 @@ async fn http_method_call_works() { assert_eq!(&response, "hello"); } +#[tokio::test] +async fn http_method_call_cleanup_on_abort() { + init_logger(); + + let notif = Arc::new(tokio::sync::Notify::new()); + let notified = notif.notified(); + let server_addr = server_with_context((Some(notif.clone()),)).await; + let uri = format!("http://{}", server_addr); + let client = HttpClientBuilder::default().build(&uri).unwrap(); + tokio::time::timeout( + Duration::from_millis(100), + client.request::("blocking_hello_with_cleanup_notif", rpc_params![]), + ) + .await + .unwrap_err(); + drop(client); + notified.await +} + +#[tokio::test] +async fn ws_method_call_cleanup_on_abort() { + init_logger(); + + let notif = Arc::new(tokio::sync::Notify::new()); + let notified = notif.notified(); + let server_addr = server_with_context((Some(notif.clone()),)).await; + let uri = format!("ws://{}", server_addr); + let client = WsClientBuilder::default().build(&uri).await.unwrap(); + tokio::time::timeout( + Duration::from_millis(100), + client.request::("blocking_hello_with_cleanup_notif", rpc_params![]), + ) + .await + .unwrap_err(); + drop(client); + notified.await +} + #[tokio::test] async fn http_method_call_str_id_works() { init_logger();