fix buggy "merge_two_streams"

crates/librqbit/src/merge_streams.rs

@@ -7,8 +7,8 @@ use futures::Stream;
 
 struct MergedStreams<S1, S2> {
     poll_count: AtomicU64,
-    s1: S1,
-    s2: S2,
+    s1: futures::stream::Fuse<S1>,
+    s2: futures::stream::Fuse<S2>,
 }
 
 pub fn merge_streams<
@@ -19,10 +19,11 @@ pub fn merge_streams<
     s1: S1,
     s2: S2,
 ) -> impl Stream<Item = I> + Unpin + 'static {
+    use futures::stream::StreamExt;
     MergedStreams {
         poll_count: AtomicU64::new(0),
-        s1,
-        s2,
+        s1: s1.fuse(),
+        s2: s2.fuse(),
     }
 }
 
@@ -34,8 +35,16 @@ fn poll_two<I, S1: Stream<Item = I> + Unpin, S2: Stream<Item = I> + Unpin>(
     use futures::StreamExt;
     let s1p = s1.poll_next_unpin(cx);
     match s1p {
-        Poll::Ready(r) => Poll::Ready(r),
-        Poll::Pending => s2.poll_next_unpin(cx),
+        Poll::Ready(Some(v)) => Poll::Ready(Some(v)),
+        Poll::Ready(None) => match s2.poll_next_unpin(cx) {
+            Poll::Ready(Some(v)) => Poll::Ready(Some(v)),
+            Poll::Ready(None) => Poll::Ready(None),
+            Poll::Pending => Poll::Pending,
+        },
+        Poll::Pending => match s2.poll_next_unpin(cx) {
+            Poll::Ready(Some(v)) => Poll::Ready(Some(v)),
+            Poll::Ready(None) | Poll::Pending => Poll::Pending,
+        },
     }
 }