Use DatagramBuffer in PacketBuilder

The PacketBuilder no longer has any business with the TransmitBuilder,
it can do everything with just the DatagramBuffer.

Author: flub

quinn-proto/src/connection/mod.rs

@@ -662,7 +662,12 @@ impl Connection {
                         builder.pad_to(transmit.segment_size() as u16);
                     }
 
-                    builder.finish_and_track(now, self, sent_frames.take(), &mut transmit);
+                    builder.finish_and_track(
+                        now,
+                        self,
+                        sent_frames.take(),
+                        &mut transmit.datagram_mut(),
+                    );
 
                     if transmit.num_datagrams() == 1 {
                         transmit.clip_datagram_size();
@@ -707,7 +712,12 @@ impl Connection {
                 // datagram.
                 // Finish current packet without adding extra padding
                 if let Some(builder) = builder_storage.take() {
-                    builder.finish_and_track(now, self, sent_frames.take(), &mut transmit);
+                    builder.finish_and_track(
+                        now,
+                        self,
+                        sent_frames.take(),
+                        &mut transmit.datagram_mut(),
+                    );
                 }
             }
 
@@ -738,7 +748,7 @@ impl Connection {
                 now,
                 space_id,
                 self.rem_cids.active(),
-                &mut transmit,
+                &mut transmit.datagram_mut(),
                 ack_eliciting,
                 self,
             )?);
@@ -835,7 +845,7 @@ impl Connection {
                             non_retransmits: true,
                             ..SentFrames::default()
                         }),
