quinn: move ConnectionRef/EndpointRef ref counts onto Shared as AtomicUsize

Avoids locking the State mutex on every Clone/Drop of a ConnectionRef
or EndpointRef: the count now lives on the lockless Shared struct as
an AtomicUsize bumped/decremented with Relaxed ordering.

Also folds in the two related upstream stream-cleanup changes and
ports all four regression tests so the port has coverage for the
drop/stop paths.

Ports (adapted to noq's double-Arc ConnectionRef and WeakConnectionHandle
machinery):
- quinn-rs/quinn#2495 @ 404db1bc9, 4b7a03949, 475b55bad
  (ref-count move; rightward-drift cleanup; RecvStream::drop early return
  + stream_drop_removes_blocked_reader test)
- quinn-rs/quinn#2609 @ 37625fe2d
  (fix: fetch_sub returns prior value, use >1 / ==1 semantics)
- quinn-rs/quinn#2541 @ 803c814
  (RecvStream::stop clears blocked_readers + recv_stream_cancel_stop_drop test)
- quinn-rs/quinn @ 07ce61cc2
  (dropped_endpoint_cleans_up / dropped_connection_cleans_up tests)

noq/Cargo.toml

@@ -100,7 +100,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 }

noq/src/connection.rs

@@ -6,7 +6,10 @@ use std::{
     net::{IpAddr, SocketAddr},
     num::NonZeroUsize,
     pin::Pin,
-    sync::{Arc, Weak},
+    sync::{
+        Arc, Weak,
+        atomic::{AtomicUsize, Ordering},
+    },
     task::{Context, Poll, Waker, ready},
 };
 
