stream: avoid yielding in AllFuture and AnyFuture (#3625)

tokio-rs · Mar 21, 2021 · 69b129b · 69b129b
1 parent 0bfcbc8
commit 69b129b
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 18 deletions.
diff --git a/tokio-stream/src/stream_ext/all.rs b/tokio-stream/src/stream_ext/all.rs
@@ -38,18 +38,21 @@ where
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let me = self.project();
-        let next = futures_core::ready!(Pin::new(me.stream).poll_next(cx));
+        let mut stream = Pin::new(me.stream);
 
-        match next {
-            Some(v) => {
-                if !(me.f)(v) {
-                    Poll::Ready(false)
-                } else {
-                    cx.waker().wake_by_ref();
-                    Poll::Pending
+        // Take a maximum of 32 items from the stream before yielding.
+        for _ in 0..32 {
+            match futures_core::ready!(stream.as_mut().poll_next(cx)) {
+                Some(v) => {
+                    if !(me.f)(v) {
+                        return Poll::Ready(false);
+                    }
                 }
+                None => return Poll::Ready(true),
             }
-            None => Poll::Ready(true),
         }
+
+        cx.waker().wake_by_ref();
+        Poll::Pending
     }
 }
diff --git a/tokio-stream/src/stream_ext/any.rs b/tokio-stream/src/stream_ext/any.rs
@@ -38,18 +38,21 @@ where
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let me = self.project();
-        let next = futures_core::ready!(Pin::new(me.stream).poll_next(cx));
+        let mut stream = Pin::new(me.stream);
 
-        match next {
-            Some(v) => {
-                if (me.f)(v) {
-                    Poll::Ready(true)
-                } else {
-                    cx.waker().wake_by_ref();
-                    Poll::Pending
+        // Take a maximum of 32 items from the stream before yielding.
+        for _ in 0..32 {
+            match futures_core::ready!(stream.as_mut().poll_next(cx)) {
+                Some(v) => {
+                    if (me.f)(v) {
+                        return Poll::Ready(true);
+                    }
                 }
+                None => return Poll::Ready(false),
             }
-            None => Poll::Ready(false),
         }
+
+        cx.waker().wake_by_ref();
+        Poll::Pending
     }
 }