make repeat, sink cooperative

tokio/src/io/util/repeat.rs

@@ -50,9 +50,11 @@ impl AsyncRead for Repeat {
     #[inline]
     fn poll_read(
         self: Pin<&mut Self>,
-        _: &mut Context<'_>,
+        cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<io::Result<()>> {
+        ready!(crate::trace::trace_leaf(cx));
+        ready!(poll_proceed_and_make_progress(cx));
         // TODO: could be faster, but should we unsafe it?
         while buf.remaining() != 0 {
             buf.put_slice(&[self.byte]);
@@ -61,6 +63,20 @@ impl AsyncRead for Repeat {
     }
 }
 
+cfg_coop! {
+    fn poll_proceed_and_make_progress(cx: &mut Context<'_>) -> Poll<()> {
+        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        coop.made_progress();
+        Poll::Ready(())
+    }
+}
+
+cfg_not_coop! {
+    fn poll_proceed_and_make_progress(_: &mut Context<'_>) -> Poll<()> {
+        Poll::Ready(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

tokio/src/io/util/sink.rs

@@ -53,9 +53,12 @@ impl AsyncWrite for Sink {
     #[inline]
     fn poll_write(
         self: Pin<&mut Self>,
-        _: &mut Context<'_>,
+        cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<Result<usize, io::Error>> {
+        ready!(crate::trace::trace_leaf(cx));
+        ready!(poll_proceed_and_make_progress(cx));
+
         Poll::Ready(Ok(buf.len()))
     }
 
@@ -76,6 +79,20 @@ impl fmt::Debug for Sink {
     }
 }
 
+cfg_coop! {
+    fn poll_proceed_and_make_progress(cx: &mut Context<'_>) -> Poll<()> {
+        let coop = ready!(crate::runtime::coop::poll_proceed(cx));
+        coop.made_progress();
+        Poll::Ready(())
+    }
+}
+
+cfg_not_coop! {
+    fn poll_proceed_and_make_progress(_: &mut Context<'_>) -> Poll<()> {
+        Poll::Ready(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

tokio/tests/io_repeat.rs

@@ -0,0 +1,18 @@
+#![warn(rust_2018_idioms)]
+#![cfg(all(feature = "full"))]
+
+use tokio::io::AsyncReadExt;
+
+#[tokio::test]
+async fn repeat_is_cooperative() {
+    tokio::select! {
+        biased;
+        _ = async {
+            loop {
+                let mut buf = [0u8; 4096];
+                tokio::io::repeat(0b101).read_exact(&mut buf).await.unwrap();
+            }
+        } => {},
+        _ = tokio::task::yield_now() => {}
+    }
+}

tokio/tests/io_sink.rs

@@ -0,0 +1,18 @@
+#![warn(rust_2018_idioms)]
+#![cfg(all(feature = "full"))]
+
+use tokio::io::AsyncWriteExt;
+
+#[tokio::test]
+async fn sink_is_cooperative() {
+    tokio::select! {
+        biased;
+        _ = async {
+            loop {
+                let buf= vec![1, 2, 3];
+                tokio::io::sink().write(&buf).await.unwrap();
+            }
+        } => {},
+        _ = tokio::task::yield_now() => {}
+    }
+}