tonic: avoid excess decode buffer reservation

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);

tonic/src/codec/decode.rs

@@ -220,7 +220,9 @@
                 )));
             }
 
-            self.buf.reserve(len);
+            if self.buf.capacity() < len {
+                self.buf.reserve(len - self.buf.len());
+            }
 
             self.state = State::ReadBody {
                 compression: compression_encoding,
@@ -318,7 +320,7 @@
     }
 
     fn response(&mut self) -> Result<(), Status> {
         if let Direction::Response(status) = self.direction {
             if let Err(Some(e)) = crate::status::infer_grpc_status(self.trailers.as_ref(), status) {
                 // If the trailers contain a grpc-status, then we should return that as the error
                 // and otherwise stop the stream (by taking the error state)
@@ -453,3 +455,82 @@
 
 #[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<Option<Self::Item>, 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<Bytes> {
+        let payload = encode_messages(messages);
+        let frames = payload
+            .chunks(chunk_size)
+            .map(|chunk| Ok::<_, Status>(Frame::data(Bytes::copy_from_slice(chunk))))
+            .collect::<Vec<_>>();
+        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
+            );
+        }
+    }
+}