fix(quic): resolve SendStream::stopped after a local reset

rs/moq-quic/README.md

@@ -90,6 +90,7 @@ Changes on top of the upstream commit, besides the renames:
 - **Stream credit on stop**: `RecvStream::stop` on a stream whose final size is already known (finished or reset) frees it and queues `MAX_STREAMS` at once, as reading the reset does, instead of leaving the credit until the next packet arrives. Not offered upstream yet; quinn `main` has the same gap.
 - **Drop resets** (`quic::driver::SendStream`): a send stream dropped before `finish` is reset with code 0, or with the peer's `STOP_SENDING` code, instead of quinn's implicit `finish`, so a cancelled writer never ends a truncated stream with a FIN. This matches moq-uring's driver and moq-net's transport contract. `quic::web` resets after the WebTransport header with the mapped code. quinn's `stream_stopped_2` finishes its stream before dropping it. Not offered upstream: it reverses a documented quinn behaviour.
 - **Reset secret** (`EndpointConfig::reset_secret`): derives the stateless reset key from an operator's secret with HKDF-SHA256, and keys a `HashedConnectionIdGenerator` from it too, since a restarted endpoint drops packets for connection IDs it can't validate before it would reset them. It also lowers `min_reset_interval` to 1ms, so a restart resets thousands of connections in seconds. Both drivers take one, so a restarted server resets its predecessor's connections. Not offered upstream: it only combines quinn's `reset_key` and `HashedConnectionIdGenerator::from_key`.
