From afb5bbd94bb27f04dbd2ab81178018a3b39efdfb Mon Sep 17 00:00:00 2001 From: David Pedersen Date: Sun, 19 Nov 2023 17:51:49 +0100 Subject: [PATCH] Add channel body --- http-body-util/Cargo.toml | 8 ++ http-body-util/src/channel.rs | 159 ++++++++++++++++++++++++++++++++++ http-body-util/src/lib.rs | 6 ++ 3 files changed, 173 insertions(+) create mode 100644 http-body-util/src/channel.rs diff --git a/http-body-util/Cargo.toml b/http-body-util/Cargo.toml index c9d36aa..adbb2b9 100644 --- a/http-body-util/Cargo.toml +++ b/http-body-util/Cargo.toml @@ -25,6 +25,11 @@ Combinators and adapters for HTTP request or response bodies. keywords = ["http"] categories = ["web-programming"] +[features] +default = [] +channel = ["dep:tokio"] +full = ["channel"] + [dependencies] bytes = "1" futures-util = { version = "0.3.14", default-features = false, features = ["alloc"] } @@ -32,5 +37,8 @@ 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] tokio = { version = "1", features = ["macros", "rt"] } diff --git a/http-body-util/src/channel.rs b/http-body-util/src/channel.rs new file mode 100644 index 0000000..fc3220f --- /dev/null +++ b/http-body-util/src/channel.rs @@ -0,0 +1,159 @@ +//! 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 tokio::sync::mpsc; + +/// A body backed by a channel. +pub struct Channel { + rx_frame: mpsc::Receiver>, + rx_error: mpsc::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) = mpsc::channel(1); + (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( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + match self.rx_frame.poll_recv(cx) { + Poll::Ready(frame) => return Poll::Ready(frame.map(Ok)), + Poll::Pending => {} + } + + match self.rx_error.poll_recv(cx) { + Poll::Ready(err) => return Poll::Ready(err.map(Err)), + 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: mpsc::Sender, +} + +impl Sender { + /// Send a frame on the channel. + pub async fn send(&self, frame: Frame) -> Result<(), SendError> { + self.tx_frame.send(frame).await.map_err(|_| SendError) + } + + /// Send data on data channel. + pub async fn send_data(&self, buf: D) -> Result<(), SendError> { + self.send(Frame::data(buf)).await + } + + /// Send trailers on trailers channel. + pub async fn send_trailers(&self, trailers: HeaderMap) -> Result<(), SendError> { + self.send(Frame::trailers(trailers)).await + } + + /// Aborts the body in an abnormal fashion. + pub fn abort(self, error: E) { + match self.tx_error.try_send(error) { + Ok(_) => {} + Err(err) => { + match err { + mpsc::error::TrySendError::Full(_) => { + // Channel::new creates the error channel with space for 1 message and we + // only send once because this method consumes `self`. So the receiver + // can't be full. + unreachable!("error receiver should never be full") + } + mpsc::error::TrySendError::Closed(_) => {} + } + } + } + } +} + +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 works() { + let (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!"); + } +} diff --git a/http-body-util/src/lib.rs b/http-body-util/src/lib.rs index 059ada6..7da96a6 100644 --- a/http-body-util/src/lib.rs +++ b/http-body-util/src/lib.rs @@ -20,6 +20,9 @@ mod full; mod limited; mod stream; +#[cfg(feature = "channel")] +pub mod channel; + mod util; use self::combinators::{BoxBody, MapErr, MapFrame, UnsyncBoxBody}; @@ -31,6 +34,9 @@ pub use self::full::Full; pub use self::limited::{LengthLimitError, Limited}; pub use self::stream::{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.