feat(ext/http): Add support for trailers w/internal API (HTTP/2 only) (#19182)

Necessary for #3326. 

Requested in #10214 as well.

1 files changed, 16 insertions, 1 deletions

diff --git a/ext/http/response_body.rs b/ext/http/response_body.rs
index e30c917c3..ea6cc5ab8 100644
--- a/ext/http/response_body.rs
+++ b/ext/http/response_body.rs
@@ -158,7 +158,11 @@ impl std::fmt::Debug for ResponseBytesInner {
 /// required by hyper. As the API requires information about request completion (including a success/fail
 /// flag), we include a very lightweight [`CompletionHandle`] for interested parties to listen on.
 #[derive(Debug, Default)]
-pub struct ResponseBytes(ResponseBytesInner, CompletionHandle);
+pub struct ResponseBytes(
+  ResponseBytesInner,
+  CompletionHandle,
+  Rc<RefCell<Option<HeaderMap>>>,
+);
 
 impl ResponseBytes {
   pub fn initialize(&mut self, inner: ResponseBytesInner) {
@@ -170,6 +174,10 @@ impl ResponseBytes {
     self.1.clone()
   }
 
+  pub fn trailers(&self) -> Rc<RefCell<Option<HeaderMap>>> {
+    self.2.clone()
+  }
+
   fn complete(&mut self, success: bool) -> ResponseBytesInner {
     if matches!(self.0, ResponseBytesInner::Done) {
       return ResponseBytesInner::Done;
@@ -250,6 +258,9 @@ impl Body for ResponseBytes {
     let res = loop {
       let res = match &mut self.0 {
         ResponseBytesInner::Done | ResponseBytesInner::Empty => {
+          if let Some(trailers) = self.2.borrow_mut().take() {
+            return std::task::Poll::Ready(Some(Ok(Frame::trailers(trailers))));
+          }
           unreachable!()
         }
         ResponseBytesInner::Bytes(..) => {
@@ -271,6 +282,9 @@ impl Body for ResponseBytes {
     };
 
     if matches!(res, ResponseStreamResult::EndOfStream) {
+      if let Some(trailers) = self.2.borrow_mut().take() {
+        return std::task::Poll::Ready(Some(Ok(Frame::trailers(trailers))));
+      }
       self.complete(true);
     }
     std::task::Poll::Ready(res.into())
@@ -278,6 +292,7 @@ impl Body for ResponseBytes {
 
   fn is_end_stream(&self) -> bool {
     matches!(self.0, ResponseBytesInner::Done | ResponseBytesInner::Empty)
+      && self.2.borrow_mut().is_none()
   }
 
   fn size_hint(&self) -> SizeHint {