From 8f63263ddec850cbe355e12e02a195a461479d91 Mon Sep 17 00:00:00 2001 From: Oleksandr Deundiak Date: Fri, 24 Jul 2026 12:47:27 -0400 Subject: [PATCH] fix(s2n-quic-transport): close all streams after Endpoint drop (#3170) Currently, after dropping the Endpoint the following is observed: - tx operations succeed - rx operations never return from an .await This is concerning because Endpoint drop can happen particularly due to IO implementation returning an error, which makes the current behaviour unexpected and hard to handle. The expected behaviour is: - tx operations return an error - rx operations return an error --- quic/s2n-quic-platform/src/io/testing.rs | 4 + quic/s2n-quic-tests/src/tests.rs | 1 + .../close_all_stream_after_endpoint_drop.rs | 199 ++++++++++++++++++ quic/s2n-quic-transport/src/connection/api.rs | 3 +- .../src/connection/api_provider.rs | 3 +- .../src/connection/connection_container.rs | 14 +- .../connection/connection_container/tests.rs | 53 ++++- .../src/connection/connection_impl.rs | 10 +- .../src/connection/connection_trait.rs | 3 +- 9 files changed, 275 insertions(+), 15 deletions(-) create mode 100644 quic/s2n-quic-tests/src/tests/close_all_stream_after_endpoint_drop.rs diff --git a/quic/s2n-quic-platform/src/io/testing.rs b/quic/s2n-quic-platform/src/io/testing.rs index b0895f3100..527903314b 100644 --- a/quic/s2n-quic-platform/src/io/testing.rs +++ b/quic/s2n-quic-platform/src/io/testing.rs @@ -222,6 +222,10 @@ impl Handle { queue_send_buffer_size: None, } } + + pub fn close_buffers(&self) { + self.buffers.close(); + } } pub struct Builder { diff --git a/quic/s2n-quic-tests/src/tests.rs b/quic/s2n-quic-tests/src/tests.rs index 3b28a549fd..291990a220 100644 --- a/quic/s2n-quic-tests/src/tests.rs +++ b/quic/s2n-quic-tests/src/tests.rs @@ -26,6 +26,7 @@ use std::{ mod blackhole; mod buffer_limit; +mod close_all_stream_after_endpoint_drop; mod connection_limits; mod connection_migration; mod deduplicate; diff --git a/quic/s2n-quic-tests/src/tests/close_all_stream_after_endpoint_drop.rs b/quic/s2n-quic-tests/src/tests/close_all_stream_after_endpoint_drop.rs new file mode 100644 index 0000000000..0ce2f6d446 --- /dev/null +++ b/quic/s2n-quic-tests/src/tests/close_all_stream_after_endpoint_drop.rs @@ -0,0 +1,199 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use bach::time::timeout; + +const CLOSE_PROPAGATION_DURATION: Duration = Duration::from_millis(5); + +/// Verifies client endpoint-drop tx stream behavior. +/// +/// After endpoint drop, tx stream send should complete with an error. +#[test] +fn client_tx_fails_after_endpoint_drop() { + let model = Model::default(); + test(model.clone(), |handle| { + let server = build_server(handle, model.clone())?; + let server_addr = start_server(server)?; + + let h = handle.clone(); + + let client = build_client(handle, model, true).unwrap(); + + primary::spawn(async move { + let connect = Connect::new(server_addr).with_server_name("localhost"); + let mut client_connection = client.connect(connect).await.unwrap(); + + let mut stream = client_connection.open_bidirectional_stream().await.unwrap(); + + let sent = Bytes::from("hello"); + stream.send(sent.clone()).await.unwrap(); + let received = stream.receive().await.unwrap().unwrap(); + assert_eq!(sent, received); + + // Drop endpoint + h.close_buffers(); + + delay(CLOSE_PROPAGATION_DURATION).await; + + let send_res = stream.send(Bytes::from("world")).await; + assert!(send_res.is_err()); + }); + + Ok(()) + }) + .unwrap(); +} + +/// Verifies client endpoint-drop rx stream behavior. +/// +/// After endpoint drop, rx stream accept should complete with an error instead of hanging. +#[test] +fn client_rx_fails_after_endpoint_drop() { + let model = Model::default(); + test(model.clone(), |handle| { + let server = build_server(handle, model.clone())?; + let server_addr = start_server(server)?; + + let h = handle.clone(); + + let client = build_client(handle, model, true).unwrap(); + + primary::spawn(async move { + let connect = Connect::new(server_addr).with_server_name("localhost"); + let mut client_connection = client.connect(connect).await.unwrap(); + + let mut stream = client_connection.open_bidirectional_stream().await.unwrap(); + + let sent = Bytes::from("hello"); + stream.send(sent.clone()).await.unwrap(); + let received = stream.receive().await.unwrap().unwrap(); + assert_eq!(sent, received); + + // Drop endpoint + h.close_buffers(); + + // The async call should return an error + let res = timeout( + CLOSE_PROPAGATION_DURATION, + async move { stream.receive().await }, + ) + .await + .unwrap(); + assert!(res.is_err()); + }); + + Ok(()) + }) + .unwrap(); +} + +/// Verifies server endpoint-drop tx stream behavior. +/// +/// After server endpoint drop, tx stream send should complete with an error. +#[test] +fn server_tx_fails_after_endpoint_drop() { + let model = Model::default(); + test(model.clone(), |handle| { + let mut server = build_server(handle, model.clone())?; + let server_addr = server.local_addr()?; + + let h = handle.clone(); + + let client = build_client(handle, model, true).unwrap(); + + run_echo_client(client, server_addr); + + primary::spawn(async move { + let mut connection = server.accept().await.unwrap(); + + let mut stream = connection + .accept_bidirectional_stream() + .await + .unwrap() + .unwrap(); + + let data = stream.receive().await.unwrap().unwrap(); + stream.send(Bytes::from("hello")).await.unwrap(); + + // Drop endpoint + h.close_buffers(); + + delay(CLOSE_PROPAGATION_DURATION).await; + + let send_res = stream.send(data).await; + assert!(send_res.is_err()); + }); + + Ok(()) + }) + .unwrap(); +} + +/// Verifies server endpoint-drop rx stream behavior. +/// +/// After server endpoint drop, rx stream receive should complete with an error instead of hanging. +#[test] +fn server_rx_fails_after_endpoint_drop() { + let model = Model::default(); + test(model.clone(), |handle| { + let mut server = build_server(handle, model.clone())?; + let server_addr = server.local_addr()?; + + let h = handle.clone(); + + let client = build_client(handle, model, true).unwrap(); + + run_echo_client(client, server_addr); + + primary::spawn(async move { + let mut connection = server.accept().await.unwrap(); + + let mut stream = connection + .accept_bidirectional_stream() + .await + .unwrap() + .unwrap(); + + let _data = stream.receive().await.unwrap().unwrap(); + stream.send(Bytes::from("hello")).await.unwrap(); + + // Drop endpoint + h.close_buffers(); + + // The async call should return an error + let res = timeout( + CLOSE_PROPAGATION_DURATION, + async move { stream.receive().await }, + ) + .await + .unwrap(); + assert!(res.is_err()); + }); + + Ok(()) + }) + .unwrap(); +} + +fn run_echo_client(client: Client, server_addr: SocketAddr) { + spawn(async move { + let connect = Connect::new(server_addr).with_server_name("localhost"); + let Ok(mut client_connection) = client.connect(connect).await else { + return; + }; + + let Ok(mut stream) = client_connection.open_bidirectional_stream().await else { + return; + }; + + let sent = Bytes::from("hello"); + let Ok(()) = stream.send(sent.clone()).await else { + return; + }; + _ = stream.receive().await; + + // Prevent dropping the connection and the stream + delay(Duration::from_secs(60)).await; + }); +} diff --git a/quic/s2n-quic-transport/src/connection/api.rs b/quic/s2n-quic-transport/src/connection/api.rs index 4e6c947f47..2cb1d12a54 100644 --- a/quic/s2n-quic-transport/src/connection/api.rs +++ b/quic/s2n-quic-transport/src/connection/api.rs @@ -163,7 +163,8 @@ impl Connection { /// This will immediately terminate all outstanding streams. #[inline] pub fn close(&self, error_code: application::Error) { - self.api.close_connection(Some(error_code)); + self.api + .close_connection(Some(connection::Error::application(error_code))); } #[inline] diff --git a/quic/s2n-quic-transport/src/connection/api_provider.rs b/quic/s2n-quic-transport/src/connection/api_provider.rs index b892805c0b..d383cca29f 100644 --- a/quic/s2n-quic-transport/src/connection/api_provider.rs +++ b/quic/s2n-quic-transport/src/connection/api_provider.rs @@ -16,7 +16,6 @@ use core::{ task::{Context, Poll}, }; use s2n_quic_core::{ - application, application::ServerName, inet::SocketAddress, query::{Query, QueryMut}, @@ -53,7 +52,7 @@ pub(crate) trait ConnectionApiProvider: Sync + Send { context: &Context, ) -> Poll>; - fn close_connection(&self, code: Option); + fn close_connection(&self, code: Option); fn server_name(&self) -> Result, connection::Error>; diff --git a/quic/s2n-quic-transport/src/connection/connection_container.rs b/quic/s2n-quic-transport/src/connection/connection_container.rs index ed92f073b5..7fda66403e 100644 --- a/quic/s2n-quic-transport/src/connection/connection_container.rs +++ b/quic/s2n-quic-transport/src/connection/connection_container.rs @@ -32,7 +32,6 @@ use intrusive_collections::{ intrusive_adapter, KeyAdapter, LinkedList, LinkedListLink, RBTree, RBTreeLink, }; use s2n_quic_core::{ - application, application::ServerName, event::supervisor, inet::SocketAddress, @@ -296,7 +295,7 @@ impl> ConnectionApiProvider for Con } } - fn close_connection(&self, error: Option) { + fn close_connection(&self, error: Option) { let _: Result<(), connection::Error> = self.api_write_call(|conn| { conn.application_close(error); Ok(()) @@ -1146,6 +1145,17 @@ impl> ConnectionContainer { } } +impl> Drop for ConnectionContainer { + // ConnectionContainer dropped means Endpoint is dropped. Close all connections. + fn drop(&mut self) { + let mut cursor = self.connection_map.front(); + while let Some(node) = cursor.get() { + node.close_connection(Some(connection::Error::endpoint_closing())); + cursor.move_next(); + } + } +} + /// Return values for iterations over a `Connection` list. /// The value instructs the iterator whether iteration will be continued. #[derive(Clone, Copy, Debug)] diff --git a/quic/s2n-quic-transport/src/connection/connection_container/tests.rs b/quic/s2n-quic-transport/src/connection/connection_container/tests.rs index ab6b42d8aa..811bb37df2 100644 --- a/quic/s2n-quic-transport/src/connection/connection_container/tests.rs +++ b/quic/s2n-quic-transport/src/connection/connection_container/tests.rs @@ -14,11 +14,12 @@ use bolero::{check, generator::*}; use bytes::Bytes; use core::{ any::Any, + cell::Cell, task::{Context, Poll}, time::Duration, }; use s2n_quic_core::{ - application, event, + event, event::builder::DatagramDropReason, inet::{DatagramInfo, SocketAddress}, io::tx, @@ -36,6 +37,12 @@ use s2n_quic_core::{ }; use std::sync::Mutex; +thread_local! { + static CLOSE_CALL_COUNT: Cell = const { Cell::new(0) }; + static APPLICATION_CLOSE_CALL_COUNT: Cell = const { Cell::new(0) }; + static APPLICATION_CLOSE_WITH_ERROR_COUNT: Cell = const { Cell::new(0) }; +} + struct TestConnection { accept_state: AcceptState, is_closed: bool, @@ -90,6 +97,7 @@ impl connection::Trait for TestConnection { ) { assert!(!self.is_closed); assert!(!self.close_timer.is_armed()); + CLOSE_CALL_COUNT.with(|count| count.set(count.get() + 1)); self.close_timer.set(timestamp + Duration::from_secs(1)); } @@ -293,7 +301,11 @@ impl connection::Trait for TestConnection { _stream_type: Option, _context: &Context, ) -> Poll, connection::Error>> { - todo!() + if self.is_closed { + return Poll::Ready(Err(connection::Error::unspecified())); + } + + Poll::Pending } fn poll_open_stream( @@ -302,11 +314,20 @@ impl connection::Trait for TestConnection { _token: &mut connection::OpenToken, _context: &Context, ) -> Poll> { - todo!() + if self.is_closed { + return Poll::Ready(Err(connection::Error::unspecified())); + } + + Poll::Ready(Ok(stream::StreamId::from_varint( + s2n_quic_core::varint::VarInt::from_u8(0), + ))) } - fn application_close(&mut self, _error: Option) { - // no-op + fn application_close(&mut self, _error: Option) { + APPLICATION_CLOSE_CALL_COUNT.with(|count| count.set(count.get() + 1)); + if _error.is_some() { + APPLICATION_CLOSE_WITH_ERROR_COUNT.with(|count| count.set(count.get() + 1)); + } } fn server_name(&self) -> Option { @@ -627,3 +648,25 @@ fn container_test() { assert!(connections.next().is_none()); }); } + +#[test] +fn drop_closes_all_connections() { + APPLICATION_CLOSE_CALL_COUNT.with(|count| count.set(0)); + APPLICATION_CLOSE_WITH_ERROR_COUNT.with(|count| count.set(0)); + + let mut id_gen = InternalConnectionIdGenerator::new(); + let (_handle, acceptor, connector, _close_handle) = endpoint::handle::Handle::new(100); + + { + let mut container: ConnectionContainer = + ConnectionContainer::new(acceptor, connector); + + for _ in 0..3 { + let id = id_gen.generate_id(); + container.insert_connection(TestConnection::default(), id); + } + } + + APPLICATION_CLOSE_CALL_COUNT.with(|count| assert_eq!(count.get(), 3)); + APPLICATION_CLOSE_WITH_ERROR_COUNT.with(|count| assert_eq!(count.get(), 3)); +} diff --git a/quic/s2n-quic-transport/src/connection/connection_impl.rs b/quic/s2n-quic-transport/src/connection/connection_impl.rs index 4c6385331d..65791d767f 100644 --- a/quic/s2n-quic-transport/src/connection/connection_impl.rs +++ b/quic/s2n-quic-transport/src/connection/connection_impl.rs @@ -33,7 +33,7 @@ use core::{ }; use s2n_codec::DecoderBufferMut; use s2n_quic_core::{ - application::{self, ServerName}, + application::ServerName, connection::{ error::Error, id::{Classification, Generator as _}, @@ -2281,7 +2281,7 @@ impl connection::Trait for ConnectionImpl { ) } - fn application_close(&mut self, error: Option) { + fn application_close(&mut self, error: Option) { if self.error.is_err() { return; } @@ -2290,7 +2290,11 @@ impl connection::Trait for ConnectionImpl { self.open_registry = None; if let Some(error) = error { - self.error = Err(connection::Error::application(error)); + self.error = Err(error); + // This will put all streams into Reset state and wake all tasks + if let Some((space, _)) = self.space_manager.application_mut() { + space.stream_manager.close(error); + } } else { // give the connection some time to flush all outstanding streams self.state = ConnectionState::Flushing; diff --git a/quic/s2n-quic-transport/src/connection/connection_trait.rs b/quic/s2n-quic-transport/src/connection/connection_trait.rs index 481cf48315..54f3de2eeb 100644 --- a/quic/s2n-quic-transport/src/connection/connection_trait.rs +++ b/quic/s2n-quic-transport/src/connection/connection_trait.rs @@ -21,7 +21,6 @@ use core::{ }; use s2n_codec::DecoderBufferMut; use s2n_quic_core::{ - application, application::ServerName, event::{self, builder::DatagramDropReason, supervisor, ConnectionPublisher, IntoEvent}, inet::{DatagramInfo, SocketAddress}, @@ -518,7 +517,7 @@ pub trait ConnectionTrait: 'static + Send + Sized { context: &Context, ) -> Poll>; - fn application_close(&mut self, error: Option); + fn application_close(&mut self, error: Option); fn server_name(&self) -> Option;