From c07a57a86cd8a3153b0424b4865552383e7498c9 Mon Sep 17 00:00:00 2001
From: "zhangli.pear" <zhangli.pear@bytedance.com>
Date: Wed, 14 Aug 2024 14:09:02 +0800
Subject: [PATCH] add data()

---
 http-body-util/src/combinators/data.rs | 29 ++++++++++++++++++++++++++
 http-body-util/src/combinators/mod.rs  |  2 ++
 http-body-util/src/lib.rs              |  9 ++++++++
 3 files changed, 40 insertions(+)
 create mode 100644 http-body-util/src/combinators/data.rs

diff --git a/http-body-util/src/combinators/data.rs b/http-body-util/src/combinators/data.rs
new file mode 100644
index 0000000..64ef639
--- /dev/null
+++ b/http-body-util/src/combinators/data.rs
@@ -0,0 +1,29 @@
+use http_body::Body;
+
+use core::future::Future;
+use core::pin::Pin;
+use core::task;
+
+#[must_use = "futures don't do anything unless polled"]
+#[derive(Debug)]
+/// Future that resolves to the next data chunk from `Body`
+pub struct Data<'a, T: ?Sized>(pub(crate) &'a mut T);
+
+impl<'a, T: Body + Unpin + ?Sized> Future for Data<'a, T> {
+    type Output = Option<Result<T::Data, T::Error>>;
+
+    fn poll(mut self: Pin<&mut Self>, ctx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
+        match Pin::new(&mut self.0).poll_frame(ctx) {
+            task::Poll::Ready(Some(Ok(frame))) => {
+                if let Ok(data) = frame.into_data() {
+                    task::Poll::Ready(Some(Ok(data)))
+                } else {
+                    task::Poll::Pending
+                }
+            },
+            task::Poll::Ready(Some(Err(e))) => task::Poll::Ready(Some(Err(e))),
+            task::Poll::Ready(None) => task::Poll::Ready(None),
+            task::Poll::Pending => task::Poll::Pending,
+        }
+    }
+}
diff --git a/http-body-util/src/combinators/mod.rs b/http-body-util/src/combinators/mod.rs
index 38d2637..d2beccb 100644
--- a/http-body-util/src/combinators/mod.rs
+++ b/http-body-util/src/combinators/mod.rs
@@ -6,6 +6,7 @@ mod frame;
 mod map_err;
 mod map_frame;
 mod with_trailers;
+mod data;
 
 pub use self::{
     box_body::{BoxBody, UnsyncBoxBody},
@@ -14,4 +15,5 @@ pub use self::{
     map_err::MapErr,
     map_frame::MapFrame,
     with_trailers::WithTrailers,
+    data::Data,
 };
diff --git a/http-body-util/src/lib.rs b/http-body-util/src/lib.rs
index dee852c..759db78 100644
--- a/http-body-util/src/lib.rs
+++ b/http-body-util/src/lib.rs
@@ -38,6 +38,15 @@ pub trait BodyExt: http_body::Body {
         combinators::Frame(self)
     }
 
+    /// Returns future that resolves to next data chunk, if any.
+    fn data(&mut self) -> combinators::Data<'_, Self>
+    where
+        Self: Unpin + Sized,
+    {
+        combinators::Data(self)
+    }
+
+
     /// Maps this body's frame to a different kind.
     fn map_frame<F, B>(self, f: F) -> MapFrame<Self, F>
     where