fix: an in-progress RPC call is not aborted when the websocket connection is closed

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.

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<Arc<tokio::sync::Notify>>,)) -> SocketAddr {
 	#[derive(Debug, Clone)]
 	struct ConnectionDetails<S> {
 		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::<u32>().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<tokio::sync::Notify>);
+			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<CustomError> 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.

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<hyper::body::Bytes>;
 
@@ -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::<String, _>("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::<String, _>("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();