proto: send STREAMS_BLOCKED when stream limit is hit (#2579)

Author: cong-or

quinn-proto/src/connection/spaces.rs

@@ -319,6 +319,7 @@ pub(super) struct LostPacket {
 pub struct Retransmits {
     pub(super) max_data: bool,
     pub(super) max_stream_id: [bool; 2],
+    pub(super) streams_blocked: [bool; 2],
     pub(super) reset_stream: Vec<(StreamId, VarInt)>,
     pub(super) stop_sending: Vec<frame::StopSending>,
     pub(super) max_stream_data: FxHashSet<StreamId>,
@@ -350,6 +351,7 @@ impl Retransmits {
     pub(super) fn is_empty(&self, streams: &StreamsState) -> bool {
         !self.max_data
             && !self.max_stream_id.into_iter().any(|x| x)
+            && !self.streams_blocked.into_iter().any(|x| x)
             && self.reset_stream.is_empty()
             && self.stop_sending.is_empty()
             && self
@@ -372,6 +374,7 @@ impl ::std::ops::BitOrAssign for Retransmits {
         self.max_data |= rhs.max_data;
         for dir in Dir::iter() {
             self.max_stream_id[dir as usize] |= rhs.max_stream_id[dir as usize];
+            self.streams_blocked[dir as usize] |= rhs.streams_blocked[dir as usize];
         }
         self.reset_stream.extend_from_slice(&rhs.reset_stream);
         self.stop_sending.extend_from_slice(&rhs.stop_sending);

quinn-proto/src/connection/streams/mod.rs

@@ -48,8 +48,8 @@ impl<'a> Streams<'a> {
             return None;
         }
 
-        // TODO: Queue STREAM_ID_BLOCKED if this fails
         if self.state.next[dir as usize] >= self.state.max[dir as usize] {
+            self.state.streams_blocked[dir as usize] = true;
             return None;
         }
 

quinn-proto/src/connection/streams/state.rs

@@ -137,6 +137,8 @@ pub struct StreamsState {
 
     /// The shrink to be applied to local_max_data when receive_window is shrunk
     receive_window_shrink_debt: u64,
+    /// Whether the locally-initiated stream limit has been hit, per direction
+    pub(super) streams_blocked: [bool; 2],
 }
 
 impl StreamsState {
@@ -181,6 +183,7 @@ impl StreamsState {
             initial_max_stream_data_bidi_local: 0u32.into(),
             initial_max_stream_data_bidi_remote: 0u32.into(),
             receive_window_shrink_debt: 0,
+            streams_blocked: [false, false],
         };
 
         for dir in Dir::iter() {
@@ -525,6 +528,32 @@ impl StreamsState {
                 Dir::Bi => stats.max_streams_bidi += 1,
             }
         }
+
+        // STREAMS_BLOCKED
+        for dir in Dir::iter() {
+            if self.streams_blocked[dir as usize] {
+                pending.streams_blocked[dir as usize] = true;
+                self.streams_blocked[dir as usize] = false;
+            }
+
+            if !pending.streams_blocked[dir as usize] || buf.len() + 9 >= max_size {
+                continue;
+            }
+
+            pending.streams_blocked[dir as usize] = false;
+            retransmits.get_or_create().streams_blocked[dir as usize] = true;
+            let limit = self.max[dir as usize];
+            trace!(limit, "STREAMS_BLOCKED ({:?})", dir);
+            buf.write(match dir {
+                Dir::Uni => frame::FrameType::STREAMS_BLOCKED_UNI,
+                Dir::Bi => frame::FrameType::STREAMS_BLOCKED_BIDI,
+            });
+            buf.write_var(limit);
+            match dir {
+                Dir::Uni => stats.streams_blocked_uni += 1,
+                Dir::Bi => stats.streams_blocked_bidi += 1,
+            }
+        }
     }
 
     pub(crate) fn write_stream_frames(
@@ -697,6 +726,7 @@ impl StreamsState {
         let current = &mut self.max[dir as usize];
         if count > *current {
             *current = count;
+            self.streams_blocked[dir as usize] = false;
             self.events.push_back(StreamEvent::Available { dir });
         }
 

quinn-proto/src/tests/mod.rs

@@ -980,6 +980,69 @@ fn stream_id_limit() {
     let _ = chunks.finalize();
 }
 
+#[test]
+fn streams_blocked() {
+    let _guard = subscribe();
+    let server = ServerConfig {
+        transport: Arc::new(TransportConfig {
+            max_concurrent_uni_streams: 1u32.into(),
+            ..TransportConfig::default()
+        }),
+        ..server_config()
+    };
+    let mut pair = Pair::new(Default::default(), server);
+    let (client_ch, server_ch) = pair.connect();
+
+    // Use up the only stream slot, then try to open another
+    let s = pair
+        .client_streams(client_ch)
+        .open(Dir::Uni)
+        .expect("first uni stream");
+    assert_eq!(pair.client_streams(client_ch).open(Dir::Uni), None);
+
+    // Send data so the STREAMS_BLOCKED piggybacks on an outgoing packet
+    pair.client_send(client_ch, s).write(b"hi").unwrap();
+    pair.drive();
+
+    assert_eq!(
+        pair.client_conn_mut(client_ch)
+            .stats()
+            .frame_tx
+            .streams_blocked_uni,
+        1
+    );
+    assert_eq!(
+        pair.server_conn_mut(server_ch)
+            .stats()
+            .frame_rx
+            .streams_blocked_uni,
+        1
+    );
+}
+
+#[test]
+fn streams_blocked_not_sent_under_limit() {
+    let _guard = subscribe();
+    let mut pair = Pair::default();
+    let (client_ch, _server_ch) = pair.connect();
+
+    // Default config allows many streams; opening one should not trigger STREAMS_BLOCKED
+    let s = pair
+        .client_streams(client_ch)
+        .open(Dir::Uni)
+        .expect("open stream");
+    pair.client_send(client_ch, s).write(b"hi").unwrap();
+    pair.drive();
+
+    assert_eq!(
+        pair.client_conn_mut(client_ch)
+            .stats()
+            .frame_tx
+            .streams_blocked_uni,
+        0
+    );
+}
+
 #[test]
 fn key_update_simple() {
     let _guard = subscribe();