Remove BufLen and switch to BufMut::remaining_mut

When writing frames into packets it needs to be known how much space
there is in the packets.  This used to be done using a max offset into
a larger buffer.  This now switches this round to use
BufMut::remaining_mut, which makes accessing this information easier
and also removes carrying this around as an extra parameter.

Author: flub

quinn-proto/src/connection/datagrams.rs

@@ -4,7 +4,7 @@ use bytes::{BufMut, Bytes};
 use thiserror::Error;
 use tracing::{debug, trace};
 
-use super::{BufLen, Connection};
+use super::Connection;
 use crate::{
     TransportError,
     frame::{Datagram, FrameStruct},
@@ -163,13 +163,13 @@ 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 (impl BufMut + BufLen), max_size: usize) -> bool {
+    pub(super) fn write(&mut self, buf: &mut impl BufMut) -> bool {
         let datagram = match self.outgoing.pop_front() {
             Some(x) => x,
             None => return false,
         };
 
-        if buf.len() + datagram.size(true) > max_size {
+        if datagram.size(true) > buf.remaining_mut() {
             // Future work: we could be more clever about cramming small datagrams into
             // mostly-full packets when a larger one is queued first
             self.outgoing.push_front(datagram);

quinn-proto/src/connection/mod.rs

@@ -848,13 +848,11 @@ impl Connection {
                 }
             }
 
-            let sent = self.populate_packet(
-                now,
-                space_id,
-                &mut transmit.datagram_mut(),
-                builder.max_size,
-                builder.exact_number,
-            );
+            let sent = {
+                let frame_space = builder.max_size - transmit.datagram_mut().len();
+                let mut buf = transmit.datagram_mut().limit(frame_space);
+                self.populate_packet(now, space_id, &mut buf, builder.exact_number)
+            };
 
             // ACK-only packets should only be sent when explicitly allowed. If we write them due to
             // any other reason, there is a bug which leads to one component announcing write
@@ -3040,8 +3038,7 @@ impl Connection {
         &mut self,
         now: Instant,
         space_id: SpaceId,
-        buf: &mut (impl BufMut + BufLen),
-        max_size: usize,
+        buf: &mut impl BufMut,
         pn: u64,
     ) -> SentFrames {
         let mut sent = SentFrames::default();
@@ -3117,7 +3114,7 @@ impl Connection {
         }
 
         // PATH_CHALLENGE
-        if buf.len() + 9 < max_size && space_id == SpaceId::Data {
+        if 9 < buf.remaining_mut() && space_id == SpaceId::Data {
             // Transmit challenges with every outgoing frame on an unvalidated path
             if let Some(token) = self.path.challenge {
                 // But only send a packet solely for that purpose at most once
@@ -3132,7 +3129,7 @@ impl Connection {
         }
 
         // PATH_RESPONSE
-        if buf.len() + 9 < max_size && space_id == SpaceId::Data {
+        if 9 < buf.remaining_mut() && space_id == SpaceId::Data {
             if let Some(token) = self.path_responses.pop_on_path(self.path.remote) {
                 sent.non_retransmits = true;
                 sent.requires_padding = true;
@@ -3144,7 +3141,7 @@ impl Connection {
         }
 
         // CRYPTO
-        while buf.len() + frame::Crypto::SIZE_BOUND < max_size && !is_0rtt {
+        while frame::Crypto::SIZE_BOUND < buf.remaining_mut() && !is_0rtt {
             let mut frame = match space.pending.crypto.pop_front() {
                 Some(x) => x,
                 None => break,
@@ -3154,8 +3151,7 @@ impl Connection {
             // Since the offset is known, we can reserve the exact size required to encode it.
             // For length we reserve 2bytes which allows to encode up to 2^14,
             // which is more than what fits into normally sized QUIC frames.
-            let max_crypto_data_size = max_size
-                - buf.len()
+            let max_crypto_data_size = buf.remaining_mut()
                 - 1 // Frame Type
                 - VarInt::size(unsafe { VarInt::from_u64_unchecked(frame.offset) })
                 - 2; // Maximum encoded length for frame size, given we send less than 2^14 bytes
@@ -3191,12 +3187,11 @@ impl Connection {
                 &mut space.pending,
                 &mut sent.retransmits,
                 &mut self.stats.frame_tx,
-                max_size,
             );
         }
 
         // NEW_CONNECTION_ID
-        while buf.len() + 44 < max_size {
+        while 44 < buf.remaining_mut() {
             let issued = match space.pending.new_cids.pop() {
                 Some(x) => x,
                 None => break,
@@ -3218,7 +3213,7 @@ impl Connection {
         }
 
         // RETIRE_CONNECTION_ID
-        while buf.len() + frame::RETIRE_CONNECTION_ID_SIZE_BOUND < max_size {
+        while frame::RETIRE_CONNECTION_ID_SIZE_BOUND < buf.remaining_mut() {
             let seq = match space.pending.retire_cids.pop() {
                 Some(x) => x,
                 None => break,
@@ -3232,8 +3227,8 @@ impl Connection {
 
         // DATAGRAM
         let mut sent_datagrams = false;
-        while buf.len() + Datagram::SIZE_BOUND < max_size && space_id == SpaceId::Data {
-            match self.datagrams.write(buf, max_size) {
+        while Datagram::SIZE_BOUND < buf.remaining_mut() && space_id == SpaceId::Data {
+            match self.datagrams.write(buf) {
                 true => {
                     sent_datagrams = true;
                     sent.non_retransmits = true;
@@ -3273,7 +3268,7 @@ impl Connection {
                 token: token.encode(&*server_config.token_key).into(),
             };
 
-            if buf.len() + new_token.size() >= max_size {
+            if new_token.size() >= buf.remaining_mut() {
                 space.pending.new_tokens.push(remote_addr);
                 break;
             }
@@ -3288,9 +3283,9 @@ impl Connection {
 
         // STREAM
         if space_id == SpaceId::Data {
-            sent.stream_frames =
-                self.streams
-                    .write_stream_frames(buf, max_size, self.config.send_fairness);
+            sent.stream_frames = self
+                .streams
+                .write_stream_frames(buf, self.config.send_fairness);
             self.stats.frame_tx.stream += sent.stream_frames.len() as u64;
         }
 
@@ -3976,29 +3971,6 @@ 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()
-    }
-}
-
-impl BufLen for bytes::buf::Limit<&mut Vec<u8>> {
-    fn len(&self) -> usize {
-        self.get_ref().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::{BufLen, stats::FrameStats},
+    connection::stats::FrameStats,
     frame::{self, FrameStruct, StreamMetaVec},
     transport_parameters::TransportParameters,
 };
@@ -411,14 +411,13 @@ impl StreamsState {
 
     pub(in crate::connection) fn write_control_frames(
         &mut self,
-        buf: &mut (impl BufMut + BufLen),
+        buf: &mut impl BufMut,
         pending: &mut Retransmits,
         retransmits: &mut ThinRetransmits,
         stats: &mut FrameStats,
-        max_size: usize,
     ) {
         // RESET_STREAM
-        while buf.len() + frame::ResetStream::SIZE_BOUND < max_size {
+        while frame::ResetStream::SIZE_BOUND < buf.remaining_mut() {
             let (id, error_code) = match pending.reset_stream.pop() {
                 Some(x) => x,
                 None => break,
@@ -442,7 +441,7 @@ impl StreamsState {
         }
 
         // STOP_SENDING
-        while buf.len() + frame::StopSending::SIZE_BOUND < max_size {
+        while frame::StopSending::SIZE_BOUND < buf.remaining_mut() {
             let frame = match pending.stop_sending.pop() {
                 Some(x) => x,
                 None => break,
@@ -461,7 +460,7 @@ impl StreamsState {
         }
 
         // MAX_DATA
-        if pending.max_data && buf.len() + 9 < max_size {
+        if pending.max_data && 9 < buf.remaining_mut() {
             pending.max_data = false;
 
             // `local_max_data` can grow bigger than `VarInt`.
@@ -484,7 +483,7 @@ impl StreamsState {
         }
 
         // MAX_STREAM_DATA
-        while buf.len() + 17 < max_size {
+        while 17 < buf.remaining_mut() {
             let id = match pending.max_stream_data.iter().next() {
                 Some(x) => *x,
                 None => break,
@@ -516,7 +515,7 @@ impl StreamsState {
 
         // MAX_STREAMS
         for dir in Dir::iter() {
-            if !pending.max_stream_id[dir as usize] || buf.len() + 9 >= max_size {
+            if !pending.max_stream_id[dir as usize] || 9 >= buf.remaining_mut() {
                 continue;
             }
 
@@ -541,19 +540,11 @@ impl StreamsState {
 
     pub(crate) fn write_stream_frames(
         &mut self,
-        buf: &mut (impl BufMut + BufLen),
-        max_buf_size: usize,
+        buf: &mut impl BufMut,
         fair: bool,
     ) -> StreamMetaVec {
         let mut stream_frames = StreamMetaVec::new();
-        while buf.len() + frame::Stream::SIZE_BOUND < max_buf_size {
-            if max_buf_size
-                .checked_sub(buf.len() + frame::Stream::SIZE_BOUND)
-                .is_none()
-            {
-                break;
-            }
-
+        while frame::Stream::SIZE_BOUND < buf.remaining_mut() {
             // Pop the stream of the highest priority that currently has pending data
             // If the stream still has some pending data left after writing, it will be reinserted, otherwise not
             let Some(stream) = self.pending.pop() else {
@@ -577,7 +568,7 @@ impl StreamsState {
 
             // Now that we know the `StreamId`, we can better account for how many bytes
             // are required to encode it.
-            let max_buf_size = max_buf_size - buf.len() - 1 - VarInt::size(id.into());
+            let max_buf_size = buf.remaining_mut() - 1 - VarInt::size(id.into());
             let (offsets, encode_length) = stream.pending.poll_transmit(max_buf_size);
             let fin = offsets.end == stream.pending.offset()
                 && matches!(stream.state, SendState::DataSent { .. });
@@ -1379,8 +1370,8 @@ mod tests {
         high.set_priority(1).unwrap();
         high.write(b"high").unwrap();
 
-        let mut buf = Vec::with_capacity(40);
-        let meta = server.write_stream_frames(&mut buf, 40, true);
+        let buf = Vec::with_capacity(40);
+        let meta = server.write_stream_frames(&mut buf.limit(40), true);
         assert_eq!(meta[0].id, id_high);
         assert_eq!(meta[1].id, id_mid);
         assert_eq!(meta[2].id, id_low);
@@ -1438,16 +1429,18 @@ mod tests {
         };
         high.set_priority(-1).unwrap();
 
-        let mut buf = Vec::with_capacity(1000);
-        let meta = server.write_stream_frames(&mut buf, 40, true);
+        let buf = Vec::with_capacity(1000);
+        let mut buf = buf.limit(40);
+        let meta = server.write_stream_frames(&mut buf, true);
+        let buf = buf.into_inner();
         assert_eq!(meta.len(), 1);
         assert_eq!(meta[0].id, id_high);
 
         // After requeuing we should end up with 2 priorities - not 3
         assert_eq!(server.pending.len(), 2);
 
         // Send the remaining data. The initial mid priority one should go first now
-        let meta = server.write_stream_frames(&mut buf, 1000, true);
+        let meta = server.write_stream_frames(&mut buf.limit(1000), true);
         assert_eq!(meta.len(), 2);
         assert_eq!(meta[0].id, id_mid);
         assert_eq!(meta[1].id, id_high);
@@ -1507,8 +1500,9 @@ mod tests {
 
             // loop until all the streams are written
             loop {
-                let buf_len = buf.len();
-                let meta = server.write_stream_frames(&mut buf, buf_len + 40, fair);
+                let mut lbuf = buf.limit(40);
+                let meta = server.write_stream_frames(&mut lbuf, fair);
+                buf = lbuf.into_inner();
                 if meta.is_empty() {
                     break;
                 }
@@ -1575,11 +1569,12 @@ mod tests {
         stream_b.write(&[b'b'; 100]).unwrap();
 
         let mut metas = vec![];
-        let mut buf = Vec::with_capacity(1024);
+        let buf = Vec::with_capacity(1024);
 
         // Write the first chunk of stream_a
-        let buf_len = buf.len();
-        let meta = server.write_stream_frames(&mut buf, buf_len + 40, false);
+        let mut buf = buf.limit(40);
+        let meta = server.write_stream_frames(&mut buf, false);
+        let mut buf = buf.into_inner();
         assert!(!meta.is_empty());
         metas.extend(meta);
 
@@ -1595,8 +1590,9 @@ mod tests {
 
         // loop until all the streams are written
         loop {
-            let buf_len = buf.len();
-            let meta = server.write_stream_frames(&mut buf, buf_len + 40, false);
+            let mut lbuf = buf.limit(40);
+            let meta = server.write_stream_frames(&mut lbuf, false);
+            buf = lbuf.into_inner();
             if meta.is_empty() {
                 break;
             }

quinn-proto/src/connection/transmit_builder.rs

@@ -2,8 +2,6 @@ use std::ops::{Deref, DerefMut};
 
 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
@@ -254,10 +252,3 @@ impl DatagramBuffer<'_> {
         self.max_offset - self.start_offset
     }
 }
-
-// Temporary compatibility with the BufLen trait.  To be removed in follow-up commits.
-impl BufLen for DatagramBuffer<'_> {
-    fn len(&self) -> usize {
-        self.deref().len()
-    }
-}