diff --git a/tokio-stream/src/iter.rs b/tokio-stream/src/iter.rs index 128be616f..37ca3b0fa 100644 --- a/tokio-stream/src/iter.rs +++ b/tokio-stream/src/iter.rs @@ -8,6 +8,7 @@ use core::task::{Context, Poll}; #[must_use = "streams do nothing unless polled"] pub struct Iter { 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> { - // 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) + } } } diff --git a/tokio-stream/tests/stream_iter.rs b/tokio-stream/tests/stream_iter.rs index 8b9ee3ce5..c98cd4853 100644 --- a/tokio-stream/tests/stream_iter.rs +++ b/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 +}