feat(s2n-quic-core): add buffer reader

quic/s2n-quic-bench/src/buffer.rs

@@ -2,7 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use criterion::{black_box, BenchmarkId, Criterion, Throughput};
-use s2n_quic_core::{buffer::ReceiveBuffer, varint::VarInt};
+use s2n_quic_core::{
+    buffer::{reader::Chunk as _, writer, ReceiveBuffer},
+    varint::VarInt,
+};
 
 pub fn benchmarks(c: &mut Criterion) {
     let mut group = c.benchmark_group("buffer");
@@ -14,6 +17,7 @@ pub fn benchmarks(c: &mut Criterion) {
 
         group.bench_with_input(BenchmarkId::new("skip", size), &input, |b, _input| {
             let mut buffer = ReceiveBuffer::new();
+            let size = VarInt::try_from(size).unwrap();
             b.iter(move || {
                 buffer.skip(black_box(size)).unwrap();
             });
@@ -25,7 +29,8 @@ pub fn benchmarks(c: &mut Criterion) {
             let len = VarInt::new(input.len() as _).unwrap();
             b.iter(move || {
                 buffer.write_at(offset, input).unwrap();
-                buffer.copy_into_buf(&mut NoOpBuf);
+                // Avoid oversampling the `pop` implementation
+                buffer.force_copy_into(&mut writer::chunk::Discard).unwrap();
                 offset += len;
             });
         });
@@ -44,34 +49,11 @@ pub fn benchmarks(c: &mut Criterion) {
                     buffer.write_at(first_offset, input).unwrap();
                     let second_offset = offset;
                     buffer.write_at(second_offset, input).unwrap();
-                    buffer.copy_into_buf(&mut NoOpBuf);
+                    // Avoid oversampling the `pop` implementation
+                    buffer.force_copy_into(&mut writer::chunk::Discard).unwrap();
                     offset = first_offset + len;
                 });
             },
         );
     }
 }
-
-/// A BufMut implementation that doesn't actually copy data into it
-///
-/// This is used to avoid oversampling the `pop` implementation for
-/// `write_at` benchmarks.
-struct NoOpBuf;
-
-unsafe impl bytes::BufMut for NoOpBuf {
-    #[inline]
-    fn remaining_mut(&self) -> usize {
-        usize::MAX
-    }
-
-    #[inline]
-    unsafe fn advance_mut(&mut self, _cnt: usize) {}
-
-    #[inline]
-    fn put_slice(&mut self, _slice: &[u8]) {}
-
-    #[inline]
-    fn chunk_mut(&mut self) -> &mut bytes::buf::UninitSlice {
-        unimplemented!()
-    }
-}

quic/s2n-quic-core/src/buffer/duplex.rs

@@ -0,0 +1,13 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use super::{Error, Reader, Writer};
+use crate::varint::VarInt;
+
+mod split;
+
+pub use split::Split;
+
+pub trait Duplex: Reader + Writer {
+    fn skip(&mut self, len: VarInt, final_offset: Option<VarInt>) -> Result<(), Error>;
+}

quic/s2n-quic-core/src/buffer/duplex/split.rs

@@ -0,0 +1,70 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::buffer::{
+    reader::Reader,
+    writer::{Chunk, Writer},
+    Duplex, Error,
+};
+
+pub struct Split<'a, C: Chunk, D: Duplex<Error = core::convert::Infallible> + ?Sized> {
+    chunk: &'a mut C,
+    duplex: &'a mut D,
+}
+
+impl<'a, C: Chunk, D: Duplex<Error = core::convert::Infallible> + ?Sized> Split<'a, C, D> {
+    #[inline]
+    pub fn new(chunk: &'a mut C, duplex: &'a mut D) -> Self {
+        Self { chunk, duplex }
+    }
+}
+
+impl<'a, C: Chunk, D: Duplex<Error = core::convert::Infallible> + ?Sized> Writer
+    for Split<'a, C, D>
+{
+    #[inline]
+    fn copy_from<R: Reader>(&mut self, reader: &mut R) -> Result<(), Error<R::Error>> {
+        let initial_offset = reader.current_offset();
+        let final_offset = reader.final_offset();
+        let is_contiguous = initial_offset == self.duplex.current_offset();
+
+        {
+            // if the chunk specializes writing zero-copy Bytes/BytesMut, then just write to the
+            // receive buffer, since that's what it stores
+            let mut should_delegate = C::SPECIALIZES_BYTES || C::SPECIALIZES_BYTES_MUT;
+
+            // if this packet is non-contiguous, then delegate to the wrapped writer
+            should_delegate |= !is_contiguous;
+
+            // if the chunk doesn't have any remaining capacity, then delegate
+            should_delegate |= !self.chunk.has_remaining_capacity();
+
+            if should_delegate {
+                self.duplex.copy_from(reader)?;
+
+                if !self.duplex.buffer_is_empty() && self.chunk.has_remaining_capacity() {
+                    self.duplex
+                        .copy_into(self.chunk)
+                        .expect("duplex error is infallible");
+                }
+
+                return Ok(());
+            }
+        }
+
+        debug_assert!(self.chunk.has_remaining_capacity());
+
+        reader.copy_into(self.chunk)?;
+        let write_len = initial_offset - reader.current_offset();
+
+        self.duplex
+            .skip(write_len, final_offset)
+            .map_err(Error::mapped)?;
+
+        if !reader.buffer_is_empty() {
+            self.duplex.copy_from(reader)?;
+        }
+
+        Ok(())
+    }
+}

quic/s2n-quic-core/src/buffer/error.rs

