1
0
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:
Rob Ede
2021-12-04 19:40:47 +00:00
committed by GitHub
parent a2d5c5a058
commit c7c02ef99d
84 changed files with 2134 additions and 1685 deletions

View File

@ -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)
}
}

View File

@ -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,
}