Bring Body::channel implementation from Hyper

Given the lack of a public implementation anymore, we don't have other choice but copying Hyper's private implementation.

Signed-off-by: David Calavera <[email protected]>

Author: calavera

examples/basic-streaming-response/Cargo.toml

@@ -6,11 +6,6 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-hyper = { version = "0.14", features = [
-    "http1",
-    "client",
-    "stream",
-] }
 lambda_runtime = { path = "../../lambda-runtime" }
 tokio = { version = "1", features = ["macros"] }
 tracing = { version = "0.1", features = ["log"] }

examples/basic-streaming-response/src/main.rs

@@ -1,5 +1,4 @@
-use hyper::body::Body;
-use lambda_runtime::{service_fn, Error, LambdaEvent, StreamResponse};
+use lambda_runtime::{body::Body, service_fn, Error, LambdaEvent, StreamResponse};
 use serde_json::Value;
 use std::{thread, time::Duration};
 

lambda-http/Cargo.toml

@@ -25,21 +25,22 @@ alb = []
 [dependencies]
 base64 = "0.21"
 bytes = { workspace = true }
+encoding_rs = "0.8"
 futures = { workspace = true }
 futures-util = { workspace = true }
 http = { workspace = true }
 http-body = { workspace = true }
 http-body-util = { workspace = true }
 hyper = { workspace = true }
 lambda_runtime = { path = "../lambda-runtime", version = "0.8.3" }
+mime = "0.3"
+percent-encoding = "2.2"
+pin-project-lite = { workspace = true }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_urlencoded = "0.7"
 tokio-stream = "0.1.2"
-mime = "0.3"
-encoding_rs = "0.8"
 url = "2.2"
-percent-encoding = "2.2"
 
 [dependencies.aws_lambda_events]
 path = "../lambda-events"

lambda-http/src/streaming.rs

@@ -63,19 +63,11 @@ where
     lambda_runtime::run(svc).await
 }
 
+pin_project_lite::pin_project! {
 pub struct BodyStream<B> {
+    #[pin]
     pub(crate) body: B,
 }
-
-impl<B> BodyStream<B>
-where
-    B: Body + Unpin + Send + 'static,
-    B::Data: Into<Bytes> + Send,
-    B::Error: Into<Error> + Send + Debug,
-{
-    fn project(self: Pin<&mut Self>) -> Pin<&mut B> {
-        unsafe { self.map_unchecked_mut(|s| &mut s.body) }
-    }
 }
 
 impl<B> Stream for BodyStream<B>
@@ -86,13 +78,16 @@ where
 {
     type Item = Result<B::Data, B::Error>;
 
-    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        match futures_util::ready!(self.project().poll_frame(cx)?) {
-            Some(frame) => match frame.into_data() {
-                Ok(data) => Poll::Ready(Some(Ok(data))),
-                Err(_frame) => Poll::Ready(None),
-            },
-            None => Poll::Ready(None),
+    #[inline]
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            match futures_util::ready!(self.as_mut().project().body.poll_frame(cx)?) {
+                Some(frame) => match frame.into_data() {
+                    Ok(data) => return Poll::Ready(Some(Ok(data))),
+                    Err(_frame) => {}
+                },
+                None => return Poll::Ready(None),
+            }
         }
     }
 }

lambda-runtime-api-client/Cargo.toml

@@ -27,7 +27,6 @@ hyper-util = { workspace = true, features = [
     "http1",
     "tokio",
 ] }
-pin-project-lite = "0.2"
 sync_wrapper = "0.1.2"
 tower = { workspace = true, features = ["util"] }
 tower-service = { workspace = true }

lambda-runtime-api-client/src/body/channel.rs

@@ -0,0 +1,110 @@
+//! Body::channel utilities. Extracted from Hyper under MIT license.
+//! https://github.com/hyperium/hyper/blob/master/LICENSE
+
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+use crate::body::{sender, watch};
+use bytes::Bytes;
+use futures_channel::mpsc;
+use futures_channel::oneshot;
+use futures_util::{stream::FusedStream, Future, Stream};
+use http::HeaderMap;
+use http_body::Body;
+use http_body::Frame;
+use http_body::SizeHint;
+pub use sender::Sender;
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub(crate) struct DecodedLength(u64);
+
+impl DecodedLength {
+    pub(crate) const CLOSE_DELIMITED: DecodedLength = DecodedLength(::std::u64::MAX);
+    pub(crate) const CHUNKED: DecodedLength = DecodedLength(::std::u64::MAX - 1);
+    pub(crate) const ZERO: DecodedLength = DecodedLength(0);
+
+    pub(crate) fn sub_if(&mut self, amt: u64) {
+        match *self {
+            DecodedLength::CHUNKED | DecodedLength::CLOSE_DELIMITED => (),
+            DecodedLength(ref mut known) => {
+                *known -= amt;
+            }
+        }
+    }
+
+    /// Converts to an Option<u64> representing a Known or Unknown length.
+    pub(crate) fn into_opt(self) -> Option<u64> {
+        match self {
+            DecodedLength::CHUNKED | DecodedLength::CLOSE_DELIMITED => None,
+            DecodedLength(known) => Some(known),
+        }
+    }
+}
+
+pub struct ChannelBody {
+    content_length: DecodedLength,
+    want_tx: watch::Sender,
+    data_rx: mpsc::Receiver<Result<Bytes, crate::Error>>,
+    trailers_rx: oneshot::Receiver<HeaderMap>,
+}
+
+pub fn channel() -> (Sender, ChannelBody) {
+    let (data_tx, data_rx) = mpsc::channel(0);
+    let (trailers_tx, trailers_rx) = oneshot::channel();
+
+    let (want_tx, want_rx) = watch::channel(sender::WANT_READY);
+
+    let tx = Sender {
+        want_rx,
+        data_tx,
+        trailers_tx: Some(trailers_tx),
+    };
+    let rx = ChannelBody {
+        content_length: DecodedLength::CHUNKED,
+        want_tx,
+        data_rx,
+        trailers_rx,
+    };
+
+    (tx, rx)
+}
+
+impl Body for ChannelBody {
+    type Data = Bytes;
+    type Error = crate::Error;
+
+    fn poll_frame(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
+        self.want_tx.send(sender::WANT_READY);
+
+        if !self.data_rx.is_terminated() {
+            if let Some(chunk) = ready!(Pin::new(&mut self.data_rx).poll_next(cx)?) {
+                self.content_length.sub_if(chunk.len() as u64);
+                return Poll::Ready(Some(Ok(Frame::data(chunk))));
+            }
+        }
+
+        // check trailers after data is terminated
+        match ready!(Pin::new(&mut self.trailers_rx).poll(cx)) {
+            Ok(t) => Poll::Ready(Some(Ok(Frame::trailers(t)))),
+            Err(_) => Poll::Ready(None),
+        }
+    }
+
+    fn is_end_stream(&self) -> bool {
+        self.content_length == DecodedLength::ZERO
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        let mut hint = SizeHint::default();
+
+        if let Some(content_length) = self.content_length.into_opt() {
+            hint.set_exact(content_length);
+        }
+
+        hint
+    }
+}

