refactor: Introduce explicit TransmitBuf::first_datagram

quinn-proto/src/connection/packet_builder.rs

@@ -25,7 +25,7 @@ pub(super) struct PacketBuilder<'a, 'b> {
     path: PathId,
     pub(super) partial_encode: PartialEncode,
     pub(super) ack_eliciting: bool,
-    pub(super) exact_number: u64,
+    pub(super) packet_number: u64,
     /// Is this packet allowed to be coalesced?
     pub(super) can_coalesce: bool,
     /// Smallest absolute position in the associated buffer that must be occupied by this packet's
@@ -93,11 +93,11 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
         }
 
         let space = &mut conn.spaces[space_id];
-        let exact_number = space.for_path(path_id).get_tx_number(&mut conn.rng);
-        let span = trace_span!("send", space = ?space_id, pn = exact_number, %path_id).entered();
+        let packet_number = space.for_path(path_id).get_tx_number(&mut conn.rng);
+        let span = trace_span!("send", space = ?space_id, pn = packet_number, %path_id).entered();
 
         let number = PacketNumber::new(
-            exact_number,
+            packet_number,
             space.for_path(path_id).largest_acked_packet.unwrap_or(0),
         );
         let header = match space_id {
@@ -171,7 +171,7 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
 
         qlog.header(
             &header,
-            Some(exact_number),
+            Some(packet_number),
             space_id,
             space_id == SpaceId::Data && conn.spaces[SpaceId::Data].crypto.is_none(),
             path_id,
@@ -182,7 +182,7 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
             space: space_id,
             path: path_id,
             partial_encode,
-            exact_number,
+            packet_number,
             can_coalesce: header.can_coalesce(),
             min_size,
             tag_len,
@@ -201,7 +201,7 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
             path: PathId::ZERO,
             partial_encode: PartialEncode::no_header(),
             ack_eliciting: true,
-            exact_number: 0,
+            packet_number: 0,
             can_coalesce: true,
             min_size: 0,
             tag_len: 0,
@@ -287,11 +287,11 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
         match pad_datagram {
             PadDatagram::No => (),
             PadDatagram::ToSize(size) => self.pad_to(size),
-            PadDatagram::ToSegmentSize => self.pad_to(self.buf.segment_size() as u16),
+            PadDatagram::ToMaxSize => self.pad_to(self.buf.max_datagram_size() as u16),
             PadDatagram::ToMinMtu => self.pad_to(MIN_INITIAL_SIZE),
         }
         let ack_eliciting = self.ack_eliciting;
-        let exact_number = self.exact_number;
+        let packet_number = self.packet_number;
         let space_id = self.space;
         let (size, padded, sent) = self.finish(conn, now);
 
@@ -311,7 +311,7 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
         };
 
         conn.paths.get_mut(&path_id).unwrap().data.sent(
-            exact_number,
+            packet_number,
             packet,
             conn.spaces[space_id].for_path(path_id),
         );
@@ -372,7 +372,7 @@ impl<'a, 'b> PacketBuilder<'a, 'b> {
         self.partial_encode.finish(
             packet_buf,
             header_crypto,
-            Some((self.exact_number, self.path, packet_crypto)),
+            Some((self.packet_number, self.path, packet_crypto)),
         );
 
         let packet_len = self.buf.len() - encode_start;
@@ -414,11 +414,11 @@ pub(super) enum PadDatagram {
     No,
     /// To a specific size
     ToSize(u16),
-    /// Pad to the current MTU/segment size
+    /// Pad to the maximum allowed size for this datagram.
     ///
-    /// For the first datagram in a transmit the MTU is the same as the
-    /// [`TransmitBuf::segment_size`].
-    ToSegmentSize,
+    /// Usually this will be the path MTU as initialised by
+    /// [`TransmitBuf::start_first_datagram`].
+    ToMaxSize,
     /// Pad to [`MIN_INITIAL_SIZE`], the minimal QUIC MTU of 1200 bytes
     ToMinMtu,
 }
@@ -437,12 +437,12 @@ impl std::ops::BitOr for PadDatagram {
             (Self::No, rhs) => rhs,
             (Self::ToSize(size), Self::No) => Self::ToSize(size),
             (Self::ToSize(a), Self::ToSize(b)) => Self::ToSize(a.max(b)),
-            (Self::ToSize(_), Self::ToSegmentSize) => Self::ToSegmentSize,
+            (Self::ToSize(_), Self::ToMaxSize) => Self::ToMaxSize,
             (Self::ToSize(_), Self::ToMinMtu) => Self::ToMinMtu,
-            (Self::ToSegmentSize, Self::No) => Self::ToSegmentSize,
-            (Self::ToSegmentSize, Self::ToSize(_)) => Self::ToSegmentSize,
-            (Self::ToSegmentSize, Self::ToSegmentSize) => Self::ToSegmentSize,
-            (Self::ToSegmentSize, Self::ToMinMtu) => Self::ToMinMtu,
+            (Self::ToMaxSize, Self::No) => Self::ToMaxSize,
+            (Self::ToMaxSize, Self::ToSize(_)) => Self::ToMaxSize,
+            (Self::ToMaxSize, Self::ToMaxSize) => Self::ToMaxSize,
+            (Self::ToMaxSize, Self::ToMinMtu) => Self::ToMinMtu,
             (Self::ToMinMtu, _) => Self::ToMinMtu,
         }
     }

quinn-proto/src/connection/streams/state.rs

@@ -565,8 +565,8 @@ impl StreamsState {
     #[cfg(test)]
     fn write_frames_for_test(&mut self, capacity: usize, fair: bool) -> frame::StreamMetaVec {
         let buf = &mut Vec::with_capacity(capacity);
-        let mut tbuf = crate::connection::TransmitBuf::new(buf, std::num::NonZeroUsize::MIN, 1_200);
-        tbuf.start_new_datagram_with_size(capacity);
+        let mut tbuf = crate::connection::TransmitBuf::new(buf, std::num::NonZeroUsize::MIN);
+        tbuf.start_first_datagram(capacity);
         let builder = &mut PacketBuilder::simple_data_buf(&mut tbuf);
         let stats = &mut FrameStats::default();
         self.write_stream_frames(builder, fair, stats);

quinn-proto/src/connection/transmit_buf.rs

@@ -35,138 +35,125 @@ pub(super) struct TransmitBuf<'a> {
     ///
     /// Note that when coalescing packets this might be before the start of the current
     /// packet.
-    datagram_start: usize,
+    datagram_start_offset: usize,
     /// The maximum offset allowed to be used for the current datagram in the buffer
     ///
     /// The first and last datagram in a batch are allowed to be smaller then the maximum
     /// size. All datagrams in between need to be exactly this size.
-    buf_capacity: usize,
+    datagram_max_offset: usize,
     /// The maximum number of datagrams allowed to write into [`TransmitBuf::buf`]
     max_datagrams: NonZeroUsize,
     /// The number of datagrams already (partially) written into the buffer
     ///
     /// Incremented by a call to [`TransmitBuf::start_new_datagram`].
     pub(super) num_datagrams: usize,
-    /// The segment size of this GSO batch
+    /// The segment size of this GSO batch, set once the second datagram is started.
     ///
     /// The segment size is the size of each datagram in the GSO batch, only the last
     /// datagram in the batch may be smaller.
     ///
-    /// For the first datagram this is set to the maximum size a datagram is allowed to be:
-    /// the current path MTU. After the first datagram is finished this is reduced to the
-    /// size of the first datagram and can no longer change.
-    segment_size: usize,
+    /// Only set once there is more than one datagram.
+    segment_size: Option<usize>,
 }
 
 impl<'a> TransmitBuf<'a> {
-    pub(super) fn new(buf: &'a mut Vec<u8>, max_datagrams: NonZeroUsize, mtu: usize) -> Self {
+    pub(super) fn new(buf: &'a mut Vec<u8>, max_datagrams: NonZeroUsize) -> Self {
         buf.clear();
         Self {
             buf,
-            datagram_start: 0,
-            buf_capacity: 0,
+            datagram_start_offset: 0,
+            datagram_max_offset: 0,
             max_datagrams,
             num_datagrams: 0,
-            segment_size: mtu,
+            segment_size: None,
         }
     }
 
-    /// Starts a datagram with a custom datagram size
+    /// Starts the first datagram in the GSO batch.
     ///
-    /// This is a specialized 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)
+    /// The size of the first datagram sets the segment size of the GSO batch.
+    pub(super) fn start_first_datagram(&mut self, max_size: usize) {
+        debug_assert_eq!(self.num_datagrams, 0, "No datagram can be stared yet");
+        debug_assert!(
+            self.buf.is_empty(),
+            "Buffer must be empty for first datagram"
+        );
+        self.datagram_max_offset = max_size;
+        self.num_datagrams = 1;
+        if self.datagram_max_offset > self.buf.capacity() {
+            // Reserve all remaining capacity right away.
+            let max_batch_capacity = max_size * self.max_datagrams.get();
+            self.buf
+                .reserve_exact(max_batch_capacity.saturating_sub(self.buf.capacity()));
+        }
     }
 
-    /// 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 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.
-        let max_capacity_hint = self.max_datagrams.get() * self.segment_size;
-        self.new_datagram_inner(self.segment_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.into());
-        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 = NonZeroUsize::MIN.saturating_add(self.num_datagrams);
-            }
-        }
-        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"
+    /// Starts a subsequent datagram in the GSO batch.
+    pub(super) fn start_datagram(&mut self) {
+        // Could be enforced with typestate, but that's probably also meh.
+        debug_assert!(
+            self.num_datagrams >= 1,
+            "Use start_first_datagram for first datagram"
         );
-        self.buf_capacity = self.datagram_start + datagram_size;
-        if self.buf_capacity > self.buf.capacity() {
-            self.buf
-                .reserve_exact(max_capacity_hint.saturating_sub(self.buf.capacity()));
+        debug_assert!(
+            self.buf.len() <= self.datagram_max_offset,
+            "Datagram exceeded max offset"
+        );
+        let segment_size = self
+            .segment_size
+            .get_or_insert_with(|| self.buf.len() - self.datagram_start_offset);
+        let segment_size = *segment_size;
+        if self.num_datagrams > 1 {
+            debug_assert_eq!(
+                self.buf.len(),
+                self.datagram_max_offset,
+                "Subsequent datagrams must be exactly the segment size"
+            );
         }
+
         self.num_datagrams += 1;
+        self.datagram_start_offset = self.buf.len();
+        self.datagram_max_offset = self.buf.len() + segment_size;
+        if self.datagram_max_offset > self.buf.capacity() {
+            // Reserve all remaining capacity right away.
+            let max_batch_capacity = segment_size * self.max_datagrams.get();
+            self.buf
+                .reserve_exact(max_batch_capacity.saturating_sub(self.buf.capacity()));
+        }
     }
 
-    /// Clips the datagram size to the current size
+    /// Mark the first datagram as completely written, setting its 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`].
+    /// maximum size it was allowed to be. 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) {
+    pub(super) fn end_first_datagram(&mut self) {
         debug_assert_eq!(self.num_datagrams, 1);
-        if self.buf.len() < self.segment_size {
+        if self.buf.len() < self.datagram_max_offset {
             trace!(
-                segment_size = self.buf.len(),
-                prev_segment_size = self.segment_size,
+                size = self.buf.len(),
+                max_size = self.datagram_max_offset,
                 "clipped datagram size"
             );
         }
-        self.segment_size = self.buf.len();
-        self.buf_capacity = self.buf.len();
+        self.datagram_max_offset = self.buf.len();
     }
 
-    /// Returns the GSO segment size
-    ///
-    /// 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 {
+    /// Returns the GSO segment size.
+    pub(super) fn segment_size(&self) -> Option<usize> {
         self.segment_size
     }
 
+    /// Returns the maximum size this datagram is allowed to be.
+    ///
+    /// Once a second datagram is started this is equivalent to the segment size.
+    pub(super) fn max_datagram_size(&self) -> usize {
+        self.datagram_max_offset - self.datagram_start_offset
+    }
+
     /// Returns the number of datagrams written into the buffer
     ///
     /// The last datagram is not necessarily finished yet.
@@ -183,20 +170,20 @@ impl<'a> TransmitBuf<'a> {
     ///
     /// In other words, this offset contains the first byte of the current datagram.
     pub(super) fn datagram_start_offset(&self) -> usize {
-        self.datagram_start
+        self.datagram_start_offset
     }
 
     /// Returns the maximum offset in the buffer allowed for the current datagram
     ///
     /// The first and last datagram in a batch are allowed to be smaller then the maximum
     /// size. All datagrams in between need to be exactly this size.
     pub(super) fn datagram_max_offset(&self) -> usize {
-        self.buf_capacity
+        self.datagram_max_offset
     }
 
     /// 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())
+        self.datagram_max_offset.saturating_sub(self.buf.len())
     }
 
     /// Returns `true` if the buffer did not have anything written into it
@@ -213,18 +200,6 @@ impl<'a> TransmitBuf<'a> {
     pub(super) fn as_mut_slice(&mut self) -> &mut [u8] {
         self.buf.as_mut_slice()
     }
-
-    /// Returns the underlying buffer and the GSO segment size, if any.
-    ///
-    /// Note that the GSO segment size is only defined if there is more than one segment.
-    pub(super) fn finish(self) -> (&'a mut Vec<u8>, Option<usize>) {
-        let gso_segment_size = if self.num_datagrams() > 1 {
-            Some(self.segment_size)
-        } else {
-            None
-        };
-        (self.buf, gso_segment_size)
-    }
 }
 
 unsafe impl BufMut for TransmitBuf<'_> {

quinn-proto/src/packet.rs

@@ -951,12 +951,12 @@ impl SpaceId {
 
     /// Returns the next higher packet space.
     ///
-    /// Keeps returning [`SpaceId::Data`] as the highest space.
-    pub(crate) fn next(&self) -> Self {
+    /// Returns `None` if at  [`SpaceId::Data`].
+    pub(crate) fn next(&self) -> Option<Self> {
         match self {
-            Self::Initial => Self::Handshake,
-            Self::Handshake => Self::Data,
-            Self::Data => Self::Data,
+            Self::Initial => Some(Self::Handshake),
+            Self::Handshake => Some(Self::Data),
+            Self::Data => None,
         }
     }
 }