diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs index e40dd0c1f..7aad43b4c 100644 --- a/quinn/src/connection.rs +++ b/quinn/src/connection.rs @@ -53,7 +53,7 @@ impl Connecting { let (on_handshake_data_send, on_handshake_data_recv) = oneshot::channel(); let (on_connected_send, on_connected_recv) = oneshot::channel(); - let conn = ConnectionRef(Arc::new(ConnectionInner { + let conn = ConnectionRef(Arc::new(Arc::new(ConnectionInner { state: Mutex::new(State::new( conn, handle, @@ -65,7 +65,7 @@ impl Connecting { runtime.clone(), )), shared: Shared::default(), - })); + }))); let driver = ConnectionDriver(conn.clone()); runtime.spawn(Box::pin( @@ -1222,10 +1222,12 @@ impl Future for OnClosed { } #[derive(Debug)] -pub(crate) struct ConnectionRef(Arc); +#[allow(clippy::redundant_allocation)] +pub(crate) struct ConnectionRef(Arc>); impl ConnectionRef { - fn from_arc(inner: Arc) -> Self { + #[allow(clippy::redundant_allocation)] + fn from_arc(inner: Arc>) -> Self { inner.state.lock("from_arc").ref_count += 1; Self(inner) } @@ -1275,7 +1277,7 @@ pub(crate) struct ConnectionInner { /// This contains a weak reference to the connection so will not itself keep the connection /// alive. #[derive(Debug, Clone)] -pub struct WeakConnectionHandle(Weak); +pub struct WeakConnectionHandle(Weak>); impl WeakConnectionHandle { /// Returns `true` if the [`Connection`] associated with this handle is still alive. diff --git a/quinn/src/tests.rs b/quinn/src/tests.rs index f286cf3cd..0e8db1895 100755 --- a/quinn/src/tests.rs +++ b/quinn/src/tests.rs @@ -1118,3 +1118,37 @@ async fn on_closed_endpoint_drop() { .expect("client timeout") .expect("client task panicked"); } + +#[tokio::test] +async fn weak_connection_handle() { + let _guard = subscribe(); + let endpoint = endpoint(); + let endpoint2 = endpoint.clone(); + let server_task = tokio::spawn(async move { + let conn = endpoint2 + .accept() + .await + .expect("endpoint") + .await + .expect("connection"); + // create a weak handle to the connection + // ensure the underlying connection is not immediately dropped + let weak = conn.weak_handle(); + assert!(weak.is_alive()); + drop(conn); + // wait to ensure the connection is fully cleaned up + endpoint2.wait_idle().await; + assert!(!weak.is_alive()); + }); + let client_task = tokio::spawn(async move { + let conn = endpoint + .connect(endpoint.local_addr().unwrap(), "localhost") + .unwrap() + .await + .expect("connect"); + conn.on_closed().await; + }); + let (server_res, client_res) = tokio::join!(server_task, client_task); + server_res.expect("server task panicked"); + client_res.expect("client task panicked"); +}