Add configurable initial codec buffer capacities

src/client.rs

@@ -135,7 +135,7 @@
 //! [`Builder`]: struct.Builder.html
 //! [`Error`]: ../struct.Error.html
 
-use crate::codec::{Codec, SendError, UserError};
+use crate::codec::{Codec, InitialBufferCapacities, SendError, UserError};
 use crate::ext::Protocol;
 use crate::frame::{Headers, Pseudo, Reason, Settings, StreamId};
 use crate::proto::{self, Error};
@@ -324,6 +324,9 @@ pub struct Builder {
     /// Maximum amount of bytes to "buffer" for writing per stream.
     max_send_buffer_size: usize,
 
+    /// Initial capacities for connection-level codec buffers.
+    initial_buffer_capacities: InitialBufferCapacities,
+
     /// Maximum number of locally reset streams to keep at a time.
     reset_stream_max: usize,
 
@@ -665,13 +668,42 @@ impl Builder {
             pending_accept_reset_stream_max: proto::DEFAULT_REMOTE_RESET_STREAM_MAX,
             initial_target_connection_window_size: None,
             initial_max_send_streams: usize::MAX,
+            initial_buffer_capacities: InitialBufferCapacities::default(),
             settings: Default::default(),
             stream_id: 1.into(),
             local_max_error_reset_streams: Some(proto::DEFAULT_LOCAL_RESET_COUNT_MAX),
             data_frame_budget: proto::DataFrameBudget::Auto,
         }
     }
 
+    /// Sets the initial capacity of the connection read buffer.
+    ///
+    /// The buffer grows as needed. If this is not set, the default capacity is
+    /// preserved.
+    pub fn initial_read_buffer_capacity(&mut self, capacity: usize) -> &mut Self {
+        self.initial_buffer_capacities.read = Some(capacity);
+        self
+    }
+
+    /// Sets the initial capacity of the connection write buffer.
+    ///
+    /// The buffer grows as needed and is clamped to the minimum required by
+    /// the frame encoder. If this is not set, the default capacity is
+    /// preserved.
+    pub fn initial_write_buffer_capacity(&mut self, capacity: usize) -> &mut Self {
+        self.initial_buffer_capacities.write = Some(capacity);
+        self
+    }
+
+    /// Sets the initial capacity of the HPACK decoder scratch buffer.
+    ///
+    /// The buffer grows as needed. If this is not set, the default capacity is
+    /// preserved.
+    pub fn initial_hpack_decode_buffer_capacity(&mut self, capacity: usize) -> &mut Self {
+        self.initial_buffer_capacities.hpack_decode = Some(capacity);
+        self
+    }
+
     /// Indicates the initial window size (in octets) for stream-level
     /// flow control for received data.
     ///
@@ -1338,7 +1370,8 @@ where
         bind_connection(&mut io).await?;
 
         // Create the codec
