diff --git a/tokio/src/io/util/copy.rs b/tokio/src/io/util/copy.rs index 3cd425b34..fbd77b5f2 100644 --- a/tokio/src/io/util/copy.rs +++ b/tokio/src/io/util/copy.rs @@ -8,6 +8,7 @@ use std::task::{Context, Poll}; #[derive(Debug)] pub(super) struct CopyBuffer { read_done: bool, + need_flush: bool, pos: usize, cap: usize, amt: u64, @@ -18,6 +19,7 @@ impl CopyBuffer { pub(super) fn new() -> Self { Self { read_done: false, + need_flush: false, pos: 0, cap: 0, amt: 0, @@ -41,7 +43,22 @@ impl CopyBuffer { if self.pos == self.cap && !self.read_done { let me = &mut *self; let mut buf = ReadBuf::new(&mut me.buf); - ready!(reader.as_mut().poll_read(cx, &mut buf))?; + + match reader.as_mut().poll_read(cx, &mut buf) { + Poll::Ready(Ok(_)) => (), + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + Poll::Pending => { + // Try flushing when the reader has no progress to avoid deadlock + // when the reader depends on buffered writer. + if self.need_flush { + ready!(writer.as_mut().poll_flush(cx))?; + self.need_flush = false; + } + + return Poll::Pending; + } + } + let n = buf.filled().len(); if n == 0 { self.read_done = true; @@ -63,6 +80,7 @@ impl CopyBuffer { } else { self.pos += i; self.amt += i as u64; + self.need_flush = true; } } diff --git a/tokio/tests/io_copy.rs b/tokio/tests/io_copy.rs index 9ed7995c7..005e17011 100644 --- a/tokio/tests/io_copy.rs +++ b/tokio/tests/io_copy.rs @@ -1,7 +1,9 @@ #![warn(rust_2018_idioms)] #![cfg(feature = "full")] -use tokio::io::{self, AsyncRead, ReadBuf}; +use bytes::BytesMut; +use futures::ready; +use tokio::io::{self, AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf}; use tokio_test::assert_ok; use std::pin::Pin; @@ -34,3 +36,52 @@ async fn copy() { assert_eq!(n, 11); assert_eq!(wr, b"hello world"); } + +#[tokio::test] +async fn proxy() { + struct BufferedWd { + buf: BytesMut, + writer: io::DuplexStream, + } + + impl AsyncWrite for BufferedWd { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + self.get_mut().buf.extend_from_slice(buf); + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + + while !this.buf.is_empty() { + let n = ready!(Pin::new(&mut this.writer).poll_write(cx, &this.buf))?; + let _ = this.buf.split_to(n); + } + + Pin::new(&mut this.writer).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.writer).poll_shutdown(cx) + } + } + + let (rd, wd) = io::duplex(1024); + let mut rd = rd.take(1024); + let mut wd = BufferedWd { + buf: BytesMut::new(), + writer: wd, + }; + + // write start bytes + assert_ok!(wd.write_all(&[0x42; 512]).await); + assert_ok!(wd.flush().await); + + let n = assert_ok!(io::copy(&mut rd, &mut wd).await); + + assert_eq!(n, 1024); +}