From 5798dc4a013350e704007271aa36768fb006aa7e Mon Sep 17 00:00:00 2001
From: tottoto <tottotodev@gmail.com>
Date: Wed, 29 May 2024 07:20:46 +0900
Subject: [PATCH] feat: Implement http_body::Body::size_hint for custom body

---
 tonic/src/codec/encode.rs                   | 10 ++++++++++
 tonic/src/transport/server/recover_error.rs |  7 +++++++
 2 files changed, 17 insertions(+)

diff --git a/tonic/src/codec/encode.rs b/tonic/src/codec/encode.rs
index d52294637..07aed1dda 100644
--- a/tonic/src/codec/encode.rs
+++ b/tonic/src/codec/encode.rs
@@ -330,6 +330,16 @@ where
         self.state.is_end_stream
     }
 
+    fn size_hint(&self) -> http_body::SizeHint {
+        let sh = self.inner.size_hint();
+        let mut size_hint = http_body::SizeHint::new();
+        size_hint.set_lower(sh.0 as u64);
+        if let Some(upper) = sh.1 {
+            size_hint.set_upper(upper as u64);
+        }
+        size_hint
+    }
+
     fn poll_data(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
diff --git a/tonic/src/transport/server/recover_error.rs b/tonic/src/transport/server/recover_error.rs
index 6d7e55bf4..fdb14a66a 100644
--- a/tonic/src/transport/server/recover_error.rs
+++ b/tonic/src/transport/server/recover_error.rs
@@ -124,4 +124,11 @@ where
             None => true,
         }
     }
+
+    fn size_hint(&self) -> http_body::SizeHint {
+        match &self.inner {
+            Some(body) => body.size_hint(),
+            None => http_body::SizeHint::with_exact(0),
+        }
+    }
 }