From d028e2e3c2f8a5d20e8c0eb2cb67f14660185d18 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Sun, 18 Aug 2024 17:37:03 +0100
Subject: [PATCH] fix buggy "merge_two_streams"

---
 crates/librqbit/src/merge_streams.rs     | 21 +++++++++++++++------
 crates/librqbit/src/session.rs           | 24 ++++++++++++++++++++----
 crates/librqbit/src/torrent_state/mod.rs |  7 ++++++-
 3 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/crates/librqbit/src/merge_streams.rs b/crates/librqbit/src/merge_streams.rs
index 608de9e..35205d4 100644
--- a/crates/librqbit/src/merge_streams.rs
+++ b/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,
+        },
     }
 }
 
diff --git a/crates/librqbit/src/session.rs b/crates/librqbit/src/session.rs
index ca98f79..c250534 100644
--- a/crates/librqbit/src/session.rs
+++ b/crates/librqbit/src/session.rs
@@ -199,10 +199,22 @@ fn merge_two_optional_streams<T>(
     s2: Option<impl Stream<Item = T> + Unpin + Send + 'static>,
 ) -> Option<BoxStream<'static, T>> {
     match (s1, s2) {
-        (Some(s1), None) => Some(Box::pin(s1)),
-        (None, Some(s2)) => Some(Box::pin(s2)),
-        (Some(s1), Some(s2)) => Some(Box::pin(merge_streams(s1, s2))),
-        (None, None) => None,
+        (Some(s1), None) => {
+            debug!("merge_two_optional_streams: using first");
+            Some(Box::pin(s1))
+        }
+        (None, Some(s2)) => {
+            debug!("merge_two_optional_streams: using second");
+            Some(Box::pin(s2))
+        }
+        (Some(s1), Some(s2)) => {
+            debug!("merge_two_optional_streams: using both");
+            Some(Box::pin(merge_streams(s1, s2)))
+        }
+        (None, None) => {
+            debug!("merge_two_optional_streams: using none");
+            None
+        }
     }
 }
 
@@ -1096,6 +1108,10 @@ impl Session {
         // Merge "initial_peers" and "peer_rx" into one stream.
         let peer_rx = merge_two_optional_streams(
             if !initial_peers.is_empty() {
+                debug!(
+                    count = initial_peers.len(),
+                    "merging initial peers into peer_rx"
+                );
                 Some(futures::stream::iter(initial_peers.into_iter()))
             } else {
                 None
diff --git a/crates/librqbit/src/torrent_state/mod.rs b/crates/librqbit/src/torrent_state/mod.rs
index 785465e..68f4d04 100644
--- a/crates/librqbit/src/torrent_state/mod.rs
+++ b/crates/librqbit/src/torrent_state/mod.rs
@@ -320,6 +320,8 @@ impl ManagedTorrent {
                                 let live =
                                     TorrentStateLive::new(paused, tx, live_cancellation_token)?;
                                 g.state = ManagedTorrentState::Live(live.clone());
+                                drop(g);
+
                                 t.state_change_notify.notify_waiters();
 
                                 spawn_fatal_errors_receiver(&t, rx, token);
@@ -343,6 +345,8 @@ impl ManagedTorrent {
                 let (tx, rx) = tokio::sync::oneshot::channel();
                 let live = TorrentStateLive::new(paused, tx, live_cancellation_token.clone())?;
                 g.state = ManagedTorrentState::Live(live.clone());
+                drop(g);
+
                 spawn_fatal_errors_receiver(self, rx, live_cancellation_token);
                 spawn_peer_adder(&live, peer_rx);
                 Ok(())
@@ -354,9 +358,10 @@ impl ManagedTorrent {
                     self.storage_factory.create_and_init(self.info())?,
                 ));
                 g.state = ManagedTorrentState::Initializing(initializing.clone());
-                self.state_change_notify.notify_waiters();
                 drop(g);
 
+                self.state_change_notify.notify_waiters();
+
                 // Recurse.
                 self.start(
                     peer_rx,