-                        &mut transmit,
+                        &mut transmit.datagram_mut(),
                     );
                     self.stats.udp_tx.on_sent(1, transmit.len());
                     return Some(Transmit {
@@ -888,7 +898,7 @@ impl Connection {
                 builder.pad_to(MIN_INITIAL_SIZE);
             }
             let last_packet_number = builder.exact_number;
-            builder.finish_and_track(now, self, sent_frames, &mut transmit);
+            builder.finish_and_track(now, self, sent_frames, &mut transmit.datagram_mut());
             self.path
                 .congestion
                 .on_sent(now, transmit.len() as u64, last_packet_number);
@@ -912,7 +922,7 @@ impl Connection {
                 now,
                 space_id,
                 self.rem_cids.active(),
-                &mut transmit,
+                &mut transmit.datagram_mut(),
                 true,
                 self,
             )?;
@@ -934,7 +944,7 @@ impl Connection {
                 non_retransmits: true,
                 ..Default::default()
             };
-            builder.finish_and_track(now, self, Some(sent_frames), &mut transmit);
+            builder.finish_and_track(now, self, Some(sent_frames), &mut transmit.datagram_mut());
 
             self.stats.path.sent_plpmtud_probes += 1;
 
@@ -1000,7 +1010,14 @@ impl Connection {
         // if a post-migration packet caused the CID to be retired, it's fair to pretend
         // this is sent first.
         debug_assert!(buf.datagram().is_empty());
-        let mut builder = PacketBuilder::new(now, SpaceId::Data, *prev_cid, buf, false, self)?;
+        let mut builder = PacketBuilder::new(
+            now,
+            SpaceId::Data,
+            *prev_cid,
+            &mut buf.datagram_mut(),
+            false,
+            self,
+        )?;
         trace!("validating previous path with PATH_CHALLENGE {:08x}", token);
         buf.datagram_mut().write(frame::FrameType::PATH_CHALLENGE);
         buf.datagram_mut().write(token);
@@ -1012,7 +1029,7 @@ impl Connection {
         // sending a datagram of this size
         builder.pad_to(MIN_INITIAL_SIZE);
 
-        builder.finish(self, buf);
+        builder.finish(self, &mut buf.datagram_mut());
         self.stats.udp_tx.on_sent(1, buf.len());
 
         Some(Transmit {

quinn-proto/src/connection/packet_builder.rs

@@ -2,7 +2,7 @@ use bytes::{BufMut, Bytes};
 use rand::Rng;
 use tracing::{trace, trace_span};
 
-use super::{Connection, SentFrames, TransmitBuilder, spaces::SentPacket};
+use super::{Connection, DatagramBuffer, SentFrames, spaces::SentPacket};
 use crate::{
     ConnectionId, Instant, TransportError, TransportErrorCode,
     connection::ConnectionSide,
@@ -36,7 +36,7 @@ impl PacketBuilder {
         now: Instant,
         space_id: SpaceId,
         dst_cid: ConnectionId,
-        buffer: &mut TransmitBuilder<'_>,
+        datagram: &mut DatagramBuffer<'_>,
         ack_eliciting: bool,
         conn: &mut Connection,
     ) -> Option<Self> {
@@ -120,9 +120,9 @@ impl PacketBuilder {
                 version,
             }),
         };
-        let partial_encode = header.encode(&mut buffer.datagram_mut());
+        let partial_encode = header.encode(datagram);
         if conn.peer_params.grease_quic_bit && conn.rng.random() {
-            buffer.datagram_mut()[partial_encode.start] ^= FIXED_BIT;
+            datagram[partial_encode.start] ^= FIXED_BIT;
         }
 
         let (sample_size, tag_len) = if let Some(ref crypto) = space.crypto {
@@ -146,10 +146,10 @@ impl PacketBuilder {
         // pn_len + payload_len + tag_len >= sample_size + 4
         // payload_len >= sample_size + 4 - pn_len - tag_len
         let min_size = Ord::max(
-            buffer.datagram().len() + (sample_size + 4).saturating_sub(number.len() + tag_len),
+            datagram.len() + (sample_size + 4).saturating_sub(number.len() + tag_len),
             partial_encode.start + dst_cid.len() + 6,
         );
-        let max_size = buffer.datagram_mut().capacity() - tag_len;
+        let max_size = datagram.capacity() - tag_len;
         debug_assert!(max_size >= min_size);
 
         Some(Self {
@@ -179,12 +179,12 @@ impl PacketBuilder {
         now: Instant,
         conn: &mut Connection,
         sent: Option<SentFrames>,
-        buffer: &mut TransmitBuilder<'_>,
+        datagram: &mut DatagramBuffer<'_>,
     ) {
         let ack_eliciting = self.ack_eliciting;
         let exact_number = self.exact_number;
         let space_id = self.space;
-        let (size, padded) = self.finish(conn, buffer);
+        let (size, padded) = self.finish(conn, datagram);
         let sent = match sent {
             Some(sent) => sent,
             None => return,
@@ -225,13 +225,13 @@ impl PacketBuilder {
     pub(super) fn finish(
         self,
         conn: &mut Connection,
-        buffer: &mut TransmitBuilder<'_>,
+        datagram: &mut DatagramBuffer<'_>,
     ) -> (usize, bool) {
-        let pad = self.min_size > buffer.datagram().len();
+        let pad = self.min_size > datagram.len();
         if pad {
-            let padding_bytes = self.min_size - buffer.datagram().len();
+            let padding_bytes = self.min_size - datagram.len();
             trace!("PADDING * {padding_bytes}");
-            buffer.datagram_mut().put_bytes(0, padding_bytes);
+            datagram.put_bytes(0, padding_bytes);
         }
 
         let space = &conn.spaces[self.space];
@@ -250,16 +250,15 @@ impl PacketBuilder {
             "Mismatching crypto tag len"
         );
 
-        buffer.datagram_mut().put_bytes(0, packet_crypto.tag_len());
+        datagram.put_bytes(0, packet_crypto.tag_len());
         let encode_start = self.partial_encode.start;
-        let mut datagram_buf = buffer.datagram_mut();
-        let packet_buf = &mut datagram_buf[encode_start..];
+        let packet_buf = &mut datagram[encode_start..];
         self.partial_encode.finish(
             packet_buf,
             header_crypto,
             Some((self.exact_number, packet_crypto)),
         );
 
-        (buffer.datagram().len() - encode_start, pad)
+        (datagram.len() - encode_start, pad)
     }
 }