+- **Stopped after a local reset**: moq-tokio's driver resolves `SendStream::stopped` with `None` once the stream is reset locally, as it already did once the reset was acknowledged and the stream freed. No stream event reports a local reset, so quinn leaves a pending `stopped` waiting, and its waker registered, until the connection closes. Not offered upstream yet; quinn 0.11 has the same gap.
 - **Connection limits** (`Limits`, `Connection::set_limits`, and the drivers' `set_limits`): one validated call changes a live connection's stream limit and windows, refusing values a peer would reject and changing nothing then. moq-tokio and moq-uring re-export the type, so both runtimes check limits the same way. Not offered upstream: it only bundles quinn's existing setters.
 
 Recheck the BBR3 changes when cherry-picking quinn#2481 updates, or if it merges.

rs/moq-tokio/src/quic/driver/connection.rs

@@ -6,7 +6,7 @@ use std::{
     net::{IpAddr, SocketAddr},
     pin::Pin,
     sync::{
-        Arc,
+        Arc, OnceLock,
         atomic::{AtomicUsize, Ordering},
     },
     task::{Context, Poll, Waker, ready},
@@ -642,6 +642,19 @@ impl Connection {
         self.0.state.lock("stats").inner.stats()
     }
 
+    /// Whether the connection has freed the send stream's state
+    #[cfg(test)]
+    pub(crate) fn send_stream_freed(&self, id: StreamId) -> bool {
+        let mut state = self.0.state.lock("send_stream_freed");
+        state.inner.send_stream(id).stopped().is_err()
+    }
+
+    /// How many streams have `SendStream::stopped` waiters registered, to check none leak
+    #[cfg(test)]
+    pub(crate) fn stopped_waiters(&self) -> usize {
+        self.0.state.lock("stopped_waiters").stopped.len()
+    }
+
     /// The wire counters the driver updates as it sends and receives.
     #[cfg(feature = "net")]
     pub(crate) fn wire(&self) -> moq_net::transport::Wire {
@@ -1092,7 +1105,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, Arc<Notify>>,
+    pub(crate) stopped: FxHashMap<StreamId, Arc<StopWaiters>>,
     /// Always set to Some before the connection becomes drained
     pub(crate) error: Option<ConnectionError>,
     sender: Pin<Box<dyn UdpSender>>,
@@ -1279,9 +1292,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_notify(id, &mut self.stopped),
-                Stream(StreamEvent::Stopped { id, .. }) => {
-                    wake_stream_notify(id, &mut self.stopped);
+                Stream(StreamEvent::Finished { id }) => self.resolve_stopped(id, None),
+                Stream(StreamEvent::Stopped { id, error_code }) => {
+                    self.resolve_stopped(id, Some(error_code));
                     wake_stream(id, &mut self.blocked_writers);
                 }
                 PathUpdated => {
@@ -1372,6 +1385,14 @@ impl State {
         self.close(0u32.into(), Bytes::new(), shared);
     }
 
+    /// Resolves the stream's `stopped` waiters with the peer's STOP_SENDING code, or `None`
+    pub(crate) fn resolve_stopped(&mut self, id: StreamId, outcome: Option<VarInt>) {
+        if let Some(waiters) = self.stopped.remove(&id) {
+            let _ = waiters.outcome.set(outcome);
+            waiters.notify.notify_waiters();
+        }
+    }
+
     pub(crate) fn check_0rtt(&self) -> Result<(), ()> {
         if self.inner.is_handshaking()
             || self.inner.accepted_0rtt()
@@ -1411,16 +1432,19 @@ 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()
-    }
+/// The `SendStream::stopped` futures waiting on one stream
+#[derive(Debug, Default)]
+pub(crate) struct StopWaiters {
+    pub(crate) notify: Notify,
+    /// What a STOP_SENDING, a delivered finish, or a local reset settled, kept here because the
+    /// stream may be freed before a woken waiter polls. Once set, nothing later counts.
+    pub(crate) outcome: OnceLock<Option<VarInt>>,
 }
 
-fn wake_all_notify(wakers: &mut FxHashMap<StreamId, Arc<Notify>>) {
+fn wake_all_notify(wakers: &mut FxHashMap<StreamId, Arc<StopWaiters>>) {
     wakers
         .drain()
-        .for_each(|(_, notify)| notify.notify_waiters())
+        .for_each(|(_, waiters)| waiters.notify.notify_waiters())
 }
 
 /// Errors that can arise when sending a datagram

rs/moq-tokio/src/quic/driver/send_stream.rs

@@ -39,6 +39,9 @@ pub struct SendStream {
     finished: bool,
     /// The offset and code a drop before `finish` resets at
     drop_reset: (u64, VarInt),
+    /// `reset` or `reset_at` succeeded, with the peer's STOP_SENDING code if it arrived first, so
+    /// `stopped` has nothing left to wait for
+    reset: Option<Option<VarInt>>,
 }
 
 impl SendStream {
@@ -49,6 +52,7 @@ impl SendStream {
             is_0rtt,
             finished: false,
             drop_reset: (0, VarInt::from_u32(0)),
+            reset: None,
         }
     }
 
@@ -216,7 +220,11 @@ impl SendStream {
         if self.is_0rtt && conn.check_0rtt().is_err() {
             return Ok(());
         }
-        conn.inner.send_stream(self.stream).reset(error_code)?;
+        let mut stream = conn.inner.send_stream(self.stream);
+        let stopped = stream.stopped().ok().flatten();
+        stream.reset(error_code)?;
+        self.reset = Some(stopped);
+        conn.resolve_stopped(self.stream, stopped);
         conn.wake();
         Ok(())
     }
@@ -235,9 +243,11 @@ impl SendStream {
         if self.is_0rtt && conn.check_0rtt().is_err() {
             return Ok(());
         }
-        conn.inner
-            .send_stream(self.stream)
-            .reset_at(offset, error_code)?;
+        let mut stream = conn.inner.send_stream(self.stream);
+        let stopped = stream.stopped().ok().flatten();
+        stream.reset_at(offset, error_code)?;
+        self.reset = Some(stopped);
+        conn.resolve_stopped(self.stream, stopped);
         conn.wake();
         Ok(())
     }
@@ -276,8 +286,10 @@ impl SendStream {
     ///
     /// Yields `Some` with the stop error code if the peer stops the stream. Yields `None` if the
     /// local side [`finish()`](Self::finish)es the stream and then the peer acknowledges receipt
-    /// of all stream data (although not necessarily the processing of it), after which the peer
-    /// closing the stream is no longer meaningful.
+    /// of all stream data (although not necessarily the processing of it), or once the local side
+    /// resets it with [`reset()`](Self::reset) or [`reset_at()`](Self::reset_at), or by dropping
+    /// it unfinished, after which the peer closing the stream is no longer meaningful. A peer's
+    /// stop code still wins over a local reset if it arrived first, but not if it arrives after.
     ///
     /// 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
@@ -291,23 +303,43 @@ impl SendStream {
         let conn = self.conn.clone();
         let stream = self.stream;
         let is_0rtt = self.is_0rtt;
+        // Join the waiters now rather than on the first poll, so a reset in between, which no
+        // stream event reports, still reaches this future.
+        let waiters = {
+            let mut conn = conn.state.lock("SendStream::stopped");
+            // A reset settles it: a STOP_SENDING after it is recorded but no longer counts.
+            if let Some(stopped) = self.reset {
+                Err(Ok(stopped))
+            } else if let Some(output) = send_stream_stopped(&mut conn, stream, is_0rtt) {
+                Err(output)
+            } else {
+                Ok(conn.stopped.entry(stream).or_default().clone())
+            }
+        };
         async move {
+            let mut waiters = match waiters {
+                Ok(waiters) => waiters,
+                Err(output) => return output,
+            };
             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;
+                // we need to declare `waiters` outside of the block, and update it inside.
                 {
                     let mut conn = conn.state.lock("SendStream::stopped");
+                    if let Some(&stopped) = waiters.outcome.get() {
+                        return Ok(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()
+                    // A wakeup that resolved nothing removed the waiters; join the current ones
+                    waiters = conn.stopped.entry(stream).or_default().clone();
+                    waiters.notify.notified()
                 }
                 .await
             }
@@ -390,18 +422,20 @@ impl Drop for SendStream {
         }
         // Fails harmlessly when already reset.
         let mut stream = conn.inner.send_stream(self.stream);
-        let reset = match stream.stopped() {
+        let stopped = stream.stopped().ok().flatten();
+        let reset = match stopped {
             // Echo the code, as RFC 9000 section 3.5 asks, even after `finish`: the peer reads
             // nothing more, so retransmitting the rest is wasted.
-            Ok(Some(code)) => stream.reset(code),
-            _ if self.finished => return,
-            _ => {
+            Some(code) => stream.reset(code),
+            None if self.finished => return,
+            None => {
                 let (offset, code) = self.drop_reset;
                 let offset = offset.min(stream.offset());
                 stream.reset_at(offset, code)
             }
         };
         if reset.is_ok() {
+            conn.resolve_stopped(self.stream, stopped);
             conn.wake();
         }
     }