Test cleanup of connections/endpoints when all references are dropped

Author: Ralith

quinn/Cargo.toml

@@ -79,7 +79,7 @@ directories-next = { workspace = true }
 rand = { workspace = true }
 rcgen = { workspace = true }
 clap = { workspace = true }
-tokio = { workspace = true, features = ["rt", "rt-multi-thread", "time", "macros"] }
+tokio = { workspace = true, features = ["rt", "rt-multi-thread", "time", "macros", "test-util"] }
 tracing-subscriber = { workspace = true }
 tracing-futures = { workspace = true }
 url = { workspace = true }

quinn/src/tests.rs

@@ -1057,6 +1057,43 @@ async fn recv_stream_cancel_stop_drop() {
     );
 }
 
+#[tokio::test(start_paused = true)]
+async fn dropped_endpoint_cleans_up() {
+    let _guard = subscribe();
+
+    let mut endpoint_factory = EndpointFactory::new();
+    let cid_generator = Arc::new(|| -> Box<dyn proto::ConnectionIdGenerator> {
+        Box::<proto::HashedConnectionIdGenerator>::default()
+    });
+    endpoint_factory
+        .endpoint_config
+        .cid_generator(cid_generator.clone());
+    let endpoint = endpoint_factory.endpoint();
+    drop(endpoint_factory);
+    assert_eq!(Arc::strong_count(&cid_generator), 2);
+    drop(endpoint);
+    // Let the driver task run; paused runtimes are guaranteed to drain pending work on sleep.
+    tokio::time::sleep(Duration::from_millis(1)).await;
+    assert_eq!(Arc::strong_count(&cid_generator), 1);
+}
+
+#[tokio::test]
+async fn dropped_connection_cleans_up() {
+    let _guard = subscribe();
+    let endpoint = endpoint();
+    join!(
+        async {
+            endpoint
+                .connect(endpoint.local_addr().unwrap(), "localhost")
+                .unwrap()
+                .await
+                .unwrap()
+        },
+        async { endpoint.accept().await.unwrap().await.unwrap() }
+    );
+    endpoint.wait_idle().await;
+}
+
 #[derive(Default)]
 struct WakeCounter {
     wakes: AtomicUsize,