tonic: add compression threshold for message encoding

Messages smaller than 1024 bytes are no longer compressed even when
compression is enabled. This optimization avoids the overhead of
compression for small messages where the compression ratio would be
minimal or potentially increase the message size.

The threshold is set to 1024 bytes (matching UNCOMPRESSED_MIN_BODY_SIZE
used in tests). Messages below this threshold are sent uncompressed
with the compression flag unset in the gRPC header.

Author: nossrannug

tests/compression/src/lib.rs

@@ -66,10 +66,15 @@ impl test_server::Test for Svc {
         &self,
         _req: Request<()>,
     ) -> Result<Response<Self::CompressOutputServerStreamStream>, Status> {
-        let data = [0_u8; UNCOMPRESSED_MIN_BODY_SIZE].to_vec();
-        let stream = tokio_stream::iter(std::iter::repeat(SomeData { data }))
-            .take(2)
-            .map(Ok::<_, Status>);
+        // Messages smaller than 1024 don't get compressed and we want to
+        // test that the first message doesn't get compressed
+        let small = vec![0u8; UNCOMPRESSED_MIN_BODY_SIZE - 100];
+        let big   = vec![0u8; UNCOMPRESSED_MIN_BODY_SIZE];
+
+        let stream = tokio_stream::iter([
+            Ok::<_, Status>(SomeData { data: small }),
+            Ok::<_, Status>(SomeData { data: big }),
+        ]);
         Ok(self.prepare_response(Response::new(Box::pin(stream))))
     }
 

tests/compression/src/server_stream.rs

@@ -58,7 +58,8 @@ async fn client_enabled_server_enabled(encoding: CompressionEncoding) {
         .await
         .expect("stream empty")
         .expect("item was error");
-    assert!(response_bytes_counter.load(SeqCst) < UNCOMPRESSED_MIN_BODY_SIZE);
+    // The first message shouldn't get compressed because it's below the threshold
+    assert!(response_bytes_counter.load(SeqCst) > UNCOMPRESSED_MIN_BODY_SIZE - 100);
 
     stream
         .next()

tonic/src/codec/encode.rs

@@ -142,13 +142,16 @@ fn encode_item<T>(
 where
     T: Encoder<Error = Status>,
 {
+    const COMPRESSION_THRESHOLD: usize = 1024;
     let offset = buf.len();
 
     buf.reserve(HEADER_SIZE);
     unsafe {
         buf.advance_mut(HEADER_SIZE);
     }
 
+    let mut was_compressed = false;
+
     if let Some(encoding) = compression_encoding {
         uncompression_buf.clear();
 
@@ -158,24 +161,31 @@ where
 
         let uncompressed_len = uncompression_buf.len();
 
-        compress(
-            CompressionSettings {
-                encoding,
-                buffer_growth_interval: buffer_settings.buffer_size,
-            },
-            uncompression_buf,
-            buf,
-            uncompressed_len,
-        )
-        .map_err(|err| Status::internal(format!("Error compressing: {err}")))?;
+        if uncompressed_len >= COMPRESSION_THRESHOLD {
+            compress(
+                CompressionSettings {
+                    encoding,
+                    buffer_growth_interval: buffer_settings.buffer_size,
+                },
+                uncompression_buf,
+                buf,
+                uncompressed_len,
+            )
+            .map_err(|err| Status::internal(format!("Error compressing: {err}")))?;
+            was_compressed = true;
+        } else {
+            buf.reserve(uncompressed_len);
+            buf.extend_from_slice(&uncompression_buf[..]);
+        }
     } else {
         encoder
             .encode(item, &mut EncodeBuf::new(buf))
             .map_err(|err| Status::internal(format!("Error encoding: {err}")))?;
     }
 
     // now that we know length, we can write the header
-    finish_encoding(compression_encoding, max_message_size, &mut buf[offset..])
+    let final_compression = if was_compressed { compression_encoding } else { None };
+    finish_encoding(final_compression, max_message_size, &mut buf[offset..])
 }
 
 fn finish_encoding(