From f0566b836a5adff981387506b5f3194a684d8674 Mon Sep 17 00:00:00 2001 From: Adam Fiksen Date: Mon, 14 Sep 2026 16:16:10 -0700 Subject: [PATCH] Add configurable initial codec buffer capacities Expose opt-in client and server builder controls for the initial read, write, and HPACK decode scratch buffer capacities. Preserve existing defaults and allow every buffer to grow as needed; clamp the write buffer to the encoder minimum. --- src/client.rs | 59 ++++++++++++++++++++++++- src/codec/framed_read.rs | 40 ++++++++++++++++- src/codec/framed_write.rs | 68 ++++++++++++++++++++++++++++- src/codec/mod.rs | 90 ++++++++++++++++++++++++++++++++++++--- src/hpack/decoder.rs | 17 +++++++- src/server.rs | 59 ++++++++++++++++++++++++- 6 files changed, 319 insertions(+), 14 deletions(-) diff --git a/src/client.rs b/src/client.rs index 39c5765bf..0201098e6 100644 --- a/src/client.rs +++ b/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,6 +668,7 @@ 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), @@ -672,6 +676,34 @@ impl Builder { } } + /// 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)); + } +} diff --git a/src/codec/framed_read.rs b/src/codec/framed_read.rs index 1ec0e2c1d..e0bcb455e 100644 --- a/src/codec/framed_read.rs +++ b/src/codec/framed_read.rs @@ -64,6 +64,17 @@ impl FramedRead { FramedRead { inner, decoder } } + pub(super) fn with_hpack_buffer_capacity( + inner: InnerFramedRead, + buffer_capacity: usize, + ) -> FramedRead { + 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 FramedRead { 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, diff --git a/src/codec/framed_write.rs b/src/codec/framed_write.rs index efd1966cf..d1ca54b24 100644 --- a/src/codec/framed_write.rs +++ b/src/codec/framed_write.rs @@ -83,22 +83,27 @@ where B: Buf, { pub fn new(inner: T) -> FramedWrite { + Self::with_capacity(inner, DEFAULT_BUFFER_CAPACITY) + } + + pub(super) fn with_capacity(inner: T, capacity: usize) -> FramedWrite { 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 FramedWrite { 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 AsyncRead for FramedWrite { @@ -366,6 +376,60 @@ impl AsyncRead for FramedWrite { // We never project the Pin to `B`. impl Unpin for FramedWrite {} +#[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> { + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + 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::*; diff --git a/src/codec/mod.rs b/src/codec/mod.rs index d7aecee91..d8fecb545 100644 --- a/src/codec/mod.rs +++ b/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 { inner: FramedRead>, } +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct InitialBufferCapacities { + pub(crate) read: Option, + pub(crate) write: Option, + pub(crate) hpack_decode: Option, +} + impl Codec 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 Codec { /// Updates the max received frame size. /// diff --git a/src/hpack/decoder.rs b/src/hpack/decoder.rs index 8f10e2e4c..7f6840a8b 100644 --- a/src/hpack/decoder.rs +++ b/src/hpack/decoder.rs @@ -154,14 +154,23 @@ struct StringMarker { impl Decoder { /// Creates a new `Decoder` with all settings set to default values. pub fn new(size: usize) -> Decoder { + Self::with_buffer_capacity(size, 4096) + } + + pub(crate) fn with_buffer_capacity(size: usize, buffer_capacity: usize) -> Decoder { Decoder { max_size_update: None, last_max_update: size, table: Table::new(size), - buffer: BytesMut::with_capacity(4096), + buffer: BytesMut::with_capacity(buffer_capacity), } } + #[cfg(test)] + pub(crate) fn buffer_capacity(&self) -> usize { + self.buffer.capacity() + } + /// Queues a potential size update #[allow(dead_code)] pub fn queue_size_update(&mut self, size: usize) { @@ -839,6 +848,12 @@ pub fn get_static(idx: usize) -> Header { mod test { use super::*; + #[test] + fn decode_buffer_capacity_is_configurable_without_changing_the_default() { + assert_eq!(Decoder::new(4096).buffer.capacity(), 4096); + assert_eq!(Decoder::with_buffer_capacity(4096, 0).buffer.capacity(), 0); + } + #[test] fn test_peek_u8() { let b = 0xff; diff --git a/src/server.rs b/src/server.rs index 316d6a180..2896db9a6 100644 --- a/src/server.rs +++ b/src/server.rs @@ -115,7 +115,7 @@ //! [`SendStream`]: ../struct.SendStream.html //! [`TcpListener`]: https://docs.rs/tokio-core/0.1/tokio_core/net/struct.TcpListener.html -use crate::codec::{Codec, UserError}; +use crate::codec::{Codec, InitialBufferCapacities, UserError}; use crate::frame::{self, Pseudo, PushPromiseHeaderError, Reason, Settings, StreamId}; use crate::proto::{self, Config, Error, Prioritized}; use crate::{FlowControl, PingPong, RecvStream, SendStream}; @@ -253,6 +253,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 due to protocol error across /// the lifetime of the connection. /// @@ -386,7 +389,8 @@ where let entered = span.enter(); // 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); @@ -660,11 +664,40 @@ impl Builder { settings: Settings::default(), initial_target_connection_window_size: None, max_send_buffer_size: proto::DEFAULT_MAX_SEND_BUFFER_SIZE, + initial_buffer_capacities: InitialBufferCapacities::default(), 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. /// @@ -1788,3 +1821,25 @@ where } } } + +#[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)); + } +}