feat: Make the future returned by SendStream::stopped static

[Frando]

Currently, the future returned by SendStream::stopped is bound to &mut self, and therefore cannot be stored and polled while also sending on the stream. This PR makes the future returned by SendStream::stopped be 'static, so that it can be stored in futures combinators while also still sending on the stream.

This is useful when you want to clean up local state once the corresponding receiver to a send stream is closed or dropped, independently of actually using the stream. Especially if sends occur infrequently, tracking the stopped notification allows to clean up state at the earliest time and not only once you try to send again.

There's a complication though: Once SendStream::stopped returns a 'static future, we'd have to account for the fact that there may be multiple futures waiting for a wakeup once the stream is stopped. Therefore, we'd need to store a list of wakers instead of a single waker per stream. I used a Slab, open to other ideas how to achieve this.

For the intended functionality I don't actually need the possibility to create multiple futures waiting for the stopped wakeup. However, I didn't yet come up with an API that encodes that only the last future returned from SendStream::stopped would be woken up reliably - if you have ideas here, please share.

I added two tests with various combinations of using stopped while stopping and/or dropping the send stream.

[Ralith]

This is technically a breaking change since we're going from &mut self to &self, though it's the sort of esoteric breakage we could theoretically sneak through.

[gretchenfrage]

I understand now that the loop is necessary. Btw, I really like how you included tests for this.

I still think it can simplify things to have the helper function only try to return the Poll::Ready value, and not also have it saddled with trying to clone the Notify:

Diff

diff --git a/quinn/src/mutex.rs b/quinn/src/mutex.rs
index 7c24df46..df44b20f 100644
--- a/quinn/src/mutex.rs
+++ b/quinn/src/mutex.rs
@@ -110,7 +110,7 @@ mod tracking {
 }
 
 #[cfg(feature = "lock_tracking")]
-pub(crate) use tracking::Mutex;
+pub(crate) use tracking::{Mutex, MutexGuard};
 
 #[cfg(not(feature = "lock_tracking"))]
 mod non_tracking {
@@ -160,4 +160,4 @@ mod non_tracking {
 }
 
 #[cfg(not(feature = "lock_tracking"))]
-pub(crate) use non_tracking::Mutex;
+pub(crate) use non_tracking::{Mutex, MutexGuard};
diff --git a/quinn/src/send_stream.rs b/quinn/src/send_stream.rs
index 759734ea..87581dc8 100644
--- a/quinn/src/send_stream.rs
+++ b/quinn/src/send_stream.rs
@@ -1,7 +1,6 @@
 use std::{
     future::Future,
     io,
-    ops::ControlFlow,
     pin::Pin,
     sync::Arc,
     task::{Context, Poll, ready},
@@ -10,11 +9,11 @@ use std::{
 use bytes::Bytes;
 use proto::{ClosedStream, ConnectionError, FinishError, StreamId, Written};
 use thiserror::Error;
-use tokio::sync::Notify;
 
 use crate::{
     VarInt,
     connection::{ConnectionRef, State},
+    mutex::MutexGuard,
 };
 
 /// A stream that can only be used to send data
@@ -219,10 +218,10 @@ impl SendStream {
                 let notify;
                 {
                     let mut conn = conn.state.lock("SendStream::stopped");
-                    notify = match stopped_or_notify(&mut conn, stream, is_0rtt) {
-                        ControlFlow::Break(res) => return res,
-                        ControlFlow::Continue(notify) => notify,
-                    };
+                    if let Some(result) = stopped_output(&mut conn, stream, is_0rtt) {
+                        break result;
+                    }
+                    notify = Arc::clone(&conn.stopped.entry(stream).or_default());
                     notify.notified()
                 }
                 .await
@@ -251,26 +250,18 @@ impl SendStream {
     }
 }
 
-fn stopped_or_notify(
-    conn: &mut State,
+fn stopped_output(
+    conn: &mut MutexGuard<State>,
     stream: StreamId,
     is_0rtt: bool,
-) -> ControlFlow<Result<Option<VarInt>, StoppedError>, Arc<Notify>> {
-    use ControlFlow::*;
+) -> Option<Result<Option<VarInt>, StoppedError>> {
     if is_0rtt && conn.check_0rtt().is_err() {
-        return Break(Err(StoppedError::ZeroRttRejected));
+        return Some(Err(StoppedError::ZeroRttRejected));
     }
     match conn.inner.send_stream(stream).stopped() {
-        Err(ClosedStream { .. }) => Break(Ok(None)),
-        Ok(Some(error_code)) => Break(Ok(Some(error_code))),
-        Ok(None) => {
-            if let Some(e) = &conn.error {
-                Break(Err(e.clone().into()))
-            } else {
-                let notify = conn.stopped.entry(stream).or_default().clone();
-                Continue(notify)
-            }
-        }
+        Err(ClosedStream { .. }) => Some(Ok(None)),
+        Ok(Some(error_code)) => Some(Ok(Some(error_code))),
+        Ok(None) => conn.error.clone().map(|e| Err(e.into())),
     }
 }
 

(Consider that this is net-negative 9 LOC, avoids needing to import control flow, and allows an if-let-else chain to be replaced a simple option map.)

[Frando]

I pushed another commit that removes the ControlFlow and instead returns Option as suggested by @gretchenfrage.

[Frando]

I also rebased onto main (had conflicts in the import section), combined the commits into one and added a proper commit message.

[djc]

This is technically a breaking change since we're going from &mut self to &self, though it's the sort of esoteric breakage we could theoretically sneak through.

Not sure "esoteric" is the right word here, but I think rustc will gladly allow reborrowing &mut as & so I guess it will be mostly fine?

[gretchenfrage]

Thanks! Haven't re-looked in detail, but leaving a comment review which should clear my "changes requested" state. The fact that the other maintainers approved is sufficient.

Apparently me leaving a comment review doesn't supersede my request changes review as I previously thought.

[gretchenfrage]

And it seems like we all agree that we're going to consider this not to be semver breaking in a way that matters