io: make duplex stream cooperative (#4470)

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() => {}
+    }
+}