From a753eac23f29dc0321b87999834d4a1fbc673531 Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Sun, 22 Aug 2021 22:03:56 +0200 Subject: [PATCH] Remove boxing from `StreamBody` (#241) I just had a thought: Why should `response::Headers` be generic, but `body::StreamBody` should not? `StreamBody` previously boxed the stream to erase the generics. So we had `response::Headers` but `body::StreamBody`, without generics. After thinking about it I think it actually makes sense for responses to remain generic because you're able to use `impl IntoResponse` so you don't have to name the generics. Whereas in the case of `BodyStream` (an extractor) you cannot use `impl Trait` so it makes sense to box the inner body to make the type easier to name. Besides, `BodyStream` is mostly useful when the request body isn't `hyper::Body`, as that already implements `Stream`. --- CHANGELOG.md | 1 + src/body/stream_body.rs | 145 ++++++++++++++++++++++------------- src/extract/request_parts.rs | 24 +++++- src/response/mod.rs | 1 - 4 files changed, 113 insertions(+), 58 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7cf8992e..48aef5e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - **added:** Add `OriginalUri` for extracting original request URI in nested services ([#197](https://github.com/tokio-rs/axum/pull/197)) - **added:** Implement `FromRequest` for `http::Extensions` ([#169](https://github.com/tokio-rs/axum/pull/169)) - **added:** Make `RequestParts::{new, try_into_request}` public so extractors can be used outside axum ([#194](https://github.com/tokio-rs/axum/pull/194)) + - **added:** Implement `FromRequest` for `axum::body::Body` ([#241](https://github.com/tokio-rs/axum/pull/241)) - **changed:** Removed `extract::UrlParams` and `extract::UrlParamsMap`. Use `extract::Path` instead ([#154](https://github.com/tokio-rs/axum/pull/154)) - **changed:** `extractor_middleware` now requires `RequestBody: Default` ([#167](https://github.com/tokio-rs/axum/pull/167)) - **changed:** Convert `RequestAlreadyExtracted` to an enum with each possible error variant ([#167](https://github.com/tokio-rs/axum/pull/167)) diff --git a/src/body/stream_body.rs b/src/body/stream_body.rs index 5e054d05..a01a983a 100644 --- a/src/body/stream_body.rs +++ b/src/body/stream_body.rs @@ -1,9 +1,12 @@ -use crate::{BoxError, Error}; +use crate::{response::IntoResponse, BoxError, Error}; use bytes::Bytes; -use futures_util::stream::{self, Stream, TryStreamExt}; -use http::HeaderMap; +use futures_util::{ + ready, + stream::{self, TryStream}, +}; +use http::{HeaderMap, Response}; use http_body::Body; -use std::convert::Infallible; +use pin_project_lite::pin_project; use std::{ fmt, pin::Pin, @@ -11,80 +14,108 @@ use std::{ }; use sync_wrapper::SyncWrapper; -/// An [`http_body::Body`] created from a [`Stream`]. -/// -/// # Example -/// -/// ``` -/// use axum::{ -/// Router, -/// handler::get, -/// body::StreamBody, -/// }; -/// use futures::stream; -/// -/// async fn handler() -> StreamBody { -/// let chunks: Vec> = vec![ -/// Ok("Hello,"), -/// Ok(" "), -/// Ok("world!"), -/// ]; -/// let stream = stream::iter(chunks); -/// StreamBody::new(stream) -/// } -/// -/// let app = Router::new().route("/", get(handler)); -/// # async { -/// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap(); -/// # }; -/// ``` -/// -/// [`Stream`]: futures_util::stream::Stream -// this should probably be extracted to `http_body`, eventually... -pub struct StreamBody { - stream: SyncWrapper> + Send>>>, +pin_project! { + /// An [`http_body::Body`] created from a [`Stream`]. + /// + /// If purpose of this type is to be used in responses. If you want to + /// extract the request body as a stream consider using + /// [`extract::BodyStream`]. + /// + /// # Example + /// + /// ``` + /// use axum::{ + /// Router, + /// handler::get, + /// body::StreamBody, + /// response::IntoResponse, + /// }; + /// use futures::stream; + /// + /// async fn handler() -> impl IntoResponse { + /// let chunks: Vec> = vec![ + /// Ok("Hello,"), + /// Ok(" "), + /// Ok("world!"), + /// ]; + /// let stream = stream::iter(chunks); + /// StreamBody::new(stream) + /// } + /// + /// let app = Router::new().route("/", get(handler)); + /// # async { + /// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap(); + /// # }; + /// ``` + /// + /// [`Stream`]: futures_util::stream::Stream + pub struct StreamBody { + #[pin] + stream: SyncWrapper, + } } -impl StreamBody { +impl StreamBody { /// Create a new `StreamBody` from a [`Stream`]. /// /// [`Stream`]: futures_util::stream::Stream - pub fn new(stream: S) -> Self + pub fn new(stream: S) -> Self where - S: Stream> + Send + 'static, - T: Into + 'static, - E: Into + 'static, + S: TryStream + Send + 'static, + S::Ok: Into, + S::Error: Into, { - let stream = stream - .map_ok(Into::into) - .map_err(|err| Error::new(err.into())); Self { - stream: SyncWrapper::new(Box::pin(stream)), + stream: SyncWrapper::new(stream), } } } -impl Default for StreamBody { - fn default() -> Self { - Self::new(stream::empty::>()) +impl IntoResponse for StreamBody +where + S: TryStream + Send + 'static, + S::Ok: Into, + S::Error: Into, +{ + type Body = Self; + type BodyError = Error; + + fn into_response(self) -> Response { + Response::new(self) } } -impl fmt::Debug for StreamBody { +impl Default for StreamBody>> { + fn default() -> Self { + Self::new(stream::empty()) + } +} + +impl fmt::Debug for StreamBody { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_tuple("StreamBody").finish() } } -impl Body for StreamBody { +impl Body for StreamBody +where + S: TryStream, + S::Ok: Into, + S::Error: Into, +{ type Data = Bytes; type Error = Error; fn poll_data( - mut self: Pin<&mut Self>, + self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll>> { - Pin::new(self.stream.get_mut()).poll_next(cx) + let stream = self.project().stream.get_pin_mut(); + match 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), + } } fn poll_trailers( @@ -97,7 +128,11 @@ impl Body for StreamBody { #[test] fn stream_body_traits() { - crate::tests::assert_send::(); - crate::tests::assert_sync::(); - crate::tests::assert_unpin::(); + use futures_util::stream::Empty; + + type EmptyStream = StreamBody>>; + + crate::tests::assert_send::(); + crate::tests::assert_sync::(); + crate::tests::assert_unpin::(); } diff --git a/src/extract/request_parts.rs b/src/extract/request_parts.rs index aeafbc17..e7b68951 100644 --- a/src/extract/request_parts.rs +++ b/src/extract/request_parts.rs @@ -1,5 +1,5 @@ use super::{rejection::*, take_body, Extension, FromRequest, RequestParts}; -use crate::{BoxError, Error}; +use crate::{body::Body, BoxError, Error}; use async_trait::async_trait; use bytes::Bytes; use futures_util::stream::Stream; @@ -171,6 +171,11 @@ where /// Extractor that extracts the request body as a [`Stream`]. /// +/// Note if your request body is [`body::Body`] you can extract that directly +/// and since it already implements [`Stream`] you don't need this type. The +/// purpose of this type is to extract other types of request bodies as a +/// [`Stream`]. +/// /// # Example /// /// ```rust,no_run @@ -194,6 +199,7 @@ where /// ``` /// /// [`Stream`]: https://docs.rs/futures/latest/futures/stream/trait.Stream.html +/// [`body::Body`]: crate::body::Body pub struct BodyStream( SyncWrapper + Send + 'static>>>, ); @@ -238,6 +244,9 @@ fn body_stream_traits() { /// Extractor that extracts the raw request body. /// +/// Note that [`body::Body`] can be extracted directly. This purpose of this +/// type is to extract other types of request bodies. +/// /// # Example /// /// ```rust,no_run @@ -257,8 +266,10 @@ fn body_stream_traits() { /// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap(); /// # }; /// ``` +/// +/// [`body::Body`]: crate::body::Body #[derive(Debug, Default, Clone)] -pub struct RawBody(pub B); +pub struct RawBody(pub B); #[async_trait] impl FromRequest for Bytes @@ -280,6 +291,15 @@ where } } +#[async_trait] +impl FromRequest for Body { + type Rejection = BodyAlreadyExtracted; + + async fn from_request(req: &mut RequestParts) -> Result { + req.take_body().ok_or(BodyAlreadyExtracted) + } +} + #[async_trait] impl FromRequest for String where diff --git a/src/response/mod.rs b/src/response/mod.rs index 3dcf2c6f..fbed2fb8 100644 --- a/src/response/mod.rs +++ b/src/response/mod.rs @@ -198,7 +198,6 @@ macro_rules! impl_into_response_for_body { impl_into_response_for_body!(hyper::Body); impl_into_response_for_body!(Full); impl_into_response_for_body!(Empty); -impl_into_response_for_body!(crate::body::StreamBody); impl IntoResponse for http_body::combinators::BoxBody where