@@ -1237,7 +1240,7 @@ pub(crate) struct ConnectionRef(Arc<Arc<ConnectionInner>>);
 impl ConnectionRef {
     #[allow(clippy::redundant_allocation)]
     fn from_arc(inner: Arc<Arc<ConnectionInner>>) -> Self {
-        inner.lock_without_waking("from_arc").ref_count += 1;
+        inner.shared.ref_count.fetch_add(1, Ordering::Relaxed);
         Self(inner)
     }
 
@@ -1258,16 +1261,18 @@ impl Clone for ConnectionRef {
 
 impl Drop for ConnectionRef {
     fn drop(&mut self) {
+        if self.shared.ref_count.fetch_sub(1, Ordering::Relaxed) > 1 {
+            return;
+        }
+
         let conn = &mut *self.lock_without_waking("drop");
-        if let Some(x) = conn.ref_count.checked_sub(1) {
-            conn.ref_count = x;
-            if x == 0 && !conn.inner.is_closed() {
-                // If the driver is alive, it's just it and us, so we'd better shut it down. If it's
-                // not, we can't do any harm. If there were any streams being opened, then either
-                // the connection will be closed for an unrelated reason or a fresh reference will
-                // be constructed for the newly opened stream.
-                conn.implicit_close(&self.shared);
-            }
+
+        if !conn.inner.is_closed() {
+            // If the driver is alive, it's just it and us, so we'd better shut it down. If it's
+            // not, we can't do any harm. If there were any streams being opened, then either
+            // the connection will be closed for an unrelated reason or a fresh reference will
+            // be constructed for the newly opened stream.
+            conn.implicit_close(&self.shared);
         }
     }
 }
@@ -1372,6 +1377,8 @@ pub(crate) struct Shared {
     datagram_received: Notify,
     datagrams_unblocked: Notify,
     closed: Notify,
+    /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
+    ref_count: AtomicUsize,
 }
 
 pub(crate) struct State {
@@ -1405,8 +1412,6 @@ pub(crate) struct State {
     /// When the last reference to a path is dropped via [`Self::decrement_path_refs`] its value is cleared.
     pub(crate) final_path_stats: FxHashMap<PathId, PathStats>,
     pub(crate) path_events: tokio::sync::broadcast::Sender<PathEvent>,
-    /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
-    ref_count: usize,
     sender: Pin<Box<dyn UdpSender>>,
     pub(crate) runtime: Arc<dyn Runtime>,
     send_buffer: Vec<u8>,
@@ -1448,7 +1453,6 @@ impl State {
             stopped: FxHashMap::default(),
             open_path: FxHashMap::default(),
             error: None,
-            ref_count: 0,
             sender,
             runtime,
             send_buffer: Vec::new(),

noq/src/endpoint.rs

@@ -8,7 +8,10 @@ use std::{
     num::NonZeroUsize,
     pin::Pin,
     str,
-    sync::{Arc, Mutex},
+    sync::{
+        Arc, Mutex,
+        atomic::{AtomicUsize, Ordering},
+    },
     task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
 };
 
@@ -424,7 +427,8 @@ impl Future for EndpointDriver {
         // - all `Endpoint` structs are dropped and all connections are drained,
         // - or `Endpoint::close` has been called and all connections are drained.
         if endpoint.recv_state.connections.is_empty()
-            && (endpoint.ref_count == 0 || endpoint.recv_state.connections.close.is_some())
+            && (self.0.shared.ref_count.load(Ordering::Relaxed) == 0
+                || endpoint.recv_state.connections.close.is_some())
         {
             trace!("endpoint driver stopping");
             Poll::Ready(Ok(()))
@@ -524,8 +528,6 @@ pub(crate) struct State {
     driver: Option<Waker>,
     ipv6: bool,
     events: mpsc::UnboundedReceiver<(ConnectionHandle, EndpointEvent)>,
-    /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
-    ref_count: usize,
     driver_lost: bool,
     runtime: Arc<dyn Runtime>,
     stats: EndpointStats,
@@ -536,6 +538,8 @@ pub(crate) struct State {
 pub(crate) struct Shared {
     incoming: Notify,
     idle: Notify,
+    /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
+    ref_count: AtomicUsize,
 }
 
 impl State {
@@ -772,6 +776,7 @@ impl EndpointRef {
             shared: Shared {
                 incoming: Notify::new(),
                 idle: Notify::new(),
+                ref_count: AtomicUsize::new(0),
             },
             state: Mutex::new(State {
                 socket,
@@ -781,7 +786,6 @@ impl EndpointRef {
                 ipv6,
                 events,
                 driver: None,
-                ref_count: 0,
                 driver_lost: false,
                 recv_state,
                 runtime,
@@ -794,23 +798,22 @@ impl EndpointRef {
 
 impl Clone for EndpointRef {
     fn clone(&self) -> Self {
-        self.0.state.lock().unwrap().ref_count += 1;
+        self.0.shared.ref_count.fetch_add(1, Ordering::Relaxed);
         Self(self.0.clone())
     }
 }
 
 impl Drop for EndpointRef {
     fn drop(&mut self) {
+        if self.0.shared.ref_count.fetch_sub(1, Ordering::Relaxed) > 1 {
+            return;
+        }
+
         let endpoint = &mut *self.0.state.lock().unwrap();
-        if let Some(x) = endpoint.ref_count.checked_sub(1) {
-            endpoint.ref_count = x;
-            if x == 0 {
-                // If the driver is about to be on its own, ensure it can shut down if the last
-                // connection is gone.
-                if let Some(task) = endpoint.driver.take() {
-                    task.wake();
-                }
-            }
+        // If the driver is about to be on its own, ensure it can shut down if the last
+        // connection is gone.
+        if let Some(task) = endpoint.driver.take() {
+            task.wake();
         }
     }
 }

noq/src/recv_stream.rs

@@ -275,6 +275,9 @@ impl RecvStream {
         }
         conn.inner.recv_stream(self.stream).stop(error_code)?;
         self.all_data_read = true;
+        // Clean up shared state that might be left over from a cancelled read
+        // operation, so `drop` doesn't have to
+        conn.blocked_readers.remove(&self.stream);
         Ok(())
     }
 
@@ -581,6 +584,18 @@ impl tokio::io::AsyncRead for RecvStream {
 
 impl Drop for RecvStream {
     fn drop(&mut self) {
+        if self.all_data_read {
+            debug_assert!(
+                !self
+                    .conn
+                    .lock_without_waking("RecvStream:drop")
+                    .blocked_readers
+                    .contains_key(&self.stream),
+                "Stream {} should not have a blocked reader when all data read is true",
+                &self.stream
+            );
+            return;
+        }
         let mut conn = self.conn.lock_and_wake("RecvStream::drop");
 
         // clean up any previously registered wakers
@@ -590,10 +605,9 @@ impl Drop for RecvStream {
             conn.skip_waking();
             return;
         }
-        if !self.all_data_read {
-            // Ignore ClosedStream errors
-            let _ = conn.inner.recv_stream(self.stream).stop(0u32.into());
-        }
+
+        // Ignore ClosedStream errors
+        let _ = conn.inner.recv_stream(self.stream).stop(0u32.into());
     }
 }
 

noq/src/tests.rs

@@ -11,8 +11,13 @@ use std::{
     convert::TryInto,
     io,
     net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, UdpSocket},
+    pin::pin,
     str,
-    sync::Arc,
+    sync::{
+        Arc,
+        atomic::{AtomicUsize, Ordering},
+    },
+    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
 };
 
 use crate::runtime::TokioRuntime;
@@ -1213,6 +1218,43 @@ async fn weak_connection_handle() {
     client_res.expect("client task panicked");
 }
 
+#[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("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();
+    tokio::join!(
+        async {
+            endpoint
+                .connect(endpoint.local_addr().unwrap(), "localhost")
+                .unwrap()
+                .await
+                .unwrap()
+        },
+        async { endpoint.accept().await.unwrap().await.unwrap() }
+    );
+    endpoint.wait_idle().await;
+}
+
 /// Test that accessing stats from `Path` works as expected.
 #[tokio::test]
 async fn path_clone_stats_after_abandon() {
@@ -1481,3 +1523,151 @@ async fn nat_traversal_wakes_connection_driver() -> TestResult {
     tokio::join!(server_task, client_task);
     Ok(())
 }
+
+#[tokio::test]
+async fn stream_drop_removes_blocked_reader() {
+    let _guard = subscribe();
+
+    for drop_stream in [false, true] {
+        let endpoint_factory = EndpointFactory::new();
+        let server = endpoint_factory.endpoint("server");
+        let server_address = server.local_addr().unwrap();
+        let client = endpoint_factory.endpoint("client");
+
+        let server_task = tokio::spawn(async move {
+            let conn = server.accept().await.unwrap().await.unwrap();
+            let mut stream = conn.accept_uni().await.unwrap();
+
+            // read "hello"
+            let mut buf = [0u8; 5];
+            stream.read_exact(&mut buf).await.unwrap();
+
+            let (waker, wake_counter) = new_count_waker();
+            let mut cx = Context::from_waker(&waker);
+            // do a blocking read which will add the stream in conn.blocked_readers
+            {
+                let mut buf = [0u8; 64];
+                let read_fut = stream.read(&mut buf);
+                tokio::pin!(read_fut);
+                assert!(matches!(read_fut.as_mut().poll(&mut cx), Poll::Pending));
+            }
+
+            if !drop_stream {
+                assert_eq!(wake_counter.wakes(), 0);
+                // We have a blocked reader, closing the connection should wake it. We use this as
+                // a proxy to assert that the stream is in conn.blocked_readers.
+                conn.close(0u32.into(), b"done");
+                assert_eq!(wake_counter.wakes(), 1);
+            } else {
+                // dropping the stream should remove it from conn.blocked_readers, so we don't
+                // expect any wakeups
+                drop(stream);
+                assert_eq!(wake_counter.wakes(), 0, "no wakeups should have occurred");
+                conn.close(0u32.into(), b"done");
+                assert_eq!(wake_counter.wakes(), 0, "no wakeups should have occurred");
+            }
+        });
+
+        let conn = client
+            .connect(server_address, "localhost")
+            .unwrap()
+            .await
+            .unwrap();
+        let mut stream = conn.open_uni().await.unwrap();
+        // need to send some data to actually start the stream
+        stream.write_all(b"hello").await.unwrap();
+
+        server_task.await.unwrap();
+    }
+}
+
+/// Test that dropping a `RecvStream` after cancelling a read and then
+/// explicitly `stop`ing it doesn't panic.
+#[tokio::test]
+async fn recv_stream_cancel_stop_drop() {
+    let _guard = subscribe();
+    let factory = EndpointFactory::new();
+    let server = factory.endpoint("server");
+    let server_addr = server.local_addr().unwrap();
+    let client = factory.endpoint("client");
+    let recv_dropped = tokio::sync::SetOnce::new();
+    tokio::join!(
+        async {
+            let conn = server.accept().await.unwrap().await.unwrap();
+            let mut recv = conn.accept_uni().await.unwrap();
+            // Create a future to read from the stream, poll it once, then immediately drop it
+            {
+                let fut = pin!(recv.read_to_end(usize::MAX));
+                let mut cx = Context::from_waker(Waker::noop());
+                assert!(fut.poll(&mut cx).is_pending());
+            }
+            recv_dropped.set(()).unwrap();
+            recv.stop(0u32.into()).unwrap();
+        },
+        async {
+            let conn = client
+                .connect(server_addr, "localhost")
+                .unwrap()
+                .await
+                .unwrap();
+            let mut send = conn.open_uni().await.unwrap();
+            _ = send.write_all(b"hello").await;
+            // Don't drop (finish) the send stream until the read has been
+            // cancelled by the server, ensuring that read_to_end can't complete
+            // immediately.
+            recv_dropped.wait().await;
+        },
+    );
+}
+
+#[derive(Default)]
+struct WakeCounter {
+    wakes: AtomicUsize,
+}
+
+impl WakeCounter {
+    fn wakes(&self) -> usize {
+        self.wakes.load(Ordering::SeqCst)
+    }
+}
+
+fn new_count_waker() -> (Waker, Arc<WakeCounter>) {
+    // instance of WakeCounter
+    let counter = Arc::new(WakeCounter::default());
+
+    // convert
+    let waker = unsafe { Waker::from_raw(raw_waker(counter.clone())) };
+    (waker, counter)
+}
+
+fn raw_waker(counter: Arc<WakeCounter>) -> RawWaker {
+    // Store an Arc<WakeCounter> behind the raw pointer.
+    let ptr = Arc::into_raw(counter) as *const ();
+    RawWaker::new(ptr, &VTABLE)
+}
+
+static VTABLE: RawWakerVTable =
+    RawWakerVTable::new(clone_waker, wake_waker, wake_by_ref_waker, drop_waker);
+
+unsafe fn clone_waker(data: *const ()) -> RawWaker {
+    let arc = unsafe { Arc::<WakeCounter>::from_raw(data as *const WakeCounter) };
+    let cloned = arc.clone();
+    std::mem::forget(arc);
+    raw_waker(cloned)
+}
+
+unsafe fn wake_waker(data: *const ()) {
+    let arc = unsafe { Arc::<WakeCounter>::from_raw(data as *const WakeCounter) };
+    arc.wakes.fetch_add(1, Ordering::SeqCst);
+    // arc drops here
+}
+
+unsafe fn wake_by_ref_waker(data: *const ()) {
+    let arc = unsafe { Arc::<WakeCounter>::from_raw(data as *const WakeCounter) };
+    arc.wakes.fetch_add(1, Ordering::SeqCst);
+    std::mem::forget(arc);
+}
+
+unsafe fn drop_waker(data: *const ()) {
+    drop(unsafe { Arc::<WakeCounter>::from_raw(data as *const WakeCounter) });
+}