Do not pass buffer around as a Vec

We want to hide the actual buffer implementation from most of these
locations.  Currently it is a TransmitBuf but it likely will become
something else in the future.  So make this generic.

Author: flub

quinn-proto/src/connection/datagrams.rs

@@ -1,10 +1,10 @@
 use std::collections::VecDeque;
 
-use bytes::Bytes;
+use bytes::{BufMut, Bytes};
 use thiserror::Error;
 use tracing::{debug, trace};
 
-use super::Connection;
+use super::{BufLen, Connection};
 use crate::{
     TransportError,
     frame::{Datagram, FrameStruct},
@@ -163,7 +163,7 @@ impl DatagramState {
     ///
     /// Returns whether a frame was written. At most `max_size` bytes will be written, including
     /// framing.
-    pub(super) fn write(&mut self, buf: &mut Vec<u8>, max_size: usize) -> bool {
+    pub(super) fn write(&mut self, buf: &mut (impl BufMut + BufLen), max_size: usize) -> bool {
         let datagram = match self.outgoing.pop_front() {
             Some(x) => x,
             None => return false,

quinn-proto/src/connection/mod.rs

@@ -7,7 +7,7 @@ use std::{
     sync::Arc,
 };
 
-use bytes::{Bytes, BytesMut};
+use bytes::{BufMut, Bytes, BytesMut};
 use frame::StreamMetaVec;
 use rand::{Rng, SeedableRng, rngs::StdRng};
 use thiserror::Error;
@@ -755,7 +755,7 @@ impl Connection {
                         self.receiving_ecn,
                         &mut SentFrames::default(),
                         &mut self.spaces[space_id],
-                        buf.buf,
+                        &mut buf,
                         &mut self.stats,
                     );
                 }
@@ -842,7 +842,7 @@ impl Connection {
             let sent = self.populate_packet(
                 now,
                 space_id,
-                buf.buf,
+                &mut buf,
                 builder.max_size,
                 builder.exact_number,
             );
@@ -3020,7 +3020,7 @@ impl Connection {
         &mut self,
         now: Instant,
         space_id: SpaceId,
-        buf: &mut Vec<u8>,
+        buf: &mut (impl BufMut + BufLen),
         max_size: usize,
         pn: u64,
     ) -> SentFrames {
@@ -3286,7 +3286,7 @@ impl Connection {
         receiving_ecn: bool,
         sent: &mut SentFrames,
         space: &mut PacketSpace,
-        buf: &mut Vec<u8>,
+        buf: &mut impl BufMut,
         stats: &mut ConnectionStats,
     ) {
         debug_assert!(!space.pending_acks.ranges().is_empty());
@@ -3963,6 +3963,23 @@ fn negotiate_max_idle_timeout(x: Option<VarInt>, y: Option<VarInt>) -> Option<Du
     }
 }
 
+/// A buffer that can tell how much has been written to it already
+///
+/// This is commonly used for when a buffer is passed and the user may not write past a
+/// given size.  It allows the user of such a buffer to know the current cursor position in
+/// the buffer.  The maximum write size is usually passed in the same unit as
+/// [`BufLen::len`]: bytes since the buffer start.
+pub(crate) trait BufLen {
+    /// Returns the number of bytes written into the buffer so far
+    fn len(&self) -> usize;
+}
+
+impl BufLen for Vec<u8> {
+    fn len(&self) -> usize {
+        self.len()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

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

@@ -15,7 +15,7 @@ use super::{
 use crate::{
     Dir, MAX_STREAM_COUNT, Side, StreamId, TransportError, VarInt,
     coding::BufMutExt,
-    connection::stats::FrameStats,
+    connection::{BufLen, stats::FrameStats},
     frame::{self, FrameStruct, StreamMetaVec},
     transport_parameters::TransportParameters,
 };
@@ -411,7 +411,7 @@ impl StreamsState {
 
     pub(in crate::connection) fn write_control_frames(
         &mut self,
-        buf: &mut Vec<u8>,
+        buf: &mut (impl BufMut + BufLen),
         pending: &mut Retransmits,
         retransmits: &mut ThinRetransmits,
         stats: &mut FrameStats,
@@ -541,7 +541,7 @@ impl StreamsState {
 
     pub(crate) fn write_stream_frames(
         &mut self,
-        buf: &mut Vec<u8>,
+        buf: &mut (impl BufMut + BufLen),
         max_buf_size: usize,
         fair: bool,
     ) -> StreamMetaVec {

quinn-proto/src/connection/transmit_buf.rs

@@ -1,5 +1,7 @@
 use bytes::BufMut;
 
+use super::BufLen;
+
 /// The buffer in which to write datagrams for [`Connection::poll_transmit`]
 ///
 /// The `poll_transmit` function writes zero or more datagrams to a buffer.  Multiple
@@ -200,3 +202,9 @@ unsafe impl BufMut for TransmitBuf<'_> {
         self.buf.chunk_mut()
     }
 }
+
+impl BufLen for TransmitBuf<'_> {
+    fn len(&self) -> usize {
+        self.len()
+    }
+}

quinn-proto/src/frame.rs

@@ -899,13 +899,13 @@ impl FrameStruct for Datagram {
 }
 
 impl Datagram {
-    pub(crate) fn encode(&self, length: bool, out: &mut Vec<u8>) {
+    pub(crate) fn encode(&self, length: bool, out: &mut impl BufMut) {
         out.write(FrameType(*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.extend_from_slice(&self.data);
+        out.put_slice(&self.data);
     }
 
     pub(crate) fn size(&self, length: bool) -> usize {

quinn-proto/src/packet.rs

@@ -6,6 +6,7 @@ use thiserror::Error;
 use crate::{
     ConnectionId,
     coding::{self, BufExt, BufMutExt},
+    connection::BufLen,
     crypto,
 };
 
@@ -281,7 +282,7 @@ pub(crate) enum Header {
 }
 
 impl Header {
-    pub(crate) fn encode(&self, w: &mut Vec<u8>) -> PartialEncode {
+    pub(crate) fn encode<W: BufMut + BufLen>(&self, w: &mut W) -> PartialEncode {
         use Header::*;
         let start = w.len();
         match *self {