feat: Make the future returned from SendStream::stopped 'static

Frando · gretchenfrage · commit f1fe1832a7ba · 2025-05-14T04:43:32.000Z
Changes the implementation of `SendStream::stopped` such that the
returned future is static and no longer lifetime-bound onto a mutable
reference to the send stream.

This allows to use the stopped future with combinators or in a separate
task while still sending on the stream concurrently.

Internally, this is done changing the implementation of the stopped
notification to use a cloneable tokio::sync::Notify instead of storing
a single waker.
diff --git a/quinn/src/connection.rs b/quinn/src/connection.rs
@@ -963,7 +963,7 @@ pub(crate) struct State {
     endpoint_events: mpsc::UnboundedSender<(ConnectionHandle, EndpointEvent)>,
     pub(crate) blocked_writers: FxHashMap<StreamId, Waker>,
     pub(crate) blocked_readers: FxHashMap<StreamId, Waker>,
-    pub(crate) stopped: FxHashMap<StreamId, Waker>,
+    pub(crate) stopped: FxHashMap<StreamId, Arc<Notify>>,
     /// Always set to Some before the connection becomes drained
     pub(crate) error: Option<ConnectionError>,
     /// Number of live handles that can be used to initiate or handle I/O; excludes the driver
@@ -1105,7 +1105,7 @@ impl State {
                         // `ZeroRttRejected` errors.
                         wake_all(&mut self.blocked_writers);
                         wake_all(&mut self.blocked_readers);
-                        wake_all(&mut self.stopped);
+                        wake_all_notify(&mut self.stopped);
                     }
                 }
                 ConnectionLost { reason } => {
@@ -1129,9 +1129,9 @@ impl State {
                     // Might mean any number of streams are ready, so we wake up everyone
                     shared.stream_budget_available[dir as usize].notify_waiters();
                 }
-                Stream(StreamEvent::Finished { id }) => wake_stream(id, &mut self.stopped),
+                Stream(StreamEvent::Finished { id }) => wake_stream_notify(id, &mut self.stopped),
                 Stream(StreamEvent::Stopped { id, .. }) => {
-                    wake_stream(id, &mut self.stopped);
+                    wake_stream_notify(id, &mut self.stopped);
                     wake_stream(id, &mut self.blocked_writers);
                 }
             }
@@ -1212,7 +1212,7 @@ impl State {
         if let Some(x) = self.on_connected.take() {
             let _ = x.send(false);
         }
-        wake_all(&mut self.stopped);
+        wake_all_notify(&mut self.stopped);
         shared.closed.notify_waiters();
     }
 
@@ -1266,6 +1266,18 @@ fn wake_all(wakers: &mut FxHashMap<StreamId, Waker>) {
     wakers.drain().for_each(|(_, waker)| waker.wake())
 }
 
