From a1add92c20c02ba39422a69392c149c6dbe5a81d Mon Sep 17 00:00:00 2001
From: Benjamin Saunders <ben.e.saunders@gmail.com>
Date: Thu, 25 Apr 2024 19:04:05 -0700
Subject: [PATCH] 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            | 34 ++++++++++++++------
 quinn-proto/src/connection/packet_builder.rs |  2 ++
 2 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs
index 021fb8f60..14e40d346 100644
--- a/quinn-proto/src/connection/mod.rs
+++ b/quinn-proto/src/connection/mod.rs
@@ -475,6 +475,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 mut segment_size = self.path.current_mtu();
 
         // Send PATH_CHALLENGE for a previous path if necessary
         if let Some(ref mut prev_path) = self.prev_path {
@@ -559,9 +560,9 @@ impl Connection {
                 && self.peer_supports_ack_frequency();
         }
 
-        // Reserving capacity can provide more capacity than we asked for.
-        // However we are not allowed to write more than MTU size. Therefore
-        // the maximum capacity is tracked separately.
+        // Reserving capacity can provide more capacity than we asked for. However, we are not
+        // allowed to write more than `segment_size`. Therefore the maximum capacity is tracked
+        // separately.
         let mut buf_capacity = 0;
 
         let mut coalesce = true;
@@ -663,17 +664,24 @@ 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 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);
+                    }
 
                     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.
+                        segment_size = buf.len() as u16;
+                    }
                 }
 
                 // Allocate space for another datagram
-                buf_capacity += self.path.current_mtu() as usize;
+                buf_capacity += segment_size as usize;
                 if buf.capacity() < buf_capacity {
                     // We reserve the maximum space for sending `max_datagrams` upfront
                     // to avoid any reallocations if more datagrams have to be appended later on.
@@ -683,12 +691,18 @@ impl Connection {
                     // (e.g. purely containing ACKs), modern memory allocators
                     // (e.g. mimalloc and jemalloc) will pool certain allocation sizes
                     // and therefore this is still rather efficient.
-                    buf.reserve(max_datagrams * self.path.current_mtu() as usize);
+                    buf.reserve(max_datagrams * segment_size as usize);
                 }
                 num_datagrams += 1;
                 coalesce = true;
                 pad_datagram = false;
                 datagram_start = buf.len();
+
+                debug_assert_eq!(
+                    datagram_start % usize::from(segment_size),
+                    0,
+                    "datagrams in a GSO batch must be aligned to the segment size"
+                );
             } else {
                 // We can append/coalesce the next packet into the current
                 // datagram.
@@ -946,7 +960,7 @@ impl Connection {
             },
             segment_size: match num_datagrams {
                 1 => None,
-                _ => Some(self.path.current_mtu() as usize),
+                _ => Some(segment_size as usize),
             },
             src_ip: self.local_ip,
         })
diff --git a/quinn-proto/src/connection/packet_builder.rs b/quinn-proto/src/connection/packet_builder.rs
index 2770fc49f..f64db1dda 100644
--- a/quinn-proto/src/connection/packet_builder.rs
+++ b/quinn-proto/src/connection/packet_builder.rs
@@ -125,6 +125,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;