mirror of https://github.com/fafhrd91/actix-web
Skip empty chucks for BodyStream and SizedStream when streaming response (#1267)
This commit is contained in:
parent
3a5b62b550
commit
02a088c70c
|
@ -5,6 +5,7 @@ use std::{fmt, mem};
|
||||||
|
|
||||||
use bytes::{Bytes, BytesMut};
|
use bytes::{Bytes, BytesMut};
|
||||||
use futures_core::Stream;
|
use futures_core::Stream;
|
||||||
|
use futures_util::ready;
|
||||||
use pin_project::{pin_project, project};
|
use pin_project::{pin_project, project};
|
||||||
|
|
||||||
use crate::error::Error;
|
use crate::error::Error;
|
||||||
|
@ -389,12 +390,19 @@ where
|
||||||
BodySize::Stream
|
BodySize::Stream
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Attempts to pull out the next value of the underlying [`Stream`].
|
||||||
|
///
|
||||||
|
/// Empty values are skipped to prevent [`BodyStream`]'s transmission being
|
||||||
|
/// ended on a zero-length chunk, but rather proceed until the underlying
|
||||||
|
/// [`Stream`] ends.
|
||||||
fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
|
fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
|
||||||
unsafe { Pin::new_unchecked(self) }
|
let mut stream = unsafe { Pin::new_unchecked(self) }.project().stream;
|
||||||
.project()
|
loop {
|
||||||
.stream
|
return Poll::Ready(match ready!(stream.as_mut().poll_next(cx)) {
|
||||||
.poll_next(cx)
|
Some(Ok(ref bytes)) if bytes.is_empty() => continue,
|
||||||
.map(|res| res.map(|res| res.map_err(std::convert::Into::into)))
|
opt => opt.map(|res| res.map_err(Into::into)),
|
||||||
|
});
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -424,17 +432,27 @@ where
|
||||||
BodySize::Sized64(self.size)
|
BodySize::Sized64(self.size)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Attempts to pull out the next value of the underlying [`Stream`].
|
||||||
|
///
|
||||||
|
/// Empty values are skipped to prevent [`SizedStream`]'s transmission being
|
||||||
|
/// ended on a zero-length chunk, but rather proceed until the underlying
|
||||||
|
/// [`Stream`] ends.
|
||||||
fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
|
fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Bytes, Error>>> {
|
||||||
unsafe { Pin::new_unchecked(self) }
|
let mut stream = unsafe { Pin::new_unchecked(self) }.project().stream;
|
||||||
.project()
|
loop {
|
||||||
.stream
|
return Poll::Ready(match ready!(stream.as_mut().poll_next(cx)) {
|
||||||
.poll_next(cx)
|
Some(Ok(ref bytes)) if bytes.is_empty() => continue,
|
||||||
|
val => val,
|
||||||
|
});
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use bytes::BufMut;
|
||||||
|
use futures::stream;
|
||||||
use futures_util::future::poll_fn;
|
use futures_util::future::poll_fn;
|
||||||
|
|
||||||
impl Body {
|
impl Body {
|
||||||
|
@ -589,4 +607,45 @@ mod tests {
|
||||||
BodySize::Sized(25)
|
BodySize::Sized(25)
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
mod body_stream {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[actix_rt::test]
|
||||||
|
async fn skips_empty_chunks() {
|
||||||
|
let body = BodyStream::new(stream::iter(
|
||||||
|
["1", "", "2"]
|
||||||
|
.iter()
|
||||||
|
.map(|&v| Ok(Bytes::from(v)) as Result<Bytes, ()>),
|
||||||
|
));
|
||||||
|
assert_eq!(read_all(body).await.unwrap(), Bytes::from("12"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
mod sized_stream {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[actix_rt::test]
|
||||||
|
async fn skips_empty_chunks() {
|
||||||
|
let body = SizedStream::new(
|
||||||
|
2,
|
||||||
|
stream::iter(
|
||||||
|
["1", "", "2"]
|
||||||
|
.iter()
|
||||||
|
.map(|&v| Ok(Bytes::from(v))),
|
||||||
|
),
|
||||||
|
);
|
||||||
|
assert_eq!(read_all(body).await.unwrap(), Bytes::from("12"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn read_all<B: MessageBody>(mut body: B) -> Result<Bytes, Error> {
|
||||||
|
use futures::StreamExt as _;
|
||||||
|
|
||||||
|
let mut bytes = BytesMut::new();
|
||||||
|
while let Some(b) = stream::poll_fn(|cx| body.poll_next(cx)).next().await {
|
||||||
|
bytes.put(b?);
|
||||||
|
}
|
||||||
|
Ok(bytes.freeze())
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue