Move allocating new datagrams into TransmitBuf

This moves the logic of updating the buffer for new datagrams into a
function on the TransmitBuf.  This centralises all the manipulation of
these variables into one logical place, ensuring all the invariants
are upheld.

Author: flub

quinn-proto/src/connection/mod.rs

@@ -661,68 +661,42 @@ impl Connection {
 
                     builder.finish_and_track(now, self, sent_frames.take(), buf.buf);
 
-                    if buf.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.
-                        buf.segment_size = buf.len();
-                        // Clip the unused capacity out of the buffer so future packets don't
-                        // overrun
-                        buf.buf_capacity = buf.len();
-
-                        // Check whether the data we planned to send will fit in the reduced segment
-                        // size. If not, bail out and leave it for the next GSO batch so we don't
-                        // end up trying to send an empty packet. We can't easily compute the right
-                        // segment size before the original call to `space_can_send`, because at
-                        // that time we haven't determined whether we're going to coalesce with the
-                        // first datagram or potentially pad it to `MIN_INITIAL_SIZE`.
-                        if space_id == SpaceId::Data {
-                            let frame_space_1rtt = buf
-                                .segment_size
-                                .saturating_sub(self.predict_1rtt_overhead(Some(pn)));
-                            if self.space_can_send(space_id, frame_space_1rtt).is_empty() {
-                                break;
-                            }
+                    if buf.num_datagrams == 1 && space_id == SpaceId::Data {
+                        // Now that we know the size of the first datagram, check whether
+                        // the data we planned to send will fit in the next segment.  If
+                        // not, bails out andleave it for the next GSO batch.  We can't
+                        // easily compute the right segment size before the original call to
+                        // `space_can_send`, because at that time we haven't determined
+                        // whether we're going to coalesce with the first datagram or
+                        // potentially pad it to `MIN_INITIAL_SIZE`.
+
+                        buf.clip_datagram_size();
+
+                        let frame_space_1rtt = buf
+                            .segment_size
+                            .saturating_sub(self.predict_1rtt_overhead(Some(pn)));
+                        if self.space_can_send(space_id, frame_space_1rtt).is_empty() {
+                            break;
                         }
                     }
                 }
 
-                // Allocate space for another datagram
-                let next_datagram_size_limit = match self.spaces[space_id].loss_probes {
-                    0 => buf.segment_size,
+                // Start the next datagram
+                match self.spaces[space_id].loss_probes {
+                    0 => buf.start_new_datagram(),
                     _ => {
                         self.spaces[space_id].loss_probes -= 1;
                         // Clamp the datagram to at most the minimum MTU to ensure that loss probes
                         // can get through and enable recovery even if the path MTU has shrank
                         // unexpectedly.
-                        usize::from(INITIAL_MTU)
+                        buf.start_new_datagram_with_size(std::cmp::min(
+                            usize::from(INITIAL_MTU),
+                            buf.segment_size,
+                        ));
                     }
                 };
-                buf.buf_capacity += next_datagram_size_limit;
-                if buf.buf.capacity() < buf.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.
-                    // Benchmarks have shown shown a 5-10% throughput improvement
-                    // compared to continuously resizing the datagram buffer.
-                    // While this will lead to over-allocation for small transmits
-                    // (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.buf.reserve(buf.max_datagrams * buf.segment_size);
-                }
-                buf.num_datagrams += 1;
                 coalesce = true;
                 pad_datagram = false;
-                buf.datagram_start = buf.len();
-
-                debug_assert_eq!(
-                    buf.datagram_start % buf.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.
@@ -924,8 +898,8 @@ impl Connection {
                 .mtud
                 .poll_transmit(now, self.packet_number_filter.peek(&self.spaces[space_id]))?;
 
-            buf.buf_capacity = probe_size as usize;
-            buf.buf.reserve(buf.buf_capacity);
+            debug_assert_eq!(buf.num_datagrams, 0);
+            buf.start_new_datagram_with_size(probe_size as usize);
 
             debug_assert_eq!(buf.datagram_start, 0);
             let mut builder =
@@ -949,7 +923,6 @@ impl Connection {
             builder.finish_and_track(now, self, Some(sent_frames), buf.buf);
 
             self.stats.path.sent_plpmtud_probes += 1;
-            buf.num_datagrams = 1;
 
             trace!(?probe_size, "writing MTUD probe");
         }
@@ -986,11 +959,7 @@ impl Connection {
     }
 
     /// Send PATH_CHALLENGE for a previous path if necessary
-    fn send_path_challenge(
-        &mut self,
-        now: Instant,
-        mut buf: &mut TransmitBuf<'_>,
-    ) -> Option<Transmit> {
+    fn send_path_challenge(&mut self, now: Instant, buf: &mut TransmitBuf<'_>) -> Option<Transmit> {
         let (prev_cid, prev_path) = self.prev_path.as_mut()?;
         if !prev_path.challenge_pending {
             return None;
@@ -1005,9 +974,7 @@ impl Connection {
             SpaceId::Data,
             "PATH_CHALLENGE queued without 1-RTT keys"
         );
-        buf.buf.reserve(MIN_INITIAL_SIZE as usize);
-
-        buf.buf_capacity = buf.buf.capacity();
+        buf.start_new_datagram_with_size(MIN_INITIAL_SIZE as usize);
 
         // Use the previous CID to avoid linking the new path with the previous path. We
         // don't bother accounting for possible retirement of that prev_cid because this is

quinn-proto/src/connection/transmit_buf.rs

@@ -74,6 +74,80 @@ impl<'a> TransmitBuf<'a> {
     pub(super) fn is_empty(&self) -> bool {
         self.len() == 0
     }
+
+    /// Clips the datagram size to the current size
+    ///
+    /// Only valid for the first datagram, when the datagram might be smaller than the
+    /// segment size.  Needed before estimating the available space in the next datagram
+    /// based on [`TransmitBuf::segment_size`].
+    ///
+    /// Use [`TransmitBuf::start_new_datagram_with_size`] if you need to reduce the size of
+    /// the last datagram in a batch.
+    pub(super) fn clip_datagram_size(&mut self) {
+        debug_assert_eq!(self.num_datagrams, 1);
+        self.segment_size = self.buf.len();
+    }
+
+    /// Starts a new datagram in the transmit buffer
+    ///
+    /// If this starts the second datagram the segment size will be set to the size of the
+    /// first datagram.
+    ///
+    /// If the underlying buffer does not have enough capacity yet this will allocate enough
+    /// capacity for all the datagrams allowed in a single batch.  Use
+    /// [`TransmitBuf::start_new_datagram_with_size`] if you know you will need less.
+    pub(super) fn start_new_datagram(&mut self) {
+        // We reserve the maximum space for sending `max_datagrams` upfront to avoid any
+        // reallocations if more datagrams have to be appended later on.  Benchmarks have
+        // shown a 5-10% throughput improvment compared to continuously resizing the
+        // datagram buffer.  While this will lead to over-allocation for small transmits
+        // (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.
+        let max_capacity_hint = self.max_datagrams * self.segment_size;
+        self.new_datagram_inner(self.segment_size, max_capacity_hint)
+    }
+
+    /// Starts a datagram with a custom datagram size
+    ///
+    /// This is a specialised version of [`TransmitBuf::start_new_datagram`] which sets the
+    /// datagram size.  Useful for e.g. PATH_CHALLENGE, tail-loss probes or MTU probes.
+    ///
+    /// After the first datagram you can never increase the segment size.  If you decrease
+    /// the size of a datagram in a batch, it must be the last datagram of the batch.
+    pub(super) fn start_new_datagram_with_size(&mut self, datagram_size: usize) {
+        // Only reserve space for this datagram, usually it is the last one in the batch.
+        let max_capacity_hint = datagram_size;
+        self.new_datagram_inner(datagram_size, max_capacity_hint)
+    }
+
+    fn new_datagram_inner(&mut self, datagram_size: usize, max_capacity_hint: usize) {
+        debug_assert!(self.num_datagrams < self.max_datagrams);
+        if self.num_datagrams == 1 {
+            // Set the segment size to the size of the first datagram.
+            self.segment_size = self.buf.len();
+        }
+        if self.num_datagrams >= 1 {
+            debug_assert!(datagram_size <= self.segment_size);
+            if datagram_size < self.segment_size {
+                // If This is a GSO batch and this datagram is smaller than the segment
+                // size, this must be the last datagram in the batch.
+                self.max_datagrams = self.num_datagrams + 1;
+            }
+        }
+        self.datagram_start = self.buf.len();
+        debug_assert_eq!(
+            self.datagram_start % self.segment_size,
+            0,
+            "datagrams in a GSO batch must be aligned to the segment size"
+        );
+        self.buf_capacity = self.datagram_start + datagram_size;
+        if self.buf_capacity > self.buf.capacity() {
+            self.buf
+                .reserve_exact(max_capacity_hint - self.buf.capacity());
+        }
+        self.num_datagrams += 1;
+    }
 }
 
 unsafe impl BufMut for TransmitBuf<'_> {