diff --git a/http-body-util/Cargo.toml b/http-body-util/Cargo.toml index 019418d..79d5bdb 100644 --- a/http-body-util/Cargo.toml +++ b/http-body-util/Cargo.toml @@ -26,6 +26,11 @@ keywords = ["http"] categories = ["web-programming"] rust-version = "1.61" +[features] +default = [] +channel = ["dep:tokio"] +full = ["channel"] + [dependencies] bytes = "1" futures-core = { version = "0.3", default-features = false } @@ -33,6 +38,9 @@ http = "1" http-body = { version = "1", path = "../http-body" } pin-project-lite = "0.2" +# optional dependencies +tokio = { version = "1", features = ["sync"], optional = true } + [dev-dependencies] futures-util = { version = "0.3", default-features = false } tokio = { version = "1", features = ["macros", "rt", "sync", "rt-multi-thread"] } diff --git a/http-body-util/src/channel.rs b/http-body-util/src/channel.rs new file mode 100644 index 0000000..e963d19 --- /dev/null +++ b/http-body-util/src/channel.rs @@ -0,0 +1,234 @@ +//! A body backed by a channel. + +use std::{ + fmt::Display, + pin::Pin, + task::{Context, Poll}, +}; + +use bytes::Buf; +use http::HeaderMap; +use http_body::{Body, Frame}; +use pin_project_lite::pin_project; +use tokio::sync::{mpsc, oneshot}; + +pin_project! { + /// A body backed by a channel. + pub struct Channel { + rx_frame: mpsc::Receiver>, + #[pin] + rx_error: oneshot::Receiver, + } +} + +impl Channel { + /// Create a new channel body. + /// + /// The channel will buffer up to the provided number of messages. Once the buffer is full, + /// attempts to send new messages will wait until a message is received from the channel. The + /// provided buffer capacity must be at least 1. + pub fn new(buffer: usize) -> (Sender, Self) { + let (tx_frame, rx_frame) = mpsc::channel(buffer); + let (tx_error, rx_error) = oneshot::channel(); + (Sender { tx_frame, tx_error }, Self { rx_frame, rx_error }) + } +} + +impl Body for Channel +where + D: Buf, +{ + type Data = D; + type Error = E; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.project(); + + match this.rx_frame.poll_recv(cx) { + Poll::Ready(frame @ Some(_)) => return Poll::Ready(frame.map(Ok)), + Poll::Ready(None) | Poll::Pending => {} + } + + use core::future::Future; + match this.rx_error.poll(cx) { + Poll::Ready(Ok(error)) => return Poll::Ready(Some(Err(error))), + Poll::Ready(Err(_)) => return Poll::Ready(None), + Poll::Pending => {} + } + + Poll::Pending + } +} + +impl std::fmt::Debug for Channel { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Channel") + .field("rx_frame", &self.rx_frame) + .field("rx_error", &self.rx_error) + .finish() + } +} + +/// A sender half created through [`Channel::new`]. +pub struct Sender { + tx_frame: mpsc::Sender>, + tx_error: oneshot::Sender, +} + +impl Sender { + /// Send a frame on the channel. + pub async fn send(&mut self, frame: Frame) -> Result<(), SendError> { + self.tx_frame.send(frame).await.map_err(|_| SendError) + } + + /// Send data on data channel. + pub async fn send_data(&mut self, buf: D) -> Result<(), SendError> { + self.send(Frame::data(buf)).await + } + + /// Send trailers on trailers channel. + pub async fn send_trailers(&mut self, trailers: HeaderMap) -> Result<(), SendError> { + self.send(Frame::trailers(trailers)).await + } + + /// Aborts the body in an abnormal fashion. + pub fn abort(self, error: E) { + self.tx_error.send(error).ok(); + } +} + +impl std::fmt::Debug for Sender { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Sender") + .field("tx_frame", &self.tx_frame) + .field("tx_error", &self.tx_error) + .finish() + } +} + +/// The error returned if [`Sender`] fails to send because the receiver is closed. +#[derive(Debug)] +#[non_exhaustive] +pub struct SendError; + +impl Display for SendError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "failed to send frame") + } +} + +impl std::error::Error for SendError {} + +#[cfg(test)] +mod tests { + use bytes::Bytes; + use http::{HeaderName, HeaderValue}; + + use crate::BodyExt; + + use super::*; + + #[tokio::test] + async fn empty() { + let (tx, body) = Channel::::new(1024); + drop(tx); + + let collected = body.collect().await.unwrap(); + assert!(collected.trailers().is_none()); + assert!(collected.to_bytes().is_empty()); + } + + #[tokio::test] + async fn can_send_data() { + let (mut tx, body) = Channel::::new(1024); + + tokio::spawn(async move { + tx.send_data(Bytes::from("Hel")).await.unwrap(); + tx.send_data(Bytes::from("lo!")).await.unwrap(); + }); + + let collected = body.collect().await.unwrap(); + assert!(collected.trailers().is_none()); + assert_eq!(collected.to_bytes(), "Hello!"); + } + + #[tokio::test] + async fn can_send_trailers() { + let (mut tx, body) = Channel::::new(1024); + + tokio::spawn(async move { + let mut trailers = HeaderMap::new(); + trailers.insert( + HeaderName::from_static("foo"), + HeaderValue::from_static("bar"), + ); + tx.send_trailers(trailers).await.unwrap(); + }); + + let collected = body.collect().await.unwrap(); + assert_eq!(collected.trailers().unwrap()["foo"], "bar"); + assert!(collected.to_bytes().is_empty()); + } + + #[tokio::test] + async fn can_send_both_data_and_trailers() { + let (mut tx, body) = Channel::::new(1024); + + tokio::spawn(async move { + tx.send_data(Bytes::from("Hel")).await.unwrap(); + tx.send_data(Bytes::from("lo!")).await.unwrap(); + let mut trailers = HeaderMap::new(); + trailers.insert( + HeaderName::from_static("foo"), + HeaderValue::from_static("bar"), + ); + tx.send_trailers(trailers).await.unwrap(); + }); + + let collected = body.collect().await.unwrap(); + assert_eq!(collected.trailers().unwrap()["foo"], "bar"); + assert_eq!(collected.to_bytes(), "Hello!"); + } + + /// A stand-in for an error type, for unit tests. + type Error = &'static str; + /// An example error message. + const MSG: Error = "oh no"; + + #[tokio::test] + async fn aborts_before_trailers() { + let (mut tx, body) = Channel::::new(1024); + + tokio::spawn(async move { + tx.send_data(Bytes::from("Hel")).await.unwrap(); + tx.send_data(Bytes::from("lo!")).await.unwrap(); + tx.abort(MSG); + }); + + let err = body.collect().await.unwrap_err(); + assert_eq!(err, MSG); + } + + #[tokio::test] + async fn aborts_after_trailers() { + let (mut tx, body) = Channel::::new(1024); + + tokio::spawn(async move { + tx.send_data(Bytes::from("Hel")).await.unwrap(); + tx.send_data(Bytes::from("lo!")).await.unwrap(); + let mut trailers = HeaderMap::new(); + trailers.insert( + HeaderName::from_static("foo"), + HeaderValue::from_static("bar"), + ); + tx.send_trailers(trailers).await.unwrap(); + tx.abort(MSG); + }); + + let err = body.collect().await.unwrap_err(); + assert_eq!(err, MSG); + } +} diff --git a/http-body-util/src/lib.rs b/http-body-util/src/lib.rs index dee852c..28709fb 100644 --- a/http-body-util/src/lib.rs +++ b/http-body-util/src/lib.rs @@ -15,6 +15,9 @@ mod full; mod limited; mod stream; +#[cfg(feature = "channel")] +pub mod channel; + mod util; use self::combinators::{BoxBody, MapErr, MapFrame, UnsyncBoxBody}; @@ -26,6 +29,9 @@ pub use self::full::Full; pub use self::limited::{LengthLimitError, Limited}; pub use self::stream::{BodyDataStream, BodyStream, StreamBody}; +#[cfg(feature = "channel")] +pub use self::channel::Channel; + /// An extension trait for [`http_body::Body`] adding various combinators and adapters pub trait BodyExt: http_body::Body { /// Returns a future that resolves to the next [`Frame`], if any.