@@ -0,0 +1,46 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum Error<Reader = core::convert::Infallible> {
+    /// An invalid data range was provided
+    OutOfRange,
+    /// The provided final size was invalid for the buffer's state
+    InvalidFin,
+    /// The provided reader failed
+    ReaderError(Reader),
+}
+
+impl<Reader> From<Reader> for Error<Reader> {
+    #[inline]
+    fn from(reader: Reader) -> Self {
+        Self::ReaderError(reader)
+    }
+}
+
+impl Error {
+    #[inline]
+    pub fn mapped<Reader>(error: Error) -> Error<Reader> {
+        match error {
+            Error::OutOfRange => Error::OutOfRange,
+            Error::InvalidFin => Error::InvalidFin,
+            Error::ReaderError(_) => unreachable!(),
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+impl<Reader: std::error::Error> std::error::Error for Error<Reader> {}
+
+impl<Reader: core::fmt::Display> core::fmt::Display for Error<Reader> {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        match self {
+            Self::OutOfRange => write!(f, "write extends out of the maximum possible offset"),
+            Self::InvalidFin => write!(
+                f,
+                "write modifies the final offset in a non-compliant manner"
+            ),
+            Self::ReaderError(reader) => write!(f, "the provided reader failed with: {reader}"),
+        }
+    }
+}

quic/s2n-quic-core/src/buffer/mod.rs

@@ -1,6 +1,14 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
+pub mod duplex;
+mod error;
+pub mod reader;
 mod receive_buffer;
+pub mod writer;
 
+pub use duplex::Duplex;
+pub use error::Error;
+pub use reader::Reader;
 pub use receive_buffer::*;
+pub use writer::Writer;

quic/s2n-quic-core/src/buffer/reader.rs

@@ -0,0 +1,63 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::varint::VarInt;
+
+pub mod chunk;
+mod empty;
+pub mod incremental;
+mod limit;
+mod slice;
+
+pub use chunk::Chunk;
+pub use empty::Empty;
+pub use incremental::Incremental;
+pub use limit::Limit;
+pub use slice::Slice;
+
+pub trait Reader: Chunk {
+    /// Returns the currently read offset for the stream
+    fn current_offset(&self) -> VarInt;
+
+    /// Returns the final offset for the stream
+    fn final_offset(&self) -> Option<VarInt>;
+
+    /// Returns `true` if the reader has the final offset buffered
+    #[inline]
+    fn has_buffered_fin(&self) -> bool {
+        self.final_offset().map_or(false, |fin| {
+            let buffered_end = self
+                .current_offset()
+                .as_u64()
+                .saturating_add(self.buffered_len() as u64);
+            fin == buffered_end
+        })
+    }
+
+    /// Returns `true` if the reader is finished producing data
+    #[inline]
+    fn is_consumed(&self) -> bool {
+        self.final_offset()
+            .map_or(false, |fin| fin == self.current_offset())
+    }
+
+    /// Limits the maximum offset that the caller can read from the reader
+    #[inline]
+    fn with_max_data(&mut self, max_data: VarInt) -> Limit<Self> {
+        let max_buffered_len = max_data.saturating_sub(self.current_offset());
+        let max_buffered_len = max_buffered_len.as_u64().min(self.buffered_len() as u64) as usize;
+        self.with_limit(max_buffered_len)
+    }
+
+    /// Limits the maximum amount of data that the caller can read from the reader
+    #[inline]
+    fn with_limit(&mut self, max_buffered_len: usize) -> Limit<Self> {
+        Limit::new(self, max_buffered_len)
+    }
+
+    /// Temporarily clears the buffer for the reader, while preserving the offsets
+    #[inline]
+    fn with_empty_buffer(&self) -> Empty<Self> {
+        Empty::new(self)
+    }
+}

quic/s2n-quic-core/src/buffer/reader/chunk.rs

@@ -0,0 +1,64 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+mod buf;
+mod bytes;
+mod full_copy;
+mod io_slice;
+mod slice;
+mod trailer;
+
+#[cfg(test)]
+mod tests;
+
+pub use buf::Buf;
+pub use full_copy::FullCopy;
+pub use io_slice::IoSlice;
+pub use trailer::Trailer;
+
+pub trait Chunk {
+    type Error;
+
+    /// Returns the length of the chunk
+    fn buffered_len(&self) -> usize;
+
+    /// Returns if the chunk is empty
+    #[inline]
+    fn buffer_is_empty(&self) -> bool {
+        self.buffered_len() == 0
+    }
+
+    /// Reads the current trailer for the chunk
+    fn read_trailer(&mut self, watermark: usize) -> Result<Trailer<'_>, Self::Error>;
+
+    /// Copies the chunk of bytes into `dest`, with a trailing set of bytes.
+    ///
+    /// Implementations should either fill the `dest` completely or exhaust the buffered data.
+    ///
+    /// The chunk may optionally return a `Trailer`, which can be used by the caller to defer
+    /// copying the trailing chunk until later.
+    fn partial_copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<Trailer<'_>, Self::Error>
+    where
+        Dest: crate::buffer::writer::Chunk;
+
+    /// Forces the entire chunk to be copied
+    ///
+    /// The returned `Trailer` will always be empty.
+    #[inline]
+    fn full_copy(&mut self) -> FullCopy<Self> {
+        FullCopy::new(self)
+    }
+
+    /// Copies the chunk of bytes into `dest`.
+    ///
+    /// Implementations should either fill the `dest` completely or exhaust the buffered data.
+    #[inline]
+    fn copy_into<Dest>(&mut self, dest: &mut Dest) -> Result<(), Self::Error>
+    where
+        Dest: crate::buffer::writer::Chunk,
+    {
+        let mut trailer = self.partial_copy_into(dest)?;
+        let _: Result<(), core::convert::Infallible> = trailer.copy_into(dest);
+        Ok(())
+    }
+}