Move pad_datagram to PacketBuilder::finish_and_track

You always need to remember to handle pad_datagram if needed.  While
really this always happens just before the call to finish_and_track.
Instead this can be done in finish_and_track without any logic change,
and this helps avoiding mistakes.

Author: flub

quinn-proto/src/connection/mod.rs

@@ -583,7 +583,7 @@ impl Connection {
             }
 
             // If the datagram is full, we need to start a new one.
-            if buf.len() == buf.datagram_max_offset() {
+            if buf.datagram_remaining_mut() == 0 {
                 // Is 1 more datagram allowed?
                 if buf.num_datagrams() >= buf.max_datagrams() {
                     // No more datagrams allowed
@@ -656,7 +656,7 @@ impl Connection {
                 pad_datagram = false;
             }
 
-            debug_assert!(buf.datagram_max_offset() - buf.len() >= MIN_PACKET_SPACE);
+            debug_assert!(buf.datagram_remaining_mut() >= MIN_PACKET_SPACE);
 
             //
             // From here on, we've determined that a packet will definitely be sent.
@@ -750,10 +750,7 @@ impl Connection {
                         ),
                     }
                 }
-                if pad_datagram {
-                    builder.pad_to(MIN_INITIAL_SIZE);
-                }
-                builder.finish_and_track(now, self, path_id, sent_frames);
+                builder.finish_and_track(now, self, path_id, sent_frames, pad_datagram);
                 if space_id == self.highest_space {
                     // Don't send another close packet
                     self.close = false;
@@ -788,6 +785,7 @@ impl Connection {
                             non_retransmits: true,
                             ..SentFrames::default()
                         },
+                        false,
                     );
                     self.stats.udp_tx.on_sent(1, buf.len());
                     return Some(Transmit {
@@ -850,17 +848,18 @@ impl Connection {
                 // Are we allowed to coalesce AND is there enough space for another *packet*
                 // in this datagram?
                 if coalesce
-                    && builder.buf.segment_size() - builder.predict_packet_size() > MIN_PACKET_SPACE
+                    && builder
+                        .buf
+                        .datagram_remaining_mut()
+                        .saturating_sub(builder.predict_packet_end())
+                        > MIN_PACKET_SPACE
                 {
                     // We can append/coalesce the next packet into the current
                     // datagram. Finish the current packet without adding extra padding.
-                    builder.finish_and_track(now, self, path_id, sent_frames);
+                    builder.finish_and_track(now, self, path_id, sent_frames, false);
                 } else {
                     // We need a new datagram for the next packet.  Finish the current
                     // packet with padding.
-                    if pad_datagram {
-                        builder.pad_to(MIN_INITIAL_SIZE);
-                    }
                     if builder.buf.num_datagrams() > 1 {
                         // If too many padding bytes would be required to continue the
                         // GSO batch after this packet, end the GSO batch here. Ensures
@@ -877,16 +876,14 @@ impl Connection {
                         // are the only packets for which we might grow `buf_capacity`
                         // by less than `segment_size`.
                         const MAX_PADDING: usize = 16;
-                        let packet_len_unpadded = builder.predict_packet_size();
-                        if packet_len_unpadded + MAX_PADDING < builder.buf.segment_size()
-                            || builder.buf.datagram_start_offset() + builder.buf.segment_size()
-                                > builder.buf.datagram_max_offset()
+                        if builder.buf.datagram_remaining_mut()
+                            > builder.predict_packet_end() + MAX_PADDING
                         {
                             trace!(
-                                "GSO truncated by demand for {} padding bytes or loss probe",
-                                builder.buf.segment_size() - packet_len_unpadded
+                                "GSO truncated by demand for {} padding bytes",
+                                builder.buf.datagram_remaining_mut() - builder.predict_packet_end()
                             );
-                            builder.finish_and_track(now, self, path_id, sent_frames);
+                            builder.finish_and_track(now, self, path_id, sent_frames, pad_datagram);
                             break;
                         }
 
@@ -895,7 +892,7 @@ impl Connection {
                         builder.pad_to(builder.buf.segment_size() as u16);
                     }
 
-                    builder.finish_and_track(now, self, path_id, sent_frames);
+                    builder.finish_and_track(now, self, path_id, sent_frames, pad_datagram);
 
                     if buf.num_datagrams() == 1 {
                         buf.clip_datagram_size();
@@ -924,10 +921,7 @@ impl Connection {
                 }
             } else {
                 // Nothing more to send.  This was the last packet.
-                if pad_datagram {
-                    builder.pad_to(MIN_INITIAL_SIZE);
-                }
-                builder.finish_and_track(now, self, path_id, sent_frames);
+                builder.finish_and_track(now, self, path_id, sent_frames, pad_datagram);
                 break;
             }
         }
@@ -987,7 +981,7 @@ impl Connection {
                 non_retransmits: true,
                 ..Default::default()
             };
-            builder.finish_and_track(now, self, path_id, sent_frames);
+            builder.finish_and_track(now, self, path_id, sent_frames, false);
 
             self.stats.path.sent_plpmtud_probes += 1;
 

quinn-proto/src/connection/packet_builder.rs

@@ -4,7 +4,7 @@ use tracing::{trace, trace_span};
 
 use super::{Connection, PathId, SentFrames, TransmitBuf, spaces::SentPacket};
 use crate::{
-    ConnectionId, Instant, TransportError, TransportErrorCode,
+    ConnectionId, Instant, MIN_INITIAL_SIZE, TransportError, TransportErrorCode,
     connection::ConnectionSide,
     frame::{self, Close},
     packet::{FIXED_BIT, Header, InitialHeader, LongType, PacketNumber, PartialEncode, SpaceId},
@@ -197,12 +197,16 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
     }
 
     pub(super) fn finish_and_track(
-        self,
+        mut self,
         now: Instant,
         conn: &mut Connection,
         path_id: PathId,
         sent: SentFrames,
+        pad_datagram: bool,
     ) {
+        if pad_datagram {
+            self.pad_to(MIN_INITIAL_SIZE);
+        }
         let ack_eliciting = self.ack_eliciting;
         let exact_number = self.exact_number;
         let space_id = self.space;
@@ -283,15 +287,17 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
             Some((self.exact_number, self.path, packet_crypto)),
         );
 
-        (self.buf.len() - encode_start, pad)
+        let packet_len = self.buf.len() - encode_start;
+        trace!(size = %packet_len, short_header = %self.short_header, "wrote packet");
+        (packet_len, pad)
     }
 
-    /// Predicts the size of the packet if it were finished now without additional padding
+    /// The number of additional bytes the current packet would take up if it was finished now
     ///
     /// This will include any padding which is required to make the size large enough to be
     /// encrypted correctly.
-    pub(super) fn predict_packet_size(&self) -> usize {
-        self.buf.len().max(self.min_size) - self.buf.datagram_start_offset() + self.tag_len
+    pub(super) fn predict_packet_end(&self) -> usize {
+        self.buf.len().max(self.min_size) + self.tag_len - self.buf.len()
     }
 
     /// Returns the remaining space in the packet that can be taken up by QUIC frames

quinn-proto/src/connection/transmit_buf.rs

@@ -1,4 +1,5 @@
 use bytes::BufMut;
+use tracing::trace;
 
 use crate::packet::BufLen;
 
@@ -138,6 +139,13 @@ impl<'a> TransmitBuf<'a> {
     /// the last datagram in a batch.
     pub(super) fn clip_datagram_size(&mut self) {
         debug_assert_eq!(self.num_datagrams, 1);
+        if self.buf.len() < self.segment_size {
+            trace!(
+                segment_size = self.buf.len(),
+                prev_segment_size = self.segment_size,
+                "clipped datagram size"
+            );
+        }
         self.segment_size = self.buf.len();
         self.buf_capacity = self.buf.len();
     }
@@ -147,6 +155,11 @@ impl<'a> TransmitBuf<'a> {
     /// This is also the maximum size datagrams are allowed to be. The first and last
     /// datagram in a batch are allowed to be smaller however. After the first datagram the
     /// segment size is clipped to the size of the first datagram.
+    ///
+    /// If the last datagram was created using [`TransmitBuf::start_new_datagram_with_size`]
+    /// the the segment size will be greater than the current datagram is allowed to be.
+    /// Thus [`TransmitBuf::datagram_remaining_mut`] should be used if you need to know the
+    /// amount of data that can be written into the datagram.
     pub(super) fn segment_size(&self) -> usize {
         self.segment_size
     }
@@ -178,6 +191,11 @@ impl<'a> TransmitBuf<'a> {
         self.buf_capacity
     }
 
+    /// Returns the number of bytes that may still be written into this datagram
+    pub(super) fn datagram_remaining_mut(&self) -> usize {
+        self.buf_capacity.saturating_sub(self.buf.len())
+    }
+
     /// Returns `true` if the buffer did not have anything written into it
     pub(super) fn is_empty(&self) -> bool {
         self.len() == 0