feat: Implement http_body::Body::size_hint for custom body (#1713)

hyperium · May 29, 2024 · 9728c01 · 9728c01
1 parent 9df3867
commit 9728c01
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git 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
@@ -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),
+        }
+    }
 }