From 66a2ab17b0f1c25d61918389c3a32e26d736638c Mon Sep 17 00:00:00 2001
From: Dirkjan Ochtman <dirkjan@ochtman.nl>
Date: Thu, 5 Sep 2024 13:28:04 +0200
Subject: [PATCH] tonic: fold encode_client() into EncodeBody::new_client()

---
 tonic/src/client/grpc.rs  |  5 +++--
 tonic/src/codec/encode.rs | 39 +++++++++++++++------------------------
 tonic/src/codec/mod.rs    |  2 +-
 tonic/src/codec/prost.rs  |  4 +++-
 4 files changed, 22 insertions(+), 28 deletions(-)

diff --git a/tonic/src/client/grpc.rs b/tonic/src/client/grpc.rs
index 2aed93737..ba4c84669 100644
--- a/tonic/src/client/grpc.rs
+++ b/tonic/src/client/grpc.rs
@@ -1,9 +1,10 @@
 use crate::codec::compression::{CompressionEncoding, EnabledCompressionEncodings};
+use crate::codec::EncodeBody;
 use crate::metadata::GRPC_CONTENT_TYPE;
 use crate::{
     body::BoxBody,
     client::GrpcService,
-    codec::{encode_client, Codec, Decoder, Streaming},
+    codec::{Codec, Decoder, Streaming},
     request::SanitizeHeaders,
     Code, Request, Response, Status,
 };
@@ -295,7 +296,7 @@ impl<T> Grpc<T> {
     {
         let request = request
             .map(|s| {
-                encode_client(
+                EncodeBody::new_client(
                     codec.encoder(),
                     s.map(Ok),
                     self.config.send_compression_encodings,
diff --git a/tonic/src/codec/encode.rs b/tonic/src/codec/encode.rs
index 4c94f0e8f..52952ff5b 100644
--- a/tonic/src/codec/encode.rs
+++ b/tonic/src/codec/encode.rs
@@ -13,28 +13,6 @@ use std::{
 };
 use tokio_stream::{adapters::Fuse, Stream, StreamExt};
 
-/// Turns a stream of grpc messages into [EncodeBody] which is used by grpc clients for
-/// turning the messages into http frames for sending over the network.
-pub fn encode_client<T, U>(
-    encoder: T,
-    source: U,
-    compression_encoding: Option<CompressionEncoding>,
-    max_message_size: Option<usize>,
-) -> EncodeBody<T, U>
-where
-    T: Encoder<Error = Status>,
-    U: Stream,
-{
-    let stream = EncodedBytes::new(
-        encoder,
-        source,
-        compression_encoding,
-        SingleMessageCompressionOverride::default(),
-        max_message_size,
-    );
-    EncodeBody::new_client(stream)
-}
-
 /// Combinator for efficient encoding of messages into reasonably sized buffers.
 /// EncodedBytes encodes ready messages from its delegate stream into a BytesMut,
 /// splitting off and yielding a buffer when either:
@@ -251,9 +229,22 @@ struct EncodeState {
 }
 
 impl<T: Encoder, U: Stream> EncodeBody<T, U> {
-    fn new_client(inner: EncodedBytes<T, U>) -> Self {
+    /// Turns a stream of grpc messages into [EncodeBody] which is used by grpc clients for
+    /// turning the messages into http frames for sending over the network.
+    pub fn new_client(
+        encoder: T,
+        source: U,
+        compression_encoding: Option<CompressionEncoding>,
+        max_message_size: Option<usize>,
+    ) -> Self {
         Self {
-            inner,
+            inner: EncodedBytes::new(
+                encoder,
+                source,
+                compression_encoding,
+                SingleMessageCompressionOverride::default(),
+                max_message_size,
+            ),
             state: EncodeState {
                 error: None,
                 role: Role::Client,
diff --git a/tonic/src/codec/mod.rs b/tonic/src/codec/mod.rs
index 9f50d7cfa..d14a6ddf4 100644
--- a/tonic/src/codec/mod.rs
+++ b/tonic/src/codec/mod.rs
@@ -16,7 +16,7 @@ use std::io;
 pub use self::buffer::{DecodeBuf, EncodeBuf};
 pub use self::compression::{CompressionEncoding, EnabledCompressionEncodings};
 pub use self::decode::Streaming;
-pub use self::encode::{encode_client, EncodeBody};
+pub use self::encode::EncodeBody;
 #[cfg(feature = "prost")]
 pub use self::prost::ProstCodec;
 
diff --git a/tonic/src/codec/prost.rs b/tonic/src/codec/prost.rs
index 62fd692d6..ec42226a7 100644
--- a/tonic/src/codec/prost.rs
+++ b/tonic/src/codec/prost.rs
@@ -150,7 +150,9 @@ fn from_decode_error(error: prost::DecodeError) -> crate::Status {
 #[cfg(test)]
 mod tests {
     use crate::codec::compression::SingleMessageCompressionOverride;
-    use crate::codec::{DecodeBuf, Decoder, EncodeBody, EncodeBuf, Encoder, Streaming, HEADER_SIZE};
+    use crate::codec::{
+        DecodeBuf, Decoder, EncodeBody, EncodeBuf, Encoder, Streaming, HEADER_SIZE,
+    };
     use crate::Status;
     use bytes::{Buf, BufMut, BytesMut};
     use http_body::Body;