io: make duplex stream cooperative (tokio-rs#4470)

Add coop checks on pipe poll_read and poll_write.

Fixes: tokio-rs#4470
Refs: tokio-rs#4291, tokio-rs#4300

tokio/src/io/util/mem.rs

@@ -185,6 +185,7 @@ impl AsyncRead for Pipe {
  cx: &mut task::Context<'_>,
  buf: &mut ReadBuf<'_>,
  ) -> Poll<std::io::Result<()>> {
+ ready!(poll_proceed_and_make_progress(cx));
  if self.buffer.has_remaining() {
  let max = self.buffer.remaining().min(buf.remaining());
  buf.put_slice(&self.buffer[..max]);
@@ -212,6 +213,7 @@ impl AsyncWrite for Pipe {
  cx: &mut task::Context<'_>,
  buf: &[u8],
  ) -> Poll<std::io::Result<usize>> {
+ ready!(poll_proceed_and_make_progress(cx));
  if self.is_closed {
  return Poll::Ready(Err(std::io::ErrorKind::BrokenPipe.into()));
  }
@@ -241,3 +243,17 @@ impl AsyncWrite for Pipe {
  Poll::Ready(Ok(()))
  }
 }
+
+cfg_coop! {
+ fn poll_proceed_and_make_progress(cx: &mut task::Context<'_>) -> Poll<()> {
+ let coop = ready!(crate::coop::poll_proceed(cx));
+ coop.made_progress();
+ Poll::Ready(())
+ }
+}
+
+cfg_not_coop! {
+ fn poll_proceed_and_make_progress(_: &mut Context<'_>) -> Poll<()> {
+ Poll::Ready(())
+ }
+}

tokio/tests/io_mem_stream.rs

@@ -100,3 +100,22 @@ async fn max_write_size() {
  // drop b only after task t1 finishes writing
  drop(b);
 }
+
+#[tokio::test]
+async fn duplex_is_cooperative() {
+ let (mut tx, mut rx) = tokio::io::duplex(1024 * 8);
+
+ tokio::select! {
+ biased;
+
+ _ = async {
+ loop {
+ let buf = [3u8; 4096];
+ let _ = tx.write_all(&buf).await;
+ let mut buf = [0u8; 4096];
+ let _ = rx.read(&mut buf).await;
+ }
+ } => {},
+ _ = tokio::task::yield_now() => {}
+ }
+}