fix(quic): let an oversized datagram step aside behind control frames

rs/moq-quic/src/connection/mod.rs

@@ -781,11 +781,9 @@ impl Connection {
                         // overrun
                         buf_capacity = buf.len();
 
-                        // A GSO batch cannot have a datagram larger than its first, so a small
-                        // first packet caps every later one. End the batch when that cap cannot
-                        // even hold a packet header and frames; a control-only packet (a loss
-                        // probe, say) should not drag the data queued behind it into tiny
-                        // segments.
+                        // Every later segment is capped at the first one's size. If that is
+                        // too small to be worth sending, end the batch so the rest goes out
+                        // at full size.
                         if segment_size < MIN_PACKET_SPACE {
                             break;
                         }

rs/moq-quic/src/connection/streams/saturated.rs

@@ -718,7 +718,7 @@ fn a_datagram_that_does_not_fit_leads_the_next_packet() {
 }
 
 /// A packet clamped below a queued datagram, such as a loss probe, still carries what fits
-/// behind it, and the datagram keeps its place.
+/// behind it, even behind its own control frames, and the datagram keeps its place.
 #[test]
 fn a_clamped_packet_sends_what_fits() {
     let mut sender = Sender::new();
@@ -728,7 +728,7 @@ fn a_clamped_packet_sends_what_fits() {
     let stream = sender.open(Some(&group), 0, 2_000);
     assert!(sender.state.can_send_frames(500));
 
-    sender.packet_of(500);
+    sender.packet_after(&[PING], 500);
     assert_eq!(sender.order, [Sent::Stream(stream)]);
     assert_eq!(sender.drain_frames(1)[0], Sent::Datagram(large));
 }

rs/moq-quic/src/connection/streams/state.rs

@@ -777,9 +777,10 @@ impl StreamsState {
     ///
     /// `payload` is the packet's frame space in `buf`, which may already hold control frames.
     /// A datagram cannot be split, so one too large for the room left ends the packet and
-    /// leads the next one. Only when it would leave the packet empty, which a packet clamped
-    /// below the path MTU can, does it step aside for what fits behind it, since
-    /// [`Self::can_send_frames`] promised a frame.
+    /// leads the next one. One too large for even an empty packet, which a packet clamped
+    /// below the path MTU can be, steps aside for what fits behind it instead, since
+    /// [`Self::can_send_frames`] promised a frame and control frames alone would leave the
+    /// packet short.
     pub(crate) fn write_frames(
         &mut self,
         buf: &mut Vec<u8>,
@@ -797,7 +798,7 @@ impl StreamsState {
                 Some(Next::Datagram(group)) => {
                     let (_, frame) = self.pending.datagram(group);
                     if buf.len() + frame > max_buf_size {
-                        if buf.len() > payload.start {
+                        if payload.start + frame <= max_buf_size {
                             break;
                         }
                         self.pending.defer_datagram(group);

rs/moq-quic/src/tests/mod.rs

@@ -3419,6 +3419,57 @@ fn tail_loss_probe_keeps_ping_when_datagram_does_not_fit() {
     assert_eq!(server_stats.frame_rx.datagram, 1);
 }
 
+/// A loss probe ended short by a datagram that fills a whole probe, but not one behind a PING,
+/// ends the GSO batch instead of repeating its tiny size for the stream data queued behind.
+#[test]
+fn tail_loss_probe_ended_short_ends_the_gso_batch() {
+    let _guard = subscribe();
+
+    const PATH_MTU: u16 = 1452;
+
+    let client_config = {
+        let mut config = client_config();
+        Arc::get_mut(&mut config.transport)
+            .unwrap()
+            .initial_mtu(PATH_MTU)
+            .mtu_discovery_config(None);
+        config
+    };
+
+    let mut pair = Pair::default();
+    pair.mtu = PATH_MTU as usize;
+    let (client_ch, server_ch) = pair.connect_with(client_config);
+    pair.client_conn_mut(client_ch).disable_peer_ack_frequency();
+
+    // Establish an outstanding ack-eliciting packet and discard it.
+    pair.client_conn_mut(client_ch).ping();
+    pair.drive_client();
+    assert!(!pair.server.inbound.is_empty());
+    pair.server.inbound.clear();
+
+    // Advance to the PTO without transmitting the queued loss probe yet.
+    pair.step();
+
+    // Fills a probe capped to INITIAL_MTU exactly.
+    // Overhead: flags + CID + PN + tag + frame type + frame length = 1 + 8 + 1 + 16 + 1 + 2 = 29
+    const DATAGRAM_LEN: usize = INITIAL_MTU as usize - 29;
+    pair.client_datagrams(client_ch)
+        .send(vec![0; DATAGRAM_LEN].into(), Default::default())
+        .unwrap();
+    let s = pair
+        .client_streams(client_ch)
+        .open(Dir::Uni, Default::default())
+        .unwrap();
+    pair.client_send(client_ch, s).write(&[0; 4096]).unwrap();
+    pair.client_conn_mut(client_ch).ping();
+
+    pair.drive();
+
+    let server_stats = pair.server_conn_mut(server_ch).stats();
+    assert_eq!(server_stats.frame_rx.datagram, 1);
+    assert!(server_stats.frame_rx.stream > 0);
+}
+
 // Respect max_datagrams when TLP happens
 #[test]
 fn tail_loss_respect_max_datagrams() {