don't poll the reader again after eof while waiting for the writer to flush

Author: jbr

src/io/copy.rs

@@ -7,6 +7,12 @@ use crate::io::{self, BufRead, BufReader, Read, Write};
 use crate::task::{Context, Poll};
 use crate::utils::Context as _;
 
+// Note: There are two otherwise-identical implementations of this
+// function because unstable has removed the `?Sized` bound for the
+// reader and writer and accepts `R` and `W` instead of `&mut R` and
+// `&mut W`. If making a change to either of the implementations,
+// ensure that you copy it into the other.
+
 /// Copies the entire contents of a reader into a writer.
 ///
 /// This function will continuously read data from `reader` and then
@@ -57,6 +63,7 @@ where
  #[pin]
  writer: W,
  amt: u64,
+ reader_eof: bool
  }
  }
 
@@ -69,13 +76,20 @@ where
 
  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
  let mut this = self.project();
+
  loop {
- let buffer = futures_core::ready!(this.reader.as_mut().poll_fill_buf(cx))?;
- if buffer.is_empty() {
+ if *this.reader_eof {
  futures_core::ready!(this.writer.as_mut().poll_flush(cx))?;
  return Poll::Ready(Ok(*this.amt));
  }
 
+ let buffer = futures_core::ready!(this.reader.as_mut().poll_fill_buf(cx))?;
+
+ if buffer.is_empty() {
+ *this.reader_eof = true;
+ continue;
+ }
+
  let i = futures_core::ready!(this.writer.as_mut().poll_write(cx, buffer))?;
  if i == 0 {
  return Poll::Ready(Err(io::ErrorKind::WriteZero.into()));
@@ -89,6 +103,7 @@ where
  let future = CopyFuture {
  reader: BufReader::new(reader),
  writer,
+ reader_eof: false,
  amt: 0,
  };
  future.await.context(|| String::from("io::copy failed"))
@@ -144,6 +159,7 @@ where
  #[pin]
  writer: W,
  amt: u64,
+ reader_eof: bool
  }
  }
 
@@ -156,13 +172,20 @@ where
 
  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
  let mut this = self.project();
+
  loop {
- let buffer = futures_core::ready!(this.reader.as_mut().poll_fill_buf(cx))?;
- if buffer.is_empty() {
+ if *this.reader_eof {
  futures_core::ready!(this.writer.as_mut().poll_flush(cx))?;
  return Poll::Ready(Ok(*this.amt));
  }
 
+ let buffer = futures_core::ready!(this.reader.as_mut().poll_fill_buf(cx))?;
+
+ if buffer.is_empty() {
+ *this.reader_eof = true;
+ continue;
+ }
+
  let i = futures_core::ready!(this.writer.as_mut().poll_write(cx, buffer))?;
  if i == 0 {
  return Poll::Ready(Err(io::ErrorKind::WriteZero.into()));
@@ -176,6 +199,7 @@ where
  let future = CopyFuture {
  reader: BufReader::new(reader),
  writer,
+ reader_eof: false,
  amt: 0,
  };
  future.await.context(|| String::from("io::copy failed"))

tests/io_copy.rs

@@ -0,0 +1,72 @@
+use std::{
+ io::Result,
+ pin::Pin,
+ task::{Context, Poll},
+};
+
+struct ReaderThatPanicsAfterEof {
+ read_count: usize,
+ has_sent_eof: bool,
+ max_read: usize,
+}
+
+impl async_std::io::Read for ReaderThatPanicsAfterEof {
+ fn poll_read(
+ mut self: Pin<&mut Self>,
+ _cx: &mut Context<'_>,
+ buf: &mut [u8],
+ ) -> Poll<Result<usize>> {
+ if self.has_sent_eof {
+ panic!("this should be unreachable because we should not poll after eof (Ready(Ok(0)))")
+ } else if self.read_count >= self.max_read {
+ self.has_sent_eof = true;
+ Poll::Ready(Ok(0))
+ } else {
+ self.read_count += 1;
+ Poll::Ready(Ok(buf.len()))
+ }
+ }
+}
+
+struct WriterThatTakesAWhileToFlush {
+ max_flush: usize,
+ flush_count: usize,
+}
+
+impl async_std::io::Write for WriterThatTakesAWhileToFlush {
+ fn poll_write(self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll<Result<usize>> {
+ Poll::Ready(Ok(buf.len()))
+ }
+
+ fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<()>> {
+ self.flush_count += 1;
+ if self.flush_count >= self.max_flush {
+ Poll::Ready(Ok(()))
+ } else {
+ cx.waker().wake_by_ref();
+ Poll::Pending
+ }
+ }
+
+ fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<()>> {
+ Poll::Ready(Ok(()))
+ }
+}
+
+#[test]
+fn io_copy_does_not_poll_after_eof() {
+ async_std::task::block_on(async {
+ let mut reader = ReaderThatPanicsAfterEof {
+ has_sent_eof: false,
+ max_read: 10,
+ read_count: 0,
+ };
+
+ let mut writer = WriterThatTakesAWhileToFlush {
+ flush_count: 0,
+ max_flush: 10,
+ };
+
+ assert!(async_std::io::copy(&mut reader, &mut writer).await.is_ok());
+ })
+}