Stream::merge does not end prematurely if one stream is delayed

Author: matklad

src/stream/stream/merge.rs

@@ -38,13 +38,14 @@ where
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let this = self.project();
-        if let Poll::Ready(Some(item)) = this.left.poll_next(cx) {
-            // The first stream made progress. The Merge needs to be polled
-            // again to check the progress of the second stream.
-            cx.waker().wake_by_ref();
-            Poll::Ready(Some(item))
-        } else {
-            this.right.poll_next(cx)
+        match this.left.poll_next(cx) {
+            Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
+            Poll::Ready(None) => this.right.poll_next(cx),
+            Poll::Pending => match this.right.poll_next(cx) {
+                Poll::Ready(Some(item)) => Poll::Ready(Some(item)),
+                Poll::Ready(None) => Poll::Pending,
+                Poll::Pending => Poll::Pending,
+            }
         }
     }
 }

tests/stream.rs

@@ -0,0 +1,51 @@
+use async_std::prelude::*;
+use async_std::stream;
+use async_std::sync::channel;
+use async_std::task;
+
+#[test]
+/// Checks that streams are merged fully even if one of the components
+/// experiences delay.
+fn merging_delayed_streams_work() {
+    let (sender, receiver) = channel::<i32>(10);
+
+    // Delay the left stream.
+    let s1 = receiver.chain(stream::once(92));
+    let s2 = stream::empty();
+    let mut s = s1.merge(s2);
+    let t1 = task::spawn(async move {
+        let mut xs = Vec::new();
+        while let Some(x) = s.next().await {
+            xs.push(x);
+        }
+        xs
+    });
+
+    task::block_on(async move {
+        task::sleep(std::time::Duration::from_millis(5)).await;
+        drop(sender);
+        let xs = t1.await;
+        assert_eq!(xs, vec![92])
+    });
+
+    let (sender, receiver) = channel::<i32>(10);
+
+    let s1 = stream::empty();
+    // Delay the right stream.
+    let s2 = receiver.chain(stream::once(92));
+    let mut s = s1.merge(s2);
+    let t1 = task::spawn(async move {
+        let mut xs = Vec::new();
+        while let Some(x) = s.next().await {
+            xs.push(x);
+        }
+        xs
+    });
+
+    task::block_on(async move {
+        task::sleep(std::time::Duration::from_millis(5)).await;
+        drop(sender);
+        let xs = t1.await;
+        assert_eq!(xs, vec![92])
+    });
+}