Match GSO segment size to the first datagram, not the MTU

Reduces the amount of padding required when sending homogeneously
sized non-fragmentable data like application datagrams.

quinn-proto/src/connection/mod.rs

@@ -474,7 +474,7 @@ impl Connection {
         // Position in `buf` of the first byte of the current UDP datagram. When coalescing QUIC
         // packets, this can be earlier than the start of the current QUIC packet.
         let mut datagram_start = 0;
-        let segment_size = usize::from(self.path.current_mtu());
+        let mut segment_size = usize::from(self.path.current_mtu());
 
         // Send PATH_CHALLENGE for a previous path if necessary
         if let Some(ref mut prev_path) = self.prev_path {
@@ -666,13 +666,30 @@ impl Connection {
 
                 // Finish current packet
                 if let Some(mut builder) = builder_storage.take() {
-                    // Pad the packet to make it suitable for sending with GSO
-                    // which will always send the maximum PDU.
-                    builder.pad_to(self.path.current_mtu());
+                    if pad_datagram {
+                        builder.pad_to(MIN_INITIAL_SIZE);
+                    }
+
+                    if num_datagrams > 1 {
+                        // Pad the current packet to GSO segment size so it can be included in the
+                        // GSO batch.
+                        builder.pad_to(segment_size as u16);
+                    }
 
                     builder.finish_and_track(now, self, sent_frames.take(), buf);
 
-                    debug_assert_eq!(buf.len(), buf_capacity, "Packet must be padded");
+                    if num_datagrams == 1 {
+                        // Set the segment size for this GSO batch to the size of the first UDP
+                        // datagram in the batch. Larger data that cannot be fragmented
+                        // (e.g. application datagrams) will be included in a future batch. When
+                        // sending large enough volumes of data for GSO to be useful, we expect
+                        // packet sizes to usually be consistent, e.g. populated by max-size STREAM
+                        // frames or uniformly sized datagrams.
+                        segment_size = buf.len();
+                        // Clip the unused capacity out of the buffer so future packets don't
+                        // overrun
+                        buf_capacity = buf.len();
+                    }
                 }
 
                 // Allocate space for another datagram

quinn-proto/src/connection/packet_builder.rs

@@ -166,6 +166,8 @@ impl PacketBuilder {
         })
     }
 
+    /// Append the minimum amount of padding such that, after encryption, the packet will occupy at
+    /// least `min_size` bytes
     pub(super) fn pad_to(&mut self, min_size: u16) {
         let prev = self.min_size;
         self.min_size = self.datagram_start + (min_size as usize) - self.tag_len;

quinn-proto/src/tests/mod.rs

@@ -3012,3 +3012,33 @@ fn stream_gso() {
     let final_ios = pair.client_conn_mut(client_ch).stats().udp_tx.ios;
     assert_eq!(final_ios - initial_ios, 2);
 }
+
+#[test]
+fn datagram_gso() {
+    let _guard = subscribe();
+    let mut pair = Pair::default();
+    let (client_ch, _) = pair.connect();
+
+    let initial_ios = pair.client_conn_mut(client_ch).stats().udp_tx.ios;
+    let initial_bytes = pair.client_conn_mut(client_ch).stats().udp_tx.bytes;
+
+    // Send 10 datagrams above half the MTU, which fits inside a `tests::util::MAX_DATAGRAMS`
+    // datagram batch
+    info!("sending");
+    const DATAGRAM_LEN: usize = 1024;
+    const DATAGRAMS: usize = 10;
+    for _ in 0..DATAGRAMS {
+        pair.client_datagrams(client_ch)
+            .send(Bytes::from_static(&[0; DATAGRAM_LEN]), false)
+            .unwrap();
+    }
+    pair.drive();
+    let final_ios = pair.client_conn_mut(client_ch).stats().udp_tx.ios;
+    let final_bytes = pair.client_conn_mut(client_ch).stats().udp_tx.bytes;
+    assert_eq!(final_ios - initial_ios, 1);
+    // Expected overhead: flags + CID + PN + tag + frame type + frame length = 1 + 8 + 1 + 16 + 1 + 2 = 29
+    assert_eq!(
+        final_bytes - initial_bytes,
+        ((29 + DATAGRAM_LEN) * DATAGRAMS) as u64
+    );
+}