stream: use cooperative budgeting in tokio_stream::iter (#8218)

tokio-stream/src/iter.rs

@@ -8,6 +8,7 @@ use core::task::{Context, Poll};
 #[must_use = "streams do nothing unless polled"]
 pub struct Iter<I> {
     iter: I,
+    #[cfg(not(feature = "rt"))]
     yield_amt: usize,
 }
 
@@ -36,6 +37,7 @@ where
 {
     Iter {
         iter: i.into_iter(),
+        #[cfg(not(feature = "rt"))]
         yield_amt: 0,
     }
 }
@@ -47,17 +49,33 @@ where
     type Item = I::Item;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<I::Item>> {
-        // TODO: add coop back
-        if self.yield_amt >= 32 {
-            self.yield_amt = 0;
+        #[cfg(feature = "rt")]
+        {
+            use tokio::task::coop;
 
-            cx.waker().wake_by_ref();
+            let coop = std::task::ready!(coop::poll_proceed(cx));
+            let item = self.iter.next();
 
-            Poll::Pending
-        } else {
-            self.yield_amt += 1;
+            coop.made_progress();
 
-            Poll::Ready(self.iter.next())
+            Poll::Ready(item)
+        }
+
+        #[cfg(not(feature = "rt"))]
+        {
+            if self.yield_amt >= 32 {
+                self.yield_amt = 0;
+
+                cx.waker().wake_by_ref();
+
+                Poll::Pending
+            } else {
+                let item = self.iter.next();
+
+                self.yield_amt += 1;
+
+                Poll::Ready(item)
+            }
         }
     }
 

tokio-stream/tests/stream_iter.rs

@@ -1,7 +1,6 @@
-use tokio_stream as stream;
-use tokio_test::task;
-
 use std::iter;
+use tokio_stream::{self as stream, Stream};
+use tokio_test::{assert_pending, assert_ready, task};
 
 #[tokio::test]
 async fn coop() {
@@ -9,6 +8,7 @@ async fn coop() {
 
     for _ in 0..10_000 {
         if stream.poll_next().is_pending() {
+            tokio::task::yield_now().await;
             assert!(stream.is_woken());
             return;
         }
@@ -16,3 +16,37 @@ async fn coop() {
 
     panic!("did not yield");
 }
+
+#[tokio::test]
+async fn test_iter_coop_budget() {
+    let mut stream = task::spawn(stream::iter(iter::repeat(1)));
+
+    // Tokio's default budget is 128.
+    // Fallback yield_amt is 32.
+    let limit = if cfg!(feature = "rt") { 128 } else { 32 };
+
+    for i in 0..limit {
+        let res = stream.poll_next();
+        assert!(res.is_ready(), "Should be ready at index {i}");
+    }
+
+    // Next poll should be pending
+    assert_pending!(stream.poll_next());
+
+    tokio::task::yield_now().await;
+    assert!(stream.is_woken());
+}
+
+#[tokio::test]
+async fn test_iter_size_hint() {
+    let stream = stream::iter(vec![1, 2, 3]);
+    assert_eq!(stream.size_hint(), (3, Some(3)));
+}
+
+#[tokio::test]
+async fn test_iter_eof_behavior() {
+    let mut stream = task::spawn(stream::iter(vec![1]));
+
+    assert_ready!(stream.poll_next());
+    assert_ready!(stream.poll_next()); // EOF should be ready None
+}