feat: use BytesMut for Transmit content

quinn-proto/src/connection/datagrams.rs

@@ -1,6 +1,6 @@
 use std::collections::VecDeque;
 
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use thiserror::Error;
 use tracing::{debug, trace};
 
@@ -127,7 +127,7 @@ impl DatagramState {
         Ok(was_empty)
     }
 
-    pub(super) fn write(&mut self, buf: &mut Vec<u8>, max_size: usize) -> bool {
+    pub(super) fn write(&mut self, buf: &mut BytesMut, max_size: usize) -> bool {
         let datagram = match self.outgoing.pop_front() {
             Some(x) => x,
             None => return false,

quinn-proto/src/connection/mod.rs

@@ -439,7 +439,7 @@ impl Connection {
                     SpaceId::Data,
                     "PATH_CHALLENGE queued without 1-RTT keys"
                 );
-                let mut buf = Vec::with_capacity(self.path.current_mtu() as usize);
+                let mut buf = BytesMut::with_capacity(self.path.current_mtu() as usize);
                 let buf_capacity = self.path.current_mtu() as usize;
 
                 let mut builder = PacketBuilder::new(
@@ -469,7 +469,7 @@ impl Connection {
                 self.stats.udp_tx.bytes += buf.len() as u64;
                 return Some(Transmit {
                     destination,
-                    contents: buf,
+                    contents: buf.freeze(),
                     ecn: None,
                     segment_size: None,
                     src_ip: self.local_ip,
@@ -500,7 +500,7 @@ impl Connection {
             _ => false,
         };
 
-        let mut buf = Vec::new();
+        let mut buf = BytesMut::new();
         // 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.
@@ -836,7 +836,7 @@ impl Connection {
 
         Some(Transmit {
             destination: self.path.remote,
-            contents: buf,
+            contents: buf.freeze(),
             ecn: if self.path.sending_ecn {
                 Some(EcnCodepoint::Ect0)
             } else {
@@ -2809,7 +2809,7 @@ impl Connection {
     fn populate_packet(
         &mut self,
         space_id: SpaceId,
-        buf: &mut Vec<u8>,
+        buf: &mut BytesMut,
         max_size: usize,
     ) -> SentFrames {
         let mut sent = SentFrames::default();
@@ -2981,7 +2981,7 @@ impl Connection {
         receiving_ecn: bool,
         sent: &mut SentFrames,
         space: &mut PacketSpace,
-        buf: &mut Vec<u8>,
+        buf: &mut BytesMut,
         stats: &mut ConnectionStats,
     ) {
         debug_assert!(!space.pending_acks.ranges().is_empty());

quinn-proto/src/connection/packet_builder.rs

@@ -1,6 +1,6 @@
 use std::time::Instant;
 
-use bytes::Bytes;
+use bytes::{Bytes, BytesMut};
 use rand::Rng;
 use tracing::{trace, trace_span};
 
@@ -32,7 +32,7 @@ impl PacketBuilder {
     pub(super) fn new(
         now: Instant,
         space_id: SpaceId,
-        buffer: &mut Vec<u8>,
+        buffer: &mut BytesMut,
         buffer_capacity: usize,
         datagram_start: usize,
         ack_eliciting: bool,
@@ -167,7 +167,7 @@ impl PacketBuilder {
         now: Instant,
         conn: &mut Connection,
         sent: Option<SentFrames>,
-        buffer: &mut Vec<u8>,
+        buffer: &mut BytesMut,
     ) {
         let ack_eliciting = self.ack_eliciting;
         let exact_number = self.exact_number;
@@ -210,7 +210,7 @@ impl PacketBuilder {
     }
 
     /// Encrypt packet, returning the length of the packet and whether padding was added
-    pub(super) fn finish(self, conn: &mut Connection, buffer: &mut Vec<u8>) -> (usize, bool) {
+    pub(super) fn finish(self, conn: &mut Connection, buffer: &mut BytesMut) -> (usize, bool) {
         let pad = buffer.len() < self.min_size;
         if pad {
             trace!("PADDING * {}", self.min_size - buffer.len());

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

@@ -4,7 +4,7 @@ use std::{
     mem,
 };
 
-use bytes::BufMut;
+use bytes::{BufMut, BytesMut};
 use rustc_hash::FxHashMap;
 use tracing::{debug, trace};
 
@@ -349,7 +349,7 @@ impl StreamsState {
 
     pub(in crate::connection) fn write_control_frames(
         &mut self,
-        buf: &mut Vec<u8>,
+        buf: &mut BytesMut,
         pending: &mut Retransmits,
         retransmits: &mut ThinRetransmits,
         stats: &mut FrameStats,
@@ -475,7 +475,7 @@ impl StreamsState {
 
     pub(crate) fn write_stream_frames(
         &mut self,
-        buf: &mut Vec<u8>,
+        buf: &mut BytesMut,
         max_buf_size: usize,
     ) -> StreamMetaVec {
         let mut stream_frames = StreamMetaVec::new();
@@ -874,7 +874,7 @@ mod tests {
         connection::State as ConnState, connection::Streams, ReadableError, RecvStream, SendStream,
         TransportErrorCode, WriteError,
     };
-    use bytes::Bytes;
+    use bytes::{Bytes, BytesMut};
 
     fn make(side: Side) -> StreamsState {
         StreamsState::new(
@@ -1266,7 +1266,7 @@ mod tests {
         high.set_priority(1).unwrap();
         high.write(b"high").unwrap();
 
-        let mut buf = Vec::with_capacity(40);
+        let mut buf = BytesMut::with_capacity(40);
         let meta = server.write_stream_frames(&mut buf, 40);
         assert_eq!(meta[0].id, id_high);
         assert_eq!(meta[1].id, id_mid);
@@ -1325,7 +1325,7 @@ mod tests {
         };
         high.set_priority(-1).unwrap();
 
-        let mut buf = Vec::with_capacity(1000);
+        let mut buf = BytesMut::with_capacity(1000);
         let meta = server.write_stream_frames(&mut buf, 40);
         assert_eq!(meta.len(), 1);
         assert_eq!(meta[0].id, id_high);

quinn-proto/src/endpoint.rs

@@ -177,7 +177,7 @@ impl Endpoint {
                 }
                 trace!("sending version negotiation");
                 // Negotiate versions
-                let mut buf = Vec::<u8>::new();
+                let mut buf = BytesMut::new();
                 Header::VersionNegotiate {
                     random: self.rng.gen::<u8>() | 0x40,
                     src_cid: dst_cid,
@@ -196,7 +196,7 @@ impl Endpoint {
                 self.transmits.push_back(Transmit {
                     destination: remote,
                     ecn: None,
-                    contents: buf,
+                    contents: buf.freeze(),
                     segment_size: None,
                     src_ip: local_ip,
                 });
@@ -340,15 +340,15 @@ impl Endpoint {
             "sending stateless reset for {} to {}",
             dst_cid, addresses.remote
         );
-        let mut buf = Vec::<u8>::new();
+        let mut buf = BytesMut::new();
         // Resets with at least this much padding can't possibly be distinguished from real packets
         const IDEAL_MIN_PADDING_LEN: usize = MIN_PADDING_LEN + MAX_CID_SIZE;
         let padding_len = if max_padding_len <= IDEAL_MIN_PADDING_LEN {
             max_padding_len
         } else {
             self.rng.gen_range(IDEAL_MIN_PADDING_LEN..max_padding_len)
         };
-        buf.reserve_exact(padding_len + RESET_TOKEN_SIZE);
+        buf.reserve(padding_len + RESET_TOKEN_SIZE);
         buf.resize(padding_len, 0);
         self.rng.fill_bytes(&mut buf[0..padding_len]);
         buf[0] = 0b0100_0000 | buf[0] >> 2;
@@ -359,7 +359,7 @@ impl Endpoint {
         self.transmits.push_back(Transmit {
             destination: addresses.remote,
             ecn: None,
-            contents: buf,
+            contents: buf.freeze(),
             segment_size: None,
             src_ip: addresses.local_ip,
         });
@@ -538,7 +538,7 @@ impl Endpoint {
                     version,
                 };
 
-                let mut buf = Vec::new();
+                let mut buf = BytesMut::new();
                 let encode = header.encode(&mut buf);
                 buf.put_slice(&token);
                 buf.extend_from_slice(&server_config.crypto.retry_tag(version, &dst_cid, &buf));
@@ -547,7 +547,7 @@ impl Endpoint {
                 self.transmits.push_back(Transmit {
                     destination: addresses.remote,
                     ecn: None,
-                    contents: buf,
+                    contents: buf.freeze(),
                     segment_size: None,
                     src_ip: addresses.local_ip,
                 });
@@ -689,7 +689,7 @@ impl Endpoint {
             version,
         };
 
-        let mut buf = Vec::<u8>::new();
+        let mut buf = BytesMut::new();
         let partial_encode = header.encode(&mut buf);
         let max_len =
             INITIAL_MTU as usize - partial_encode.header_len - crypto.packet.local.tag_len();
@@ -703,7 +703,7 @@ impl Endpoint {
         self.transmits.push_back(Transmit {
             destination: addresses.remote,
             ecn: None,
-            contents: buf,
+            contents: buf.freeze(),
             segment_size: None,
             src_ip: addresses.local_ip,
         })

quinn-proto/src/frame.rs

@@ -4,7 +4,7 @@ use std::{
     ops::{Range, RangeInclusive},
 };
 
-use bytes::{Buf, BufMut, Bytes};
+use bytes::{Buf, BufMut, Bytes, BytesMut};
 use tinyvec::TinyVec;
 
 use crate::{
@@ -853,13 +853,13 @@ impl FrameStruct for Datagram {
 }
 
 impl Datagram {
-    pub(crate) fn encode<W: BufMut>(&self, length: bool, out: &mut W) {
+    pub(crate) fn encode(&self, length: bool, out: &mut BytesMut) {
         out.write(Type(*DATAGRAM_TYS.start() | u64::from(length))); // 1 byte
         if length {
             // Safe to unwrap because we check length sanity before queueing datagrams
             out.write(VarInt::from_u64(self.data.len() as u64).unwrap()); // <= 8 bytes
         }
-        out.put_slice(&self.data);
+        out.extend_from_slice(&self.data);
     }
 
     pub(crate) fn size(&self, length: bool) -> usize {

quinn-proto/src/lib.rs

@@ -38,6 +38,7 @@ mod tests;
 pub mod transport_parameters;
 mod varint;
 
+use bytes::Bytes;
 pub use varint::{VarInt, VarIntBoundsExceeded};
 
 mod connection;
@@ -277,7 +278,7 @@ pub struct Transmit {
     /// Explicit congestion notification bits to set on the packet
     pub ecn: Option<EcnCodepoint>,
     /// Contents of the datagram
-    pub contents: Vec<u8>,
+    pub contents: Bytes,
     /// The segment size if this transmission contains multiple datagrams.
     /// This is `None` if the transmit only contains a single datagram
     pub segment_size: Option<usize>,

quinn-proto/src/packet.rs

@@ -269,7 +269,7 @@ pub(crate) enum Header {
 }
 
 impl Header {
-    pub(crate) fn encode(&self, w: &mut Vec<u8>) -> PartialEncode {
+    pub(crate) fn encode(&self, w: &mut BytesMut) -> PartialEncode {
         use self::Header::*;
         let start = w.len();
         match *self {
@@ -844,7 +844,7 @@ mod tests {
 
         let dcid = ConnectionId::new(&hex!("06b858ec6f80452b"));
         let client = initial_keys(Version::V1, &dcid, Side::Client);
-        let mut buf = Vec::new();
+        let mut buf = BytesMut::new();
         let header = Header::Initial {
             number: PacketNumber::U8(0),
             src_cid: ConnectionId::new(&[]),
@@ -875,7 +875,7 @@ mod tests {
 
         let server = initial_keys(Version::V1, &dcid, Side::Server);
         let supported_versions = DEFAULT_SUPPORTED_VERSIONS.to_vec();
-        let decode = PartialDecode::new(buf.as_slice().into(), 0, &supported_versions, false)
+        let decode = PartialDecode::new(buf, 0, &supported_versions, false)
             .unwrap()
             .0;
         let mut packet = decode.finish(Some(&*server.header.remote)).unwrap();