-        let mut codec = Codec::new(io);
+        let mut codec =
+            Codec::with_initial_buffer_capacities(io, builder.initial_buffer_capacities);
 
         if let Some(max) = builder.settings.max_frame_size() {
             codec.set_max_recv_frame_size(max as usize);
@@ -1748,3 +1781,25 @@ impl proto::Peer for Peer {
         Ok(response)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn codec_buffer_capacities_are_opt_in() {
+        let default = Builder::new();
+        assert_eq!(default.initial_buffer_capacities.read, None);
+        assert_eq!(default.initial_buffer_capacities.write, None);
+        assert_eq!(default.initial_buffer_capacities.hpack_decode, None);
+
+        let mut configured = Builder::new();
+        configured
+            .initial_read_buffer_capacity(4 * 1024)
+            .initial_write_buffer_capacity(4 * 1024)
+            .initial_hpack_decode_buffer_capacity(0);
+        assert_eq!(configured.initial_buffer_capacities.read, Some(4 * 1024));
+        assert_eq!(configured.initial_buffer_capacities.write, Some(4 * 1024));
+        assert_eq!(configured.initial_buffer_capacities.hpack_decode, Some(0));
+    }
+}

src/codec/framed_read.rs

@@ -64,6 +64,17 @@ impl<T> FramedRead<T> {
         FramedRead { inner, decoder }
     }
 
+    pub(super) fn with_hpack_buffer_capacity(
+        inner: InnerFramedRead<T, LengthDelimitedCodec>,
+        buffer_capacity: usize,
+    ) -> FramedRead<T> {
+        let decoder = FrameDecoder::with_hpack_buffer_capacity(
+            inner.decoder().max_frame_length(),
+            buffer_capacity,
+        );
+        FramedRead { inner, decoder }
+    }
+
     pub fn get_ref(&self) -> &T {
         self.inner.get_ref()
     }
@@ -72,6 +83,16 @@ impl<T> FramedRead<T> {
         self.inner.get_mut()
     }
 
+    #[cfg(test)]
+    pub(super) fn read_buffer_capacity(&self) -> usize {
+        self.inner.read_buffer().capacity()
+    }
+
+    #[cfg(test)]
+    pub(super) fn hpack_buffer_capacity(&self) -> usize {
+        self.decoder.hpack.buffer_capacity()
+    }
+
     /// Returns the current max frame size setting
     #[inline]
     pub fn max_frame_size(&self) -> usize {
@@ -114,9 +135,26 @@ fn calc_max_continuation_frames(header_max: usize, frame_max: usize) -> usize {
 
 impl FrameDecoder {
     fn new(max_frame_size: usize) -> Self {
+        Self::with_hpack_decoder(
+            max_frame_size,
+            hpack::Decoder::new(DEFAULT_SETTINGS_HEADER_TABLE_SIZE),
+        )
+    }
+
+    fn with_hpack_buffer_capacity(max_frame_size: usize, buffer_capacity: usize) -> Self {
+        Self::with_hpack_decoder(
+            max_frame_size,
+            hpack::Decoder::with_buffer_capacity(
+                DEFAULT_SETTINGS_HEADER_TABLE_SIZE,
+                buffer_capacity,
+            ),
+        )
+    }
+
+    fn with_hpack_decoder(max_frame_size: usize, hpack: hpack::Decoder) -> Self {
         let max_header_list_size = DEFAULT_SETTINGS_MAX_HEADER_LIST_SIZE;
         FrameDecoder {
-            hpack: hpack::Decoder::new(DEFAULT_SETTINGS_HEADER_TABLE_SIZE),
+            hpack,
             max_header_list_size,
             max_continuation_frames: calc_max_continuation_frames(
                 max_header_list_size,

src/codec/framed_write.rs

@@ -83,22 +83,27 @@ where
     B: Buf,
 {
     pub fn new(inner: T) -> FramedWrite<T, B> {
+        Self::with_capacity(inner, DEFAULT_BUFFER_CAPACITY)
+    }
+
+    pub(super) fn with_capacity(inner: T, capacity: usize) -> FramedWrite<T, B> {
         let chain_threshold = if inner.is_write_vectored() {
             CHAIN_THRESHOLD
         } else {
             CHAIN_THRESHOLD_WITHOUT_VECTORED_IO
         };
+        let min_buffer_capacity = chain_threshold + frame::HEADER_LEN;
         FramedWrite {
             inner,
             final_flush_done: false,
             encoder: Encoder {
                 hpack: hpack::Encoder::default(),
-                buf: Cursor::new(BytesMut::with_capacity(DEFAULT_BUFFER_CAPACITY)),
+                buf: Cursor::new(BytesMut::with_capacity(capacity.max(min_buffer_capacity))),
                 next: None,
                 last_data_frame: None,
                 max_frame_size: frame::DEFAULT_MAX_FRAME_SIZE,
                 chain_threshold,
-                min_buffer_capacity: chain_threshold + frame::HEADER_LEN,
+                min_buffer_capacity,
             },
         }
     }
@@ -351,6 +356,11 @@ impl<T, B> FramedWrite<T, B> {
     pub fn get_mut(&mut self) -> &mut T {
         &mut self.inner
     }
+
+    #[cfg(test)]
+    pub(super) fn write_buffer_capacity(&self) -> usize {
+        self.encoder.buf.get_ref().capacity()
+    }
 }
 
 impl<T: AsyncRead + Unpin, B> AsyncRead for FramedWrite<T, B> {
@@ -366,6 +376,60 @@ impl<T: AsyncRead + Unpin, B> AsyncRead for FramedWrite<T, B> {
 // We never project the Pin to `B`.
 impl<T: Unpin, B> Unpin for FramedWrite<T, B> {}
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytes::Bytes;
+
+    struct TestWriter {
+        vectored: bool,
+    }
+
+    impl AsyncWrite for TestWriter {
+        fn poll_write(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+            buf: &[u8],
+        ) -> Poll<io::Result<usize>> {
+            Poll::Ready(Ok(buf.len()))
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn is_write_vectored(&self) -> bool {
+            self.vectored
+        }
+    }
+
+    #[test]
+    fn default_capacity_is_unchanged() {
+        let framed: FramedWrite<_, Bytes> = FramedWrite::new(TestWriter { vectored: true });
+        assert_eq!(framed.encoder.buf.get_ref().capacity(), DEFAULT_BUFFER_CAPACITY);
+    }
+
+    #[test]
+    fn requested_capacity_is_clamped_to_encoder_minimum() {
+        for (vectored, expected) in [
+            (true, CHAIN_THRESHOLD + frame::HEADER_LEN),
+            (
+                false,
+                CHAIN_THRESHOLD_WITHOUT_VECTORED_IO + frame::HEADER_LEN,
+            ),
+        ] {
+            let framed: FramedWrite<_, Bytes> =
+                FramedWrite::with_capacity(TestWriter { vectored }, 0);
+            assert_eq!(framed.encoder.buf.get_ref().capacity(), expected);
+            assert_eq!(framed.encoder.min_buffer_capacity, expected);
+        }
+    }
+}
+
 #[cfg(feature = "unstable")]
 mod unstable {
     use super::*;

src/codec/mod.rs

@@ -17,6 +17,7 @@ use std::pin::Pin;
 use std::task::{Context, Poll};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_util::codec::length_delimited;
+use tokio_util::codec::FramedRead as LengthDelimitedFramedRead;
 
 use std::io;
 
@@ -25,6 +26,13 @@ pub struct Codec<T, B> {
     inner: FramedRead<FramedWrite<T, B>>,
 }
 
+#[derive(Clone, Copy, Debug, Default)]
+pub(crate) struct InitialBufferCapacities {
+    pub(crate) read: Option<usize>,
+    pub(crate) write: Option<usize>,
+    pub(crate) hpack_decode: Option<usize>,
+}
+
 impl<T, B> Codec<T, B>
 where
     T: AsyncRead + AsyncWrite + Unpin,
@@ -38,18 +46,55 @@ where
 
     /// Returns a new `Codec` with the given maximum frame size
     pub fn with_max_recv_frame_size(io: T, max_frame_size: usize) -> Self {
+        Self::with_max_recv_frame_size_and_initial_buffer_capacities(
+            io,
+            max_frame_size,
+            InitialBufferCapacities::default(),
+        )
+    }
+
+    pub(crate) fn with_initial_buffer_capacities(
+        io: T,
+        capacities: InitialBufferCapacities,
+    ) -> Self {
+        Self::with_max_recv_frame_size_and_initial_buffer_capacities(
+            io,
+            frame::DEFAULT_MAX_FRAME_SIZE as usize,
+            capacities,
+        )
+    }
+
+    fn with_max_recv_frame_size_and_initial_buffer_capacities(
+        io: T,
+        max_frame_size: usize,
+        capacities: InitialBufferCapacities,
+    ) -> Self {
         // Wrap with writer
-        let framed_write = FramedWrite::new(io);
+        let framed_write = match capacities.write {
+            Some(capacity) => FramedWrite::with_capacity(io, capacity),
+            None => FramedWrite::new(io),
+        };
 
         // Delimit the frames
-        let delimited = length_delimited::Builder::new()
+        let mut builder = length_delimited::Builder::new();
+        builder
             .big_endian()
             .length_field_length(3)
             .length_adjustment(9)
-            .num_skip(0) // Don't skip the header
-            .new_read(framed_write);
-
-        let mut inner = FramedRead::new(delimited);
+            .num_skip(0); // Don't skip the header
+        let delimited = match capacities.read {
+            Some(capacity) => LengthDelimitedFramedRead::with_capacity(
+                framed_write,
+                builder.new_codec(),
+                capacity,
+            ),
+            None => builder.new_read(framed_write),
+        };
+
+        let mut inner = match capacities.hpack_decode {
+            Some(capacity) => FramedRead::with_hpack_buffer_capacity(delimited, capacity),
+            None => FramedRead::new(delimited),
+        };
 
         // Use FramedRead's method since it checks the value is within range.
         inner.set_max_frame_size(max_frame_size);
@@ -58,6 +103,39 @@ where
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use bytes::Bytes;
+
+    #[test]
+    fn default_initial_buffer_capacities_are_unchanged() {
+        let (_peer, io) = tokio::io::duplex(64);
+        let codec: Codec<_, Bytes> = Codec::new(io);
+
+        assert_eq!(codec.inner.read_buffer_capacity(), 8 * 1024);
+        assert_eq!(codec.inner.get_ref().write_buffer_capacity(), 16 * 1024);
+        assert_eq!(codec.inner.hpack_buffer_capacity(), 4 * 1024);
+    }
+
+    #[test]
+    fn initial_buffer_capacities_reach_each_codec_layer() {
+        let (_peer, io) = tokio::io::duplex(64);
+        let codec: Codec<_, Bytes> = Codec::with_initial_buffer_capacities(
+            io,
+            InitialBufferCapacities {
+                read: Some(4 * 1024),
+                write: Some(4 * 1024),
+                hpack_decode: Some(0),
+            },
+        );
+
+        assert_eq!(codec.inner.read_buffer_capacity(), 4 * 1024);
+        assert_eq!(codec.inner.get_ref().write_buffer_capacity(), 4 * 1024);
+        assert_eq!(codec.inner.hpack_buffer_capacity(), 0);
+    }
+}
+
 impl<T, B> Codec<T, B> {
     /// Updates the max received frame size.
     ///