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
GongLG · Feb 8, 2022 · 6a3b85d · 6a3b85d
1 parent a9d580f
commit 6a3b85d
Showing 1 changed file with 57 additions and 6 deletions.
diff --git a/tokio/src/io/util/mem.rs b/tokio/src/io/util/mem.rs
@@ -177,10 +177,8 @@ impl Pipe {
             waker.wake();
         }
     }
-}
 
-impl AsyncRead for Pipe {
-    fn poll_read(
+    fn poll_read_internal(
         mut self: Pin<&mut Self>,
         cx: &mut task::Context<'_>,
         buf: &mut ReadBuf<'_>,
@@ -205,10 +203,8 @@ impl AsyncRead for Pipe {
             Poll::Pending
         }
     }
-}
 
-impl AsyncWrite for Pipe {
-    fn poll_write(
+    fn poll_write_internal(
         mut self: Pin<&mut Self>,
         cx: &mut task::Context<'_>,
         buf: &[u8],
@@ -230,6 +226,61 @@ impl AsyncWrite for Pipe {
         }
         Poll::Ready(Ok(len))
     }
+}
+
+impl AsyncRead for Pipe {
+    cfg_coop! {
+        fn poll_read(
+            self: Pin<&mut Self>,
+            cx: &mut task::Context<'_>,
+            buf: &mut ReadBuf<'_>,
+        ) -> Poll<std::io::Result<()>> {
+            let coop = ready!(crate::coop::poll_proceed(cx));
+
+            let ret = self.poll_read_internal(cx, buf);
+            if ret.is_ready() {
+                coop.made_progress();
+            }
+            ret
+        }
+    }
+
+    cfg_not_coop! {
+        fn poll_read(
+            mut self: Pin<&mut Self>,
+            cx: &mut task::Context<'_>,
+            buf: &mut ReadBuf<'_>,
+        ) -> Poll<std::io::Result<()>> {
+            self.poll_read_internal(cx, buf)
+        }
+    }
+}
+
+impl AsyncWrite for Pipe {
+    cfg_coop! {
+        fn poll_write(
+            self: Pin<&mut Self>,
+            cx: &mut task::Context<'_>,
+            buf: &[u8],
+        ) -> Poll<std::io::Result<usize>> {
+            let coop = ready!(crate::coop::poll_proceed(cx));
+            let ret = self.poll_write_internal(cx, buf);
+            if ret.is_ready() {
+                coop.made_progress();
+            }
+            ret
+        }
+    }
+
+    cfg_not_coop! {
+        fn poll_write(
+            mut self: Pin<&mut Self>,
+            cx: &mut task::Context<'_>,
+            buf: &[u8],
+        ) -> Poll<std::io::Result<usize>> {
+            self.poll_write_internal(cx, buf)
+        }
+    }
 
     fn poll_flush(self: Pin<&mut Self>, _: &mut task::Context<'_>) -> Poll<std::io::Result<()>> {
         Poll::Ready(Ok(()))