mirror of
https://github.com/fafhrd91/actix-web
synced 2025-07-01 00:44:26 +02:00
body ergonomics v3 (#2468)
This commit is contained in:
@ -12,7 +12,7 @@ use actix_rt::task::{spawn_blocking, JoinHandle};
|
||||
use bytes::Bytes;
|
||||
use derive_more::Display;
|
||||
use futures_core::ready;
|
||||
use pin_project::pin_project;
|
||||
use pin_project_lite::pin_project;
|
||||
|
||||
#[cfg(feature = "compress-brotli")]
|
||||
use brotli2::write::BrotliEncoder;
|
||||
@ -23,8 +23,10 @@ use flate2::write::{GzEncoder, ZlibEncoder};
|
||||
#[cfg(feature = "compress-zstd")]
|
||||
use zstd::stream::write::Encoder as ZstdEncoder;
|
||||
|
||||
use super::Writer;
|
||||
use crate::{
|
||||
body::{AnyBody, BodySize, MessageBody},
|
||||
body::{BodySize, MessageBody},
|
||||
error::BlockingError,
|
||||
http::{
|
||||
header::{ContentEncoding, CONTENT_ENCODING},
|
||||
HeaderValue, StatusCode,
|
||||
@ -32,84 +34,92 @@ use crate::{
|
||||
ResponseHead,
|
||||
};
|
||||
|
||||
use super::Writer;
|
||||
use crate::error::BlockingError;
|
||||
|
||||
const MAX_CHUNK_SIZE_ENCODE_IN_PLACE: usize = 1024;
|
||||
|
||||
#[pin_project]
|
||||
pub struct Encoder<B> {
|
||||
eof: bool,
|
||||
#[pin]
|
||||
body: EncoderBody<B>,
|
||||
encoder: Option<ContentEncoder>,
|
||||
fut: Option<JoinHandle<Result<ContentEncoder, io::Error>>>,
|
||||
pin_project! {
|
||||
pub struct Encoder<B> {
|
||||
#[pin]
|
||||
body: EncoderBody<B>,
|
||||
encoder: Option<ContentEncoder>,
|
||||
fut: Option<JoinHandle<Result<ContentEncoder, io::Error>>>,
|
||||
eof: bool,
|
||||
}
|
||||
}
|
||||
|
||||
impl<B: MessageBody> Encoder<B> {
|
||||
fn none() -> Self {
|
||||
Encoder {
|
||||
body: EncoderBody::None,
|
||||
encoder: None,
|
||||
fut: None,
|
||||
eof: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn response(
|
||||
encoding: ContentEncoding,
|
||||
head: &mut ResponseHead,
|
||||
body: AnyBody<B>,
|
||||
) -> AnyBody<Encoder<B>> {
|
||||
body: B,
|
||||
) -> Self {
|
||||
let can_encode = !(head.headers().contains_key(&CONTENT_ENCODING)
|
||||
|| head.status == StatusCode::SWITCHING_PROTOCOLS
|
||||
|| head.status == StatusCode::NO_CONTENT
|
||||
|| encoding == ContentEncoding::Identity
|
||||
|| encoding == ContentEncoding::Auto);
|
||||
|
||||
let body = match body {
|
||||
AnyBody::None => return AnyBody::None,
|
||||
AnyBody::Bytes(buf) => {
|
||||
if can_encode {
|
||||
EncoderBody::Bytes(buf)
|
||||
} else {
|
||||
return AnyBody::Bytes(buf);
|
||||
}
|
||||
}
|
||||
AnyBody::Body(body) => EncoderBody::Stream(body),
|
||||
};
|
||||
match body.size() {
|
||||
// no need to compress an empty body
|
||||
BodySize::None => return Self::none(),
|
||||
|
||||
// we cannot assume that Sized is not a stream
|
||||
BodySize::Sized(_) | BodySize::Stream => {}
|
||||
}
|
||||
|
||||
// TODO potentially some optimisation for single-chunk responses here by trying to read the
|
||||
// payload eagerly, stopping after 2 polls if the first is a chunk and the second is None
|
||||
|
||||
if can_encode {
|
||||
// Modify response body only if encoder is not None
|
||||
// Modify response body only if encoder is set
|
||||
if let Some(enc) = ContentEncoder::encoder(encoding) {
|
||||
update_head(encoding, head);
|
||||
head.no_chunking(false);
|
||||
|
||||
return AnyBody::Body(Encoder {
|
||||
body,
|
||||
eof: false,
|
||||
fut: None,
|
||||
return Encoder {
|
||||
body: EncoderBody::Stream { body },
|
||||
encoder: Some(enc),
|
||||
});
|
||||
fut: None,
|
||||
eof: false,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
AnyBody::Body(Encoder {
|
||||
body,
|
||||
eof: false,
|
||||
fut: None,
|
||||
Encoder {
|
||||
body: EncoderBody::Stream { body },
|
||||
encoder: None,
|
||||
})
|
||||
fut: None,
|
||||
eof: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[pin_project(project = EncoderBodyProj)]
|
||||
enum EncoderBody<B> {
|
||||
Bytes(Bytes),
|
||||
Stream(#[pin] B),
|
||||
pin_project! {
|
||||
#[project = EncoderBodyProj]
|
||||
enum EncoderBody<B> {
|
||||
None,
|
||||
Stream { #[pin] body: B },
|
||||
}
|
||||
}
|
||||
|
||||
impl<B> MessageBody for EncoderBody<B>
|
||||
where
|
||||
B: MessageBody,
|
||||
{
|
||||
type Error = EncoderError<B::Error>;
|
||||
type Error = EncoderError;
|
||||
|
||||
fn size(&self) -> BodySize {
|
||||
match self {
|
||||
EncoderBody::Bytes(ref b) => b.size(),
|
||||
EncoderBody::Stream(ref b) => b.size(),
|
||||
EncoderBody::None => BodySize::None,
|
||||
EncoderBody::Stream { body } => body.size(),
|
||||
}
|
||||
}
|
||||
|
||||
@ -118,14 +128,11 @@ where
|
||||
cx: &mut Context<'_>,
|
||||
) -> Poll<Option<Result<Bytes, Self::Error>>> {
|
||||
match self.project() {
|
||||
EncoderBodyProj::Bytes(b) => {
|
||||
if b.is_empty() {
|
||||
Poll::Ready(None)
|
||||
} else {
|
||||
Poll::Ready(Some(Ok(std::mem::take(b))))
|
||||
}
|
||||
}
|
||||
EncoderBodyProj::Stream(b) => b.poll_next(cx).map_err(EncoderError::Body),
|
||||
EncoderBodyProj::None => Poll::Ready(None),
|
||||
|
||||
EncoderBodyProj::Stream { body } => body
|
||||
.poll_next(cx)
|
||||
.map_err(|err| EncoderError::Body(err.into())),
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -134,7 +141,7 @@ impl<B> MessageBody for Encoder<B>
|
||||
where
|
||||
B: MessageBody,
|
||||
{
|
||||
type Error = EncoderError<B::Error>;
|
||||
type Error = EncoderError;
|
||||
|
||||
fn size(&self) -> BodySize {
|
||||
if self.encoder.is_none() {
|
||||
@ -197,6 +204,7 @@ where
|
||||
None => {
|
||||
if let Some(encoder) = this.encoder.take() {
|
||||
let chunk = encoder.finish().map_err(EncoderError::Io)?;
|
||||
|
||||
if chunk.is_empty() {
|
||||
return Poll::Ready(None);
|
||||
} else {
|
||||
@ -222,12 +230,15 @@ fn update_head(encoding: ContentEncoding, head: &mut ResponseHead) {
|
||||
enum ContentEncoder {
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
Deflate(ZlibEncoder<Writer>),
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
Gzip(GzEncoder<Writer>),
|
||||
|
||||
#[cfg(feature = "compress-brotli")]
|
||||
Br(BrotliEncoder<Writer>),
|
||||
// We need explicit 'static lifetime here because ZstdEncoder need lifetime
|
||||
// argument, and we use `spawn_blocking` in `Encoder::poll_next` that require `FnOnce() -> R + Send + 'static`
|
||||
|
||||
// Wwe need explicit 'static lifetime here because ZstdEncoder needs a lifetime argument and we
|
||||
// use `spawn_blocking` in `Encoder::poll_next` that requires `FnOnce() -> R + Send + 'static`.
|
||||
#[cfg(feature = "compress-zstd")]
|
||||
Zstd(ZstdEncoder<'static, Writer>),
|
||||
}
|
||||
@ -240,20 +251,24 @@ impl ContentEncoder {
|
||||
Writer::new(),
|
||||
flate2::Compression::fast(),
|
||||
))),
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
ContentEncoding::Gzip => Some(ContentEncoder::Gzip(GzEncoder::new(
|
||||
Writer::new(),
|
||||
flate2::Compression::fast(),
|
||||
))),
|
||||
|
||||
#[cfg(feature = "compress-brotli")]
|
||||
ContentEncoding::Br => {
|
||||
Some(ContentEncoder::Br(BrotliEncoder::new(Writer::new(), 3)))
|
||||
}
|
||||
|
||||
#[cfg(feature = "compress-zstd")]
|
||||
ContentEncoding::Zstd => {
|
||||
let encoder = ZstdEncoder::new(Writer::new(), 3).ok()?;
|
||||
Some(ContentEncoder::Zstd(encoder))
|
||||
}
|
||||
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@ -263,10 +278,13 @@ impl ContentEncoder {
|
||||
match *self {
|
||||
#[cfg(feature = "compress-brotli")]
|
||||
ContentEncoder::Br(ref mut encoder) => encoder.get_mut().take(),
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
ContentEncoder::Deflate(ref mut encoder) => encoder.get_mut().take(),
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
ContentEncoder::Gzip(ref mut encoder) => encoder.get_mut().take(),
|
||||
|
||||
#[cfg(feature = "compress-zstd")]
|
||||
ContentEncoder::Zstd(ref mut encoder) => encoder.get_mut().take(),
|
||||
}
|
||||
@ -279,16 +297,19 @@ impl ContentEncoder {
|
||||
Ok(writer) => Ok(writer.buf.freeze()),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
ContentEncoder::Gzip(encoder) => match encoder.finish() {
|
||||
Ok(writer) => Ok(writer.buf.freeze()),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
ContentEncoder::Deflate(encoder) => match encoder.finish() {
|
||||
Ok(writer) => Ok(writer.buf.freeze()),
|
||||
Err(err) => Err(err),
|
||||
},
|
||||
|
||||
#[cfg(feature = "compress-zstd")]
|
||||
ContentEncoder::Zstd(encoder) => match encoder.finish() {
|
||||
Ok(writer) => Ok(writer.buf.freeze()),
|
||||
@ -307,6 +328,7 @@ impl ContentEncoder {
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
ContentEncoder::Gzip(ref mut encoder) => match encoder.write_all(data) {
|
||||
Ok(_) => Ok(()),
|
||||
@ -315,6 +337,7 @@ impl ContentEncoder {
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
|
||||
#[cfg(feature = "compress-gzip")]
|
||||
ContentEncoder::Deflate(ref mut encoder) => match encoder.write_all(data) {
|
||||
Ok(_) => Ok(()),
|
||||
@ -323,6 +346,7 @@ impl ContentEncoder {
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
|
||||
#[cfg(feature = "compress-zstd")]
|
||||
ContentEncoder::Zstd(ref mut encoder) => match encoder.write_all(data) {
|
||||
Ok(_) => Ok(()),
|
||||
@ -337,9 +361,9 @@ impl ContentEncoder {
|
||||
|
||||
#[derive(Debug, Display)]
|
||||
#[non_exhaustive]
|
||||
pub enum EncoderError<E> {
|
||||
pub enum EncoderError {
|
||||
#[display(fmt = "body")]
|
||||
Body(E),
|
||||
Body(Box<dyn StdError>),
|
||||
|
||||
#[display(fmt = "blocking")]
|
||||
Blocking(BlockingError),
|
||||
@ -348,18 +372,18 @@ pub enum EncoderError<E> {
|
||||
Io(io::Error),
|
||||
}
|
||||
|
||||
impl<E: StdError + 'static> StdError for EncoderError<E> {
|
||||
impl StdError for EncoderError {
|
||||
fn source(&self) -> Option<&(dyn StdError + 'static)> {
|
||||
match self {
|
||||
EncoderError::Body(err) => Some(err),
|
||||
EncoderError::Body(err) => Some(&**err),
|
||||
EncoderError::Blocking(err) => Some(err),
|
||||
EncoderError::Io(err) => Some(err),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<E: StdError + 'static> From<EncoderError<E>> for crate::Error {
|
||||
fn from(err: EncoderError<E>) -> Self {
|
||||
impl From<EncoderError> for crate::Error {
|
||||
fn from(err: EncoderError) -> Self {
|
||||
crate::Error::new_encoder().with_cause(err)
|
||||
}
|
||||
}
|
||||
|
@ -10,6 +10,9 @@ mod encoder;
|
||||
pub use self::decoder::Decoder;
|
||||
pub use self::encoder::Encoder;
|
||||
|
||||
/// Special-purpose writer for streaming (de-)compression.
|
||||
///
|
||||
/// Pre-allocates 8KiB of capacity.
|
||||
pub(self) struct Writer {
|
||||
buf: BytesMut,
|
||||
}
|
||||
|
Reference in New Issue
Block a user