+fn wake_stream_notify(stream_id: StreamId, wakers: &mut FxHashMap<StreamId, Arc<Notify>>) {
+    if let Some(notify) = wakers.remove(&stream_id) {
+        notify.notify_waiters()
+    }
+}
+
+fn wake_all_notify(wakers: &mut FxHashMap<StreamId, Arc<Notify>>) {
+    wakers
+        .drain()
+        .for_each(|(_, notify)| notify.notify_waiters())
+}
+
 /// Errors that can arise when sending a datagram
 #[derive(Debug, Error, Clone, Eq, PartialEq)]
 pub enum SendDatagramError {
diff --git a/quinn/src/send_stream.rs b/quinn/src/send_stream.rs
@@ -9,7 +9,10 @@ use bytes::Bytes;
 use proto::{ClosedStream, ConnectionError, FinishError, StreamId, Written};
 use thiserror::Error;
 
-use crate::{VarInt, connection::ConnectionRef};
+use crate::{
+    VarInt,
+    connection::{ConnectionRef, State},
+};
 
 /// A stream that can only be used to send data
 ///
@@ -199,27 +202,31 @@ impl SendStream {
     /// For a variety of reasons, the peer may not send acknowledgements immediately upon receiving
     /// data. As such, relying on `stopped` to know when the peer has read a stream to completion
     /// may introduce more latency than using an application-level response of some sort.
-    pub async fn stopped(&mut self) -> Result<Option<VarInt>, StoppedError> {
-        Stopped { stream: self }.await
-    }
-
-    fn poll_stopped(&mut self, cx: &mut Context) -> Poll<Result<Option<VarInt>, StoppedError>> {
-        let mut conn = self.conn.state.lock("SendStream::poll_stopped");
-
-        if self.is_0rtt {
-            conn.check_0rtt()
-                .map_err(|()| StoppedError::ZeroRttRejected)?;
-        }
-
-        match conn.inner.send_stream(self.stream).stopped() {
-            Err(_) => Poll::Ready(Ok(None)),
-            Ok(Some(error_code)) => Poll::Ready(Ok(Some(error_code))),
-            Ok(None) => {
-                if let Some(e) = &conn.error {
-                    return Poll::Ready(Err(e.clone().into()));
+    pub fn stopped(
+        &self,
+    ) -> impl Future<Output = Result<Option<VarInt>, StoppedError>> + Send + Sync + 'static {
+        let conn = self.conn.clone();
+        let stream = self.stream;
+        let is_0rtt = self.is_0rtt;
+        async move {
+            loop {
+                // The `Notify::notified` future needs to be created while the lock is being held,
+                // otherwise a wakeup could be missed if triggered inbetween releasing the lock
+                // and creating the future.
+                // The lock may only be held in a block without `await`s, otherwise the future
+                // becomes `!Send`. `Notify::notified` is lifetime-bound to `Notify`, therefore
+                // we need to declare `notify` outside of the block, and initialize it inside.
+                let notify;
+                {
+                    let mut conn = conn.state.lock("SendStream::stopped");
+                    if let Some(output) = send_stream_stopped(&mut conn, stream, is_0rtt) {
+                        return output;
+                    }
+
+                    notify = conn.stopped.entry(stream).or_default().clone();
+                    notify.notified()
                 }
-                conn.stopped.insert(self.stream, cx.waker().clone());
-                Poll::Pending
+                .await
             }
         }
     }
@@ -245,6 +252,25 @@ impl SendStream {
     }
 }
 
+/// Check if a send stream is stopped.
+///
+/// Returns `Some` if the stream is stopped or the connection is closed.
+/// Returns `None` if the stream is not stopped.
+fn send_stream_stopped(
+    conn: &mut State,
+    stream: StreamId,
+    is_0rtt: bool,
+) -> Option<Result<Option<VarInt>, StoppedError>> {
+    if is_0rtt && conn.check_0rtt().is_err() {
+        return Some(Err(StoppedError::ZeroRttRejected));
+    }
+    match conn.inner.send_stream(stream).stopped() {
+        Err(ClosedStream { .. }) => Some(Ok(None)),
+        Ok(Some(error_code)) => Some(Ok(Some(error_code))),
+        Ok(None) => conn.error.clone().map(|error| Err(error.into())),
+    }
+}
+
 #[cfg(feature = "futures-io")]
 impl futures_io::AsyncWrite for SendStream {
     fn poll_write(self: Pin<&mut Self>, cx: &mut Context, buf: &[u8]) -> Poll<io::Result<usize>> {
@@ -283,7 +309,6 @@ impl Drop for SendStream {
         let mut conn = self.conn.state.lock("SendStream::drop");
 
         // clean up any previously registered wakers
-        conn.stopped.remove(&self.stream);
         conn.blocked_writers.remove(&self.stream);
 
         if conn.error.is_some() || (self.is_0rtt && conn.check_0rtt().is_err()) {
@@ -302,19 +327,6 @@ impl Drop for SendStream {
     }
 }
 
-/// Future produced by `SendStream::stopped`
-struct Stopped<'a> {
-    stream: &'a mut SendStream,
-}
-
-impl Future for Stopped<'_> {
-    type Output = Result<Option<VarInt>, StoppedError>;
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
-        self.get_mut().stream.poll_stopped(cx)
-    }
-}
-
 /// Errors that arise from writing to a stream
 #[derive(Debug, Error, Clone, PartialEq, Eq)]
 pub enum WriteError {
diff --git a/quinn/src/tests.rs b/quinn/src/tests.rs
@@ -861,3 +861,87 @@ async fn multiple_conns_with_zero_length_cids() {
     .instrument(error_span!("server"));
     tokio::join!(client1, client2, server);
 }
+
+#[tokio::test]
+async fn stream_stopped() {
+    let _guard = subscribe();
+    let factory = EndpointFactory::new();
+    let server = {
+        let _guard = error_span!("server").entered();
+        factory.endpoint()
+    };
+    let server_addr = server.local_addr().unwrap();
+
+    let client = {
+        let _guard = error_span!("client1").entered();
+        factory.endpoint()
+    };
+
+    let client = async move {
+        let conn = client
+            .connect(server_addr, "localhost")
+            .unwrap()
+            .await
+            .unwrap();
+        let mut stream = conn.open_uni().await.unwrap();
+        let stopped1 = stream.stopped();
+        let stopped2 = stream.stopped();
+        let stopped3 = stream.stopped();
+
+        stream.write_all(b"hi").await.unwrap();
+        // spawn one of the futures into a task
+        let stopped1 = tokio::task::spawn(stopped1);
+        // verify that both futures resolved
+        let (stopped1, stopped2) = tokio::join!(stopped1, stopped2);
+        assert!(matches!(stopped1, Ok(Ok(Some(val))) if val == 42u32.into()));
+        assert!(matches!(stopped2, Ok(Some(val)) if val == 42u32.into()));
+        // drop the stream
+        drop(stream);
+        // verify that a future also resolves after dropping the stream
+        let stopped3 = stopped3.await;
+        assert_eq!(stopped3, Ok(Some(42u32.into())));
+    };
+    let client =
+        tokio::time::timeout(Duration::from_millis(100), client).instrument(error_span!("client"));
+    let server = async move {
+        let conn = server.accept().await.unwrap().await.unwrap();
+        let mut stream = conn.accept_uni().await.unwrap();
+        let mut buf = [0u8; 2];
+        stream.read_exact(&mut buf).await.unwrap();
+        stream.stop(42u32.into()).unwrap();
+        conn
+    }
+    .instrument(error_span!("server"));
+    let (client, conn) = tokio::join!(client, server);
+    client.expect("timeout");
+    drop(conn);
+}
+
+#[tokio::test]
+async fn stream_stopped_2() {
+    let _guard = subscribe();
+    let endpoint = endpoint();
+
+    let (conn, _server_conn) = tokio::try_join!(
+        endpoint
+            .connect(endpoint.local_addr().unwrap(), "localhost")
+            .unwrap(),
+        async { endpoint.accept().await.unwrap().await }
+    )
+    .unwrap();
+    let send_stream = conn.open_uni().await.unwrap();
+    let stopped = tokio::time::timeout(Duration::from_millis(100), send_stream.stopped())
+        .instrument(error_span!("stopped"));
+    tokio::pin!(stopped);
+    // poll the future once so that the waker is registered.
+    tokio::select! {
+        biased;
+        _x = &mut stopped => {},
+        _x = std::future::ready(()) => {}
+    }
+    // drop the send stream
+    drop(send_stream);
+    // make sure the stopped future still resolves
+    let res = stopped.await;
+    assert_eq!(res, Ok(Ok(None)));
+}
