Add a Bytes array based send interface to quinn-proto

This change modifies quinn-proto to allow callers to submit a list
of one or more owned `Bytes` chunks for transmission instead of
pure byte slices. This will provide a more efficient zero-copy
interface for applications which already make use of owned
`Bytes` buffers.

The internals of quinn-proto have been refactored in order to keep
the ability to pass `&[u8]` buffers and defer the conversion into
`Bytes` as long as possible, in order to avoid unnecessary allocations
if no data can be stored due to flow control.

Author: Matthias Einwag

quinn-proto/src/bytes_source.rs

@@ -0,0 +1,231 @@
+//! Defines types which can be lazily converted into `Bytes` chunks
+
+use bytes::Bytes;
+
+/// A source of one or more buffers which can be converted into `Bytes` buffers on demand
+///
+/// The purpose of this data type is to defer conversion as long as possible,
+/// so that no heap allocation is required in case no data is writable.
+pub trait BytesSource {
+    // Returns the total amount of bytes to write
+    fn len(&self) -> usize;
+    /// Limits the source to the new size
+    fn limit(&mut self, new_len: usize);
+    /// The amount of bytes and chunks consumed from this source
+    fn consumed(&self) -> Written;
+    /// Returns the next chunk from the source of owned chunks.
+    ///
+    /// This method will consume parts of the source.
+    // Calling it will yield `Bytes` elements up to the configured `limit`.
+    fn pop_chunk(&mut self) -> Option<Bytes>;
+}
+
+/// Indicates how many bytes and chunks had been transferred in a write operation
+#[derive(Debug, Default, PartialEq, Eq, Clone, Copy)]
+pub struct Written {
+    /// The amount of bytes which had been written
+    pub bytes: usize,
+    /// The amount of full chunks which had been written
+    ///
+    /// If a chunk was only partially written, it will not be counted by this field.
+    pub chunks: usize,
+}
+
+/// A [`BytesSource`] implementation for `&'a mut [Bytes]`
+///
+/// The type allows to dequeue [`Bytes`] chunks from an array of chunks, up to
+/// a configured limit.
+pub struct BytesArray<'a> {
+    /// The wrapped slice of `Bytes`
+    chunks: &'a mut [Bytes],
+    /// The amount of bytes which are still available for consumption
+    len: usize,
+    /// The amount of bytes and chunks consumed from this source
+    consumed: Written,
+}
+
+impl<'a> BytesArray<'a> {
+    pub fn from_chunks(chunks: &'a mut [Bytes]) -> Self {
+        let len = chunks.iter().map(|chunk| chunk.len()).sum();
+        Self {
+            chunks,
+            len,
+            consumed: Written::default(),
+        }
+    }
+}
+
+impl<'a> BytesSource for BytesArray<'a> {
+    fn len(&self) -> usize {
+        self.len
+    }
+
+    fn limit(&mut self, new_len: usize) {
+        self.len = self.len.min(new_len);
+    }
+
+    fn consumed(&self) -> Written {
+        self.consumed
+    }
+
+    fn pop_chunk(&mut self) -> Option<Bytes> {
+        // The loop exists to skip empty chunks while still marking them as
+        // consumed
+        while self.consumed.chunks < self.chunks.len() {
+            let chunk = &mut self.chunks[self.consumed.chunks];
+
+            if chunk.len() <= self.len {
+                let chunk = std::mem::take(chunk);
+                self.consumed.chunks += 1;
+                if chunk.is_empty() {
+                    continue;
+                }
+                self.len -= chunk.len();
+                self.consumed.bytes += chunk.len();
+                return Some(chunk);
+            } else if self.len > 0 {
+                let chunk = chunk.split_to(self.len);
+                self.len = 0;
+                self.consumed.bytes += chunk.len();
+                return Some(chunk);
+            } else {
+                return None;
+            }
+        }
+
+        None
+    }
+}
+
+/// A [`BytesSource`] implementation for `&[u8]`
+///
+/// The type allows to dequeue a single [`Bytes`] chunk, which will be lazily
+/// created from a reference. This allows to defer the allocation until it is
+/// known how much data needs to be copied.
+pub struct ByteSlice<'a> {
+    /// The wrapped byte slice
+    data: &'a [u8],
+    /// The amount of bytes which are still available for consumption
+    len: usize,
+    /// The amount of bytes and chunks consumed from this source
+    consumed: Written,
+}
+
+impl<'a> ByteSlice<'a> {
+    pub fn from_slice(data: &'a [u8]) -> Self {
+        let len = data.len();
+        Self {
+            data,
+            len,
+            consumed: Written::default(),
+        }
+    }
+}
+
+impl<'a> BytesSource for ByteSlice<'a> {
+    fn len(&self) -> usize {
+        self.len
+    }
+
+    fn limit(&mut self, new_len: usize) {
+        self.len = self.len.min(new_len);
+    }
+
+    fn consumed(&self) -> Written {
+        self.consumed
+    }
+
+    fn pop_chunk(&mut self) -> Option<Bytes> {
+        if self.len == 0 {
+            return None;
+        }
+
+        let chunk = Bytes::from(self.data[..self.len].to_owned());
+        self.consumed.bytes += chunk.len();
+        if self.len == self.data.len() {
+            self.consumed.chunks += 1;
+        }
+        self.len = 0;
+        Some(chunk)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn bytes_array() {
+        let full = b"Hello World 123456789 ABCDEFGHJIJKLMNOPQRSTUVWXYZ".to_owned();
+        for limit in 0..full.len() {
+            let mut chunks = [
+                Bytes::from_static(b""),
+                Bytes::from_static(b"Hello "),
+                Bytes::from_static(b"Wo"),
+                Bytes::from_static(b""),
+                Bytes::from_static(b"r"),
+                Bytes::from_static(b"ld"),
+                Bytes::from_static(b""),
+                Bytes::from_static(b" 12345678"),
+                Bytes::from_static(b"9 ABCDE"),
+                Bytes::from_static(b"F"),
+                Bytes::from_static(b"GHJIJKLMNOPQRSTUVWXYZ"),
+            ];
+            let num_chunks = chunks.len();
+            let last_chunk_len = chunks[chunks.len() - 1].len();
+
+            let mut array = BytesArray::from_chunks(&mut chunks);
+            assert_eq!(array.len(), full.len());
+
+            array.limit(limit);
+            assert_eq!(array.len(), limit);
+
+            let mut buf = Vec::new();
+            let mut chunks_popped = 0;
+            while let Some(chunk) = array.pop_chunk() {
+                buf.extend_from_slice(&chunk);
+                chunks_popped += 1;
+            }
+
+            assert_eq!(&buf[..], &full[..limit]);
+            assert_eq!(array.consumed().bytes, limit);
+
+            if limit == full.len() {
+                // Full consumption of the last chunk
+                assert_eq!(array.consumed().chunks, num_chunks);
+                // Since there are empty chunks, we consume more than there are popped
+                assert_eq!(array.consumed().chunks, chunks_popped + 3);
+            } else if limit > full.len() - last_chunk_len {
+                // Partial consumption of the last chunk
+                assert_eq!(array.consumed().chunks, num_chunks - 1);
+                assert_eq!(array.consumed().chunks, chunks_popped + 2);
+            }
+        }
+    }
+
+    #[test]
+    fn byte_slice() {
+        let full = b"Hello World 123456789 ABCDEFGHJIJKLMNOPQRSTUVWXYZ".to_owned();
+        for limit in 0..full.len() {
+            let mut array = ByteSlice::from_slice(&full[..]);
+            assert_eq!(array.len(), full.len());
+
+            array.limit(limit);
+            assert_eq!(array.len(), limit);
+
+            let mut buf = Vec::new();
+            while let Some(chunk) = array.pop_chunk() {
+                buf.extend_from_slice(&chunk);
+            }
+
+            assert_eq!(&buf[..], &full[..limit]);
+            assert_eq!(array.consumed().bytes, limit);
+
+            if limit == full.len() {
+                assert_eq!(array.consumed().chunks, 1);
+            } else {
+                assert_eq!(array.consumed().chunks, 0);
+            }
+        }
+    }
+}

quinn-proto/src/connection/mod.rs

@@ -14,6 +14,7 @@ use thiserror::Error;
 use tracing::{debug, error, trace, trace_span, warn};
 
 use crate::{
+    bytes_source::{ByteSlice, BytesArray, BytesSource},
     cid_generator::ConnectionIdGenerator,
     cid_queue::CidQueue,
     coding::BufMutExt,
@@ -29,7 +30,7 @@ use crate::{
         EndpointEventInner, IssuedCid,
     },
     transport_parameters::TransportParameters,
-    Dir, Frame, Side, StreamId, Transmit, TransportError, TransportErrorCode, VarInt,
+    Dir, Frame, Side, StreamId, Transmit, TransportError, TransportErrorCode, VarInt, Written,
     MAX_STREAM_COUNT, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE, TIMER_GRANULARITY,
 };
 
@@ -948,6 +949,32 @@ where
     ///
     /// Returns the number of bytes successfully written.
     pub fn write(&mut self, stream: StreamId, data: &[u8]) -> Result<usize, WriteError> {
+        let mut source = ByteSlice::from_slice(data);
+        self.write_source(stream, &mut source)?;
+        Ok(source.consumed().bytes)
+    }
+
+    /// Send data on the given stream
+    ///
+    /// Returns the number of bytes and chunks successfully written.
+    /// Note that this method might also write a partial chunk. In this case
+    /// [`Written::chunks`] will not count this chunk as fully written. However
+    /// the chunk will be advanced and contain only non-written data after the call.
+    pub fn write_chunks(
+        &mut self,
+        stream: StreamId,
+        data: &mut [Bytes],
+    ) -> Result<Written, WriteError> {
+        let mut source = BytesArray::from_chunks(data);
+        self.write_source(stream, &mut source)?;
+        Ok(source.consumed())
+    }
+
+    fn write_source<B: BytesSource>(
+        &mut self,
+        stream: StreamId,
+        data: &mut B,
+    ) -> Result<(), WriteError> {
         assert!(stream.dir() == Dir::Bi || stream.initiator() == self.side);
         if self.state.is_closed() {
             trace!(%stream, "write blocked; connection draining");

quinn-proto/src/connection/streams.rs

@@ -12,6 +12,7 @@ use tracing::{debug, trace};
 use super::assembler::Chunk;
 use super::spaces::{Retransmits, ThinRetransmits};
 use crate::{
+    bytes_source::BytesSource,
     coding::BufMutExt,
     connection::stats::FrameStats,
     frame::{self, FrameStruct, StreamMetaVec},
@@ -254,7 +255,11 @@ impl Streams {
     }
 
     /// Queue `data` to be written for `stream`
-    pub fn write(&mut self, id: StreamId, data: &[u8]) -> Result<usize, WriteError> {
+    pub fn write<S: BytesSource>(
+        &mut self,
+        id: StreamId,
+        source: &mut S,
+    ) -> Result<(), WriteError> {
         let limit = (self.max_data - self.data_sent).min(self.send_window - self.unacked_data);
         let stream = self.send.get_mut(&id).ok_or(WriteError::UnknownStream)?;
         if limit == 0 {
@@ -267,15 +272,16 @@ impl Streams {
         }
 
         let was_pending = stream.is_pending();
-        let len = (data.len() as u64).min(limit) as usize;
-        let len = stream.write(&data[0..len])?;
-        self.data_sent += len as u64;
-        self.unacked_data += len as u64;
-        trace!(stream = %id, "wrote {} bytes", len);
+        source.limit(limit as usize);
+        stream.write(source)?;
+        let consumed = source.consumed().bytes as u64;
+        self.data_sent += consumed;
+        self.unacked_data += consumed;
+        trace!(stream = %id, "wrote {} bytes", consumed);
         if !was_pending {
             push_pending(&mut self.pending, id, stream.priority);
         }
-        Ok(len)
+        Ok(())
     }
 
     /// Process incoming stream frame
@@ -1068,7 +1074,7 @@ enum StreamHalf {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::TransportErrorCode;
+    use crate::{connection::ByteSlice, TransportErrorCode};
 
     fn make(side: Side) -> Streams {
         Streams::new(
@@ -1286,9 +1292,15 @@ mod tests {
         let id = server.open(Dir::Uni).unwrap();
         let reason = 0u32.into();
         server.received_stop_sending(id, reason);
-        assert_eq!(server.write(id, &[]), Err(WriteError::Stopped(reason)));
+        assert_eq!(
+            server.write(id, &mut ByteSlice::from_slice(&[])),
+            Err(WriteError::Stopped(reason))
+        );
         server.reset(id).unwrap();
-        assert_eq!(server.write(id, &[]), Err(WriteError::UnknownStream));
+        assert_eq!(
+            server.write(id, &mut ByteSlice::from_slice(&[])),
+            Err(WriteError::UnknownStream)
+        );
     }
 
     #[test]
@@ -1321,9 +1333,15 @@ mod tests {
         let id_low = server.open(Dir::Bi).unwrap();
         server.set_priority(id_low, -1).unwrap();
         server.set_priority(id_high, 1).unwrap();
-        server.write(id_mid, b"mid").unwrap();
-        server.write(id_low, b"low").unwrap();
-        server.write(id_high, b"high").unwrap();
+        server
+            .write(id_mid, &mut ByteSlice::from_slice(b"mid"))
+            .unwrap();
+        server
+            .write(id_low, &mut ByteSlice::from_slice(b"low"))
+            .unwrap();
+        server
+            .write(id_high, &mut ByteSlice::from_slice(b"high"))
+            .unwrap();
         let mut buf = Vec::with_capacity(40);
         let meta = server.write_stream_frames(&mut buf, 40);
         assert_eq!(meta[0].id, id_high);