lambda-runtime-api-client/src/body.rs renamed to lambda-runtime-api-client/src/body/mod.rs

@@ -3,15 +3,26 @@
 
 use crate::{BoxError, Error};
 use bytes::Bytes;
-use futures_channel::mpsc::{self, Sender};
 use futures_util::stream::Stream;
-use futures_util::TryStream;
 use http_body::{Body as _, Frame};
 use http_body_util::{BodyExt, Collected};
-use pin_project_lite::pin_project;
 use std::pin::Pin;
 use std::task::{Context, Poll};
-use sync_wrapper::SyncWrapper;
+
+use self::channel::Sender;
+
+macro_rules! ready {
+    ($e:expr) => {
+        match $e {
+            std::task::Poll::Ready(v) => v,
+            std::task::Poll::Pending => return std::task::Poll::Pending,
+        }
+    };
+}
+
+mod channel;
+mod sender;
+mod watch;
 
 type BoxBody = http_body_util::combinators::UnsyncBoxBody<Bytes, Error>;
 
@@ -56,24 +67,24 @@ impl Body {
     }
 
     /// Create a new `Body` stream with associated Sender half.
-    pub fn channel() -> (Sender<Result<Bytes, BoxError>>, Body) {
-        let (sender, recv) = mpsc::channel::<Result<Bytes, BoxError>>(0);
-        (sender, Self::from_stream(recv))
+    pub fn channel() -> (Sender, Body) {
+        let (sender, body) = channel::channel();
+        (sender, Body::new(body))
     }
 
-    /// Create a new `Body` from a [`Stream`].
-    ///
-    /// [`Stream`]: https://docs.rs/futures-core/latest/futures_core/stream/trait.Stream.html
-    pub fn from_stream<S>(stream: S) -> Self
-    where
-        S: TryStream + Send + 'static,
-        S::Ok: Into<Bytes>,
-        S::Error: Into<BoxError>,
-    {
-        Self::new(StreamBody {
-            stream: SyncWrapper::new(stream),
-        })
-    }
+    // /// Create a new `Body` from a [`Stream`].
+    // ///
+    // /// [`Stream`]: https://docs.rs/futures-core/latest/futures_core/stream/trait.Stream.html
+    // pub fn from_stream<S>(stream: S) -> Self
+    // where
+    //     S: TryStream + Send + 'static,
+    //     S::Ok: Into<Bytes>,
+    //     S::Error: Into<BoxError>,
+    // {
+    //     Self::new(StreamBody {
+    //         stream: SyncWrapper::new(stream),
+    //     })
+    // }
 
     /// Collect the body into `Bytes`
     pub async fn collect(self) -> Result<Collected<Bytes>, Error> {
@@ -146,35 +157,3 @@ impl Stream for Body {
         }
     }
 }
-
-pin_project! {
-    struct StreamBody<S> {
-        #[pin]
-        stream: SyncWrapper<S>,
-    }
-}
-
-impl<S> http_body::Body for StreamBody<S>
-where
-    S: TryStream,
-    S::Ok: Into<Bytes>,
-    S::Error: Into<BoxError>,
-{
-    type Data = Bytes;
-    type Error = Error;
-
-    fn poll_frame(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
-        let stream = self.project().stream.get_pin_mut();
-        match futures_util::ready!(stream.try_poll_next(cx)) {
-            Some(Ok(chunk)) => Poll::Ready(Some(Ok(Frame::data(chunk.into())))),
-            Some(Err(err)) => Poll::Ready(Some(Err(Error::new(err)))),
-            None => Poll::Ready(None),
-        }
-    }
-}
-
-#[test]
-fn test_try_downcast() {
-    assert_eq!(try_downcast::<i32, _>(5_u32), Err(5_u32));
-    assert_eq!(try_downcast::<i32, _>(5_i32), Ok(5_i32));
-}