diff options
author | Carl Lerche <me@carllerche.com> | 2020-03-25 14:56:24 -0700 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-03-25 14:56:24 -0700 |
commit | 186196b911bb7cbbd67e74b4ef051d3daf2d64c1 (patch) | |
tree | 5ac43c7a780e1ee0936d89607459e59f2ea90e80 | |
parent | 57ba37c97854d32e691ea68006c8d69d58c79b23 (diff) |
stream: iter() should yield every so often. (#2343)
-rw-r--r-- | tokio/src/stream/iter.rs | 3 | ||||
-rw-r--r-- | tokio/tests/stream_iter.rs | 18 |
2 files changed, 20 insertions, 1 deletions
diff --git a/tokio/src/stream/iter.rs b/tokio/src/stream/iter.rs index 5909b4cf..36eeb561 100644 --- a/tokio/src/stream/iter.rs +++ b/tokio/src/stream/iter.rs @@ -44,7 +44,8 @@ where { type Item = I::Item; - fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<I::Item>> { + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<I::Item>> { + ready!(crate::coop::poll_proceed(cx)); Poll::Ready(self.iter.next()) } diff --git a/tokio/tests/stream_iter.rs b/tokio/tests/stream_iter.rs new file mode 100644 index 00000000..45148a7a --- /dev/null +++ b/tokio/tests/stream_iter.rs @@ -0,0 +1,18 @@ +use tokio::stream; +use tokio_test::task; + +use std::iter; + +#[tokio::test] +async fn coop() { + let mut stream = task::spawn(stream::iter(iter::repeat(1))); + + for _ in 0..10_000 { + if stream.poll_next().is_pending() { + assert!(stream.is_woken()); + return; + } + } + + panic!("did not yield"); +} |