diff --git a/tonic/benches/decode.rs b/tonic/benches/decode.rs index ecac25464..8f2a76bc6 100644 --- a/tonic/benches/decode.rs +++ b/tonic/benches/decode.rs @@ -163,6 +163,15 @@ bench!(message_count_1, 500, 505, 1); bench!(message_count_10, 500, 505, 10); bench!(message_count_20, 500, 505, 20); +// representative unary and streaming workloads +bench!(unary_1k_single_chunk, 1_000, 1_005, 1); +bench!(unary_10k_single_chunk, 10_000, 10_005, 1); +bench!(unary_64k_16k_chunks, 64 * 1024, 16 * 1024, 1); +bench!(unary_1m_16k_chunks, 1024 * 1024, 16 * 1024, 1); +bench!(streaming_1k_16k_chunks, 1_000, 16 * 1024, 100); +bench!(streaming_1k_100b_chunks, 1_000, 100, 100); +bench!(unary_1m_single_chunk, 1024 * 1024, 1024 * 1024 + 5, 1); + benchmark_group!(chunk_size, chunk_size_100, chunk_size_500, chunk_size_1005); benchmark_group!( @@ -179,4 +188,15 @@ benchmark_group!( message_count_20 ); -benchmark_main!(chunk_size, message_size, message_count); +benchmark_group!( + representative, + unary_1k_single_chunk, + unary_10k_single_chunk, + unary_64k_16k_chunks, + unary_1m_16k_chunks, + streaming_1k_16k_chunks, + streaming_1k_100b_chunks, + unary_1m_single_chunk +); + +benchmark_main!(chunk_size, message_size, message_count, representative); diff --git a/tonic/src/codec/decode.rs b/tonic/src/codec/decode.rs index f5d83cb68..474d827ad 100644 --- a/tonic/src/codec/decode.rs +++ b/tonic/src/codec/decode.rs @@ -220,7 +220,9 @@ impl StreamingInner { ))); } - self.buf.reserve(len); + if self.buf.capacity() < len { + self.buf.reserve(len - self.buf.len()); + } self.state = State::ReadBody { compression: compression_encoding, @@ -453,3 +455,82 @@ impl fmt::Debug for Streaming { #[cfg(test)] static_assertions::assert_impl_all!(Streaming<()>: Send, Sync); + +#[cfg(test)] +mod tests { + use super::*; + use bytes::Bytes; + use http_body::Frame; + use http_body_util::StreamBody; + + #[derive(Debug)] + struct BytesDecoder { + buffer_settings: BufferSettings, + } + + impl Decoder for BytesDecoder { + type Item = Bytes; + type Error = Status; + + fn decode(&mut self, src: &mut DecodeBuf<'_>) -> Result, Self::Error> { + Ok(Some(src.copy_to_bytes(src.remaining()))) + } + + fn buffer_settings(&self) -> BufferSettings { + self.buffer_settings + } + } + + fn encode_messages(messages: &[Bytes]) -> Bytes { + let mut payload = BytesMut::new(); + + for message in messages { + payload.put_u8(0); + payload.put_u32(message.len() as u32); + payload.put_slice(message); + } + + payload.freeze() + } + + async fn decode_messages( + messages: &[Bytes], + chunk_size: usize, + buffer_size: usize, + ) -> Vec { + let payload = encode_messages(messages); + let frames = payload + .chunks(chunk_size) + .map(|chunk| Ok::<_, Status>(Frame::data(Bytes::copy_from_slice(chunk)))) + .collect::>(); + let body = StreamBody::new(tokio_stream::iter(frames)); + let decoder = BytesDecoder { + buffer_settings: BufferSettings::new(buffer_size, 32 * 1024), + }; + let mut stream = Streaming::new_request(decoder, body, None, None); + let mut decoded = Vec::new(); + + while let Some(message) = stream.message().await.unwrap() { + decoded.push(message); + } + + decoded + } + + #[tokio::test] + async fn decodes_messages_with_varied_chunk_and_buffer_sizes() { + let messages = vec![ + Bytes::from_static(b"first"), + Bytes::from(vec![b'a'; 64]), + Bytes::from_static(b"last"), + ]; + let payload_len = encode_messages(&messages).len(); + + for (chunk_size, buffer_size) in [(1, 1), (7, 4), (16, 8), (payload_len, 2)] { + assert_eq!( + decode_messages(&messages, chunk_size, buffer_size).await, + messages + ); + } + } +}