diff --git a/noq/Cargo.toml b/noq/Cargo.toml index 4e50fa2e9..617078f78 100644 --- a/noq/Cargo.toml +++ b/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 } diff --git a/noq/src/connection.rs b/noq/src/connection.rs index 741d7e0e5..cb3aec858 100644 --- a/noq/src/connection.rs +++ b/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>); impl ConnectionRef { #[allow(clippy::redundant_allocation)] fn from_arc(inner: Arc>) -> 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, pub(crate) path_events: tokio::sync::broadcast::Sender, - /// Number of live handles that can be used to initiate or handle I/O; excludes the driver - ref_count: usize, sender: Pin>, pub(crate) runtime: Arc, send_buffer: Vec, @@ -1448,7 +1453,6 @@ impl State { stopped: FxHashMap::default(), open_path: FxHashMap::default(), error: None, - ref_count: 0, sender, runtime, send_buffer: Vec::new(), diff --git a/noq/src/endpoint.rs b/noq/src/endpoint.rs index 847ad3da4..820909e74 100644 --- a/noq/src/endpoint.rs +++ b/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, 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, 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(); } } } diff --git a/noq/src/recv_stream.rs b/noq/src/recv_stream.rs index 5467a9dc5..81fc858fc 100644 --- a/noq/src/recv_stream.rs +++ b/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()); } } diff --git a/noq/src/tests.rs b/noq/src/tests.rs index c6bb0f6cf..b1f8c5381 100755 --- a/noq/src/tests.rs +++ b/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 { + Box::::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) { + // 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) -> RawWaker { + // Store an Arc 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::::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::::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::::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::::from_raw(data as *const WakeCounter) }); +}