io: fix copy buffered write (#4001)

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;
                 }
             }
 

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<io::Result<usize>> {
+            self.get_mut().buf.extend_from_slice(buf);
+            Poll::Ready(Ok(buf.len()))
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+            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<io::Result<()>> {
+            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);
+}