Add Body::from_stream (#1848)

axum-core/Cargo.toml

@@ -18,6 +18,8 @@ futures-util = { version = "0.3", default-features = false, features = ["alloc"]
 http = "0.2.7"
 http-body = "0.4.5"
 mime = "0.3.16"
+pin-project-lite = "0.2.7"
+sync_wrapper = "0.1.1"
 tower-layer = "0.3"
 tower-service = "0.3"
 

axum-core/src/body.rs

@@ -1,13 +1,17 @@
 //! HTTP body utilities.
 
+use crate::response::{IntoResponse, Response};
 use crate::{BoxError, Error};
 use bytes::Bytes;
 use bytes::{Buf, BufMut};
 use futures_util::stream::Stream;
+use futures_util::TryStream;
 use http::HeaderMap;
 use http_body::Body as _;
+use pin_project_lite::pin_project;
 use std::pin::Pin;
 use std::task::{Context, Poll};
+use sync_wrapper::SyncWrapper;
 
 /// A boxed [`Body`] trait object.
 ///
@@ -107,6 +111,20 @@ impl Body {
     pub fn empty() -> Self {
         Self::new(http_body::Empty::new())
     }
+
+    /// Create a new `Body` from a [`Stream`].
+    ///
+    /// [`Stream`]: futures_util::stream::Stream
+    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),
+        })
+    }
 }
 
 impl Default for Body {
@@ -175,6 +193,49 @@ 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_data(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<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(chunk.into()))),
+            Some(Err(err)) => Poll::Ready(Some(Err(Error::new(err)))),
+            None => Poll::Ready(None),
+        }
+    }
+
+    #[inline]
+    fn poll_trailers(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
+        Poll::Ready(Ok(None))
+    }
+}
+
+impl IntoResponse for Body {
+    fn into_response(self) -> Response {
+        Response::new(self.0)
+    }
+}
+
 #[test]
 fn test_try_downcast() {
     assert_eq!(try_downcast::<i32, _>(5_u32), Err(5_u32));

axum-extra/src/body/async_read_body.rs

@@ -1,5 +1,5 @@
 use axum::{
-    body::{self, Bytes, HttpBody, StreamBody},
+    body::{Body, Bytes, HttpBody},
     http::HeaderMap,
     response::{IntoResponse, Response},
     Error,
@@ -47,51 +47,45 @@ pin_project! {
     #[cfg(feature = "async-read-body")]
     #[derive(Debug)]
     #[must_use]
-    pub struct AsyncReadBody<R> {
+    pub struct AsyncReadBody {
         #[pin]
-        read: StreamBody<ReaderStream<R>>,
+        body: Body,
     }
 }
 
-impl<R> AsyncReadBody<R> {
+impl AsyncReadBody {
     /// Create a new `AsyncReadBody`.
-    pub fn new(read: R) -> Self
+    pub fn new<R>(read: R) -> Self
     where
         R: AsyncRead + Send + 'static,
     {
         Self {
-            read: StreamBody::new(ReaderStream::new(read)),
+            body: Body::from_stream(ReaderStream::new(read)),
         }
     }
 }
 
-impl<R> HttpBody for AsyncReadBody<R>
-where
-    R: AsyncRead + Send + 'static,
-{
+impl HttpBody for AsyncReadBody {
     type Data = Bytes;
     type Error = Error;
 
     fn poll_data(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
     ) -> Poll<Option<Result<Self::Data, Self::Error>>> {
-        self.project().read.poll_data(cx)
+        self.project().body.poll_data(cx)
     }
 
     fn poll_trailers(
         self: Pin<&mut Self>,
-        _cx: &mut Context<'_>,
+        cx: &mut Context<'_>,
     ) -> Poll<Result<Option<HeaderMap>, Self::Error>> {
-        Poll::Ready(Ok(None))
+        self.project().body.poll_trailers(cx)
     }
 }
 
-impl<R> IntoResponse for AsyncReadBody<R>
-where
-    R: AsyncRead + Send + 'static,
-{
+impl IntoResponse for AsyncReadBody {
     fn into_response(self) -> Response {
-        Response::new(body::boxed(self))
+        self.body.into_response()
     }
 }

axum-extra/src/extract/multipart.rs

@@ -7,7 +7,7 @@ use axum::{
     body::{Body, Bytes},
     extract::FromRequest,
     response::{IntoResponse, Response},
-    BoxError, RequestExt,
+    RequestExt,
 };
 use futures_util::stream::Stream;
 use http::{
@@ -353,7 +353,7 @@ impl std::error::Error for InvalidBoundary {}
 mod tests {
     use super::*;
     use crate::test_helpers::*;
-    use axum::{body::Body, response::IntoResponse, routing::post, Router};
+    use axum::{response::IntoResponse, routing::post, Router};
 
     #[tokio::test]
     async fn content_type_with_encoding() {

axum-extra/src/json_lines.rs

@@ -2,7 +2,7 @@
 
 use axum::{
     async_trait,
-    body::{Body, StreamBody},
+    body::Body,
     extract::FromRequest,
     response::{IntoResponse, Response},
     BoxError,
@@ -166,7 +166,7 @@ where
             buf.write_all(b"\n")?;
             Ok::<_, BoxError>(buf.into_inner().freeze())
         });
-        let stream = StreamBody::new(stream);
+        let stream = Body::from_stream(stream);
 
         // there is no consensus around mime type yet
         // https://github.com/wardi/jsonlines/issues/36

axum/src/body/mod.rs

@@ -1,9 +1,5 @@
 //! HTTP body utilities.
 
-mod stream_body;
-
-pub use self::stream_body::StreamBody;
-
 #[doc(no_inline)]
 pub use http_body::{Body as HttpBody, Empty, Full};
 

axum/src/test_helpers/mod.rs

@@ -7,6 +7,5 @@ pub(crate) use self::test_client::*;
 
 pub(crate) fn assert_send<T: Send>() {}
 pub(crate) fn assert_sync<T: Sync>() {}
-pub(crate) fn assert_unpin<T: Unpin>() {}
 
 pub(crate) struct NotSendSync(*const ());