Reorganize _start() a bit

crates/librqbit/src/torrent_state/mod.rs

@@ -279,77 +279,25 @@ impl ManagedTorrent {
         peer_rx: Option<PeerStream>,
         start_paused: bool,
     ) -> anyhow::Result<()> {
-        let session = self
+        // State machine transitions.
-            .shared
+        //
-            .session
+        // - error -> initializing
-            .upgrade()
+        // - initializing -> paused
-            .context("session is dead, cannot start torrent")?;
+        // - paused -> live
-        let mut g = self.locked.write();
+        // - live -> paused
-        g.paused = start_paused;
+        //
-        let cancellation_token = session.cancellation_token().child_token();
+        // - initializing -> error
+        // - live -> error
 
-        let spawn_fatal_errors_receiver =
+        fn _start<'a>(
-            |state: &Arc<Self>,
+            t: &'a Arc<ManagedTorrent>,
-             rx: tokio::sync::oneshot::Receiver<anyhow::Error>,
+            peer_rx: Option<PeerStream>,
-             token: CancellationToken| {
+            start_paused: bool,
-                let span = state.shared.span.clone();
+            session: Arc<Session>,
-                let state = Arc::downgrade(state);
+            g: Option<parking_lot::RwLockWriteGuard<'a, ManagedTorrentLocked>>,
-                spawn_with_cancel(
+            token: CancellationToken,
-                    error_span!(parent: span, "fatal_errors_receiver"),
+        ) -> anyhow::Result<()> {
-                    token,
+            let mut g = g.unwrap_or_else(|| t.locked.write());
-                    async move {
-                        let e = match rx.await {
-                            Ok(e) => e,
-                            Err(_) => return Ok(()),
-                        };
-                        if let Some(state) = state.upgrade() {
-                            state.stop_with_error(e);
-                        } else {
-                            warn!("tried to stop the torrent with error, but couldn't upgrade the arc");
-                        }
-                        Ok(())
-                    },
-                );
-            };
-
-        fn spawn_peer_adder(live: &Arc<TorrentStateLive>, mut peer_rx: PeerStream) {
-            live.spawn(
-                error_span!(parent: live.torrent().span.clone(), "external_peer_adder"),
-                {
-                    let live = live.clone();
-                    async move {
-                        let live = {
-                            let weak = Arc::downgrade(&live);
-                            drop(live);
-                            weak
-                        };
-
-                        loop {
-                            match timeout(Duration::from_secs(5), peer_rx.next()).await {
-                                Ok(Some(peer)) => {
-                                    trace!(?peer, "received peer from peer_rx");
-                                    let live = match live.upgrade() {
-                                        Some(live) => live,
-                                        None => return Ok(()),
-                                    };
-                                    live.add_peer_if_not_seen(peer).context("torrent closed")?;
-                                }
-                                Ok(None) => {
-                                    debug!("peer_rx closed, closing peer adder");
-                                    return Ok(());
-                                }
-                                // If timeout, check if the torrent is live.
-                                Err(_) if live.strong_count() == 0 => {
-                                    debug!("timed out waiting for peers, torrent isn't live, closing peer adder");
-                                    return Ok(());
-                                }
-                                Err(_) => continue,
-                            }
-                        }
-                    }
-                },
-            );
-        }
 
             match &g.state {
                 ManagedTorrentState::Live(_) => {
@@ -358,9 +306,9 @@ impl ManagedTorrent {
                 ManagedTorrentState::Initializing(init) => {
                     let init = init.clone();
                     drop(g);
-                let t = self.clone();
+                    let t = t.clone();
-                let span = self.shared().span.clone();
+                    let span = t.shared().span.clone();
-                let token = cancellation_token.clone();
+                    let token = token.clone();
 
                     spawn_with_cancel(
                         error_span!(parent: span.clone(), "initialize_and_start"),
@@ -382,25 +330,14 @@ impl ManagedTorrent {
                                         return Ok(());
                                     }
 
-                                if start_paused {
                                     g.state = ManagedTorrentState::Paused(paused);
                                     t.state_change_notify.notify_waiters();
+
+                                    if start_paused {
                                         return Ok(());
                                     }
 
-                                let (tx, rx) = tokio::sync::oneshot::channel();
+                                    _start(&t, peer_rx, start_paused, session, Some(g), token)
-                                let live = TorrentStateLive::new(paused, tx, cancellation_token)?;
-                                g.state = ManagedTorrentState::Live(live.clone());
-                                drop(g);
-
-                                t.state_change_notify.notify_waiters();
-
-                                spawn_fatal_errors_receiver(&t, rx, token);
-                                if let Some(peer_rx) = peer_rx {
-                                    spawn_peer_adder(&live, peer_rx);
-                                }
-
-                                Ok(())
                                 }
                                 Err(err) => {
                                     let result = anyhow::anyhow!("{:?}", err);
@@ -414,41 +351,63 @@ impl ManagedTorrent {
                     Ok(())
                 }
                 ManagedTorrentState::Paused(_) => {
+                    if start_paused {
+                        warn!("start(start_paused=true) called, but torrent already paused");
+                        return Ok(());
+                    }
                     let paused = g.state.take().assert_paused();
                     let (tx, rx) = tokio::sync::oneshot::channel();
-                let live = TorrentStateLive::new(paused, tx, cancellation_token.clone())?;
+                    let live = TorrentStateLive::new(paused, tx, token.clone())?;
                     g.state = ManagedTorrentState::Live(live.clone());
+                    t.state_change_notify.notify_waiters();
                     drop(g);
 
-                spawn_fatal_errors_receiver(self, rx, cancellation_token);
+                    spawn_fatal_errors_receiver(t, rx, token);
                     if let Some(peer_rx) = peer_rx {
                         spawn_peer_adder(&live, peer_rx);
                     }
                     Ok(())
                 }
                 ManagedTorrentState::Error(_) => {
-                let metadata = self.metadata.load_full().expect("TODO");
+                    let metadata = t.metadata.load_full().expect("TODO");
                     let initializing = Arc::new(TorrentStateInitializing::new(
-                    self.shared.clone(),
+                        t.shared.clone(),
                         metadata.clone(),
                         g.only_files.clone(),
-                    self.shared
+                        t.shared
                             .storage_factory
-                        .create_and_init(self.shared(), &metadata)?,
+                            .create_and_init(t.shared(), &metadata)?,
                         true,
                     ));
                     g.state = ManagedTorrentState::Initializing(initializing.clone());
-                drop(g);
+                    t.state_change_notify.notify_waiters();
-
-                self.state_change_notify.notify_waiters();
 
                     // Recurse.
-                self.start(peer_rx, start_paused)
+                    _start(t, peer_rx, start_paused, session, Some(g), token)
                 }
                 ManagedTorrentState::None => bail!("bug: torrent is in empty state"),
             }
         }
 
+        let session = self
+            .shared
+            .session
+            .upgrade()
+            .context("session is dead, cannot start torrent")?;
+        let mut g = self.locked.write();
+        g.paused = start_paused;
+        let cancellation_token = session.cancellation_token().child_token();
+
+        _start(
+            self,
+            peer_rx,
+            start_paused,
+            session,
+            Some(g),
+            cancellation_token,
+        )
+    }
+
     pub fn is_paused(&self) -> bool {
         self.locked.read().paused
     }
@@ -618,3 +577,67 @@ impl ManagedTorrent {
 }
 
 pub type ManagedTorrentHandle = Arc<ManagedTorrent>;
+
+fn spawn_fatal_errors_receiver(
+    state: &Arc<ManagedTorrent>,
+    rx: tokio::sync::oneshot::Receiver<anyhow::Error>,
+    token: CancellationToken,
+) {
+    let span = state.shared.span.clone();
+    let state = Arc::downgrade(state);
+    spawn_with_cancel(
+        error_span!(parent: span, "fatal_errors_receiver"),
+        token,
+        async move {
+            let e = match rx.await {
+                Ok(e) => e,
+                Err(_) => return Ok(()),
+            };
+            if let Some(state) = state.upgrade() {
+                state.stop_with_error(e);
+            } else {
+                warn!("tried to stop the torrent with error, but couldn't upgrade the arc");
+            }
+            Ok(())
+        },
+    );
+}
+
+fn spawn_peer_adder(live: &Arc<TorrentStateLive>, mut peer_rx: PeerStream) {
+    live.spawn(
+        error_span!(parent: live.torrent().span.clone(), "external_peer_adder"),
+        {
+            let live = live.clone();
+            async move {
+                let live = {
+                    let weak = Arc::downgrade(&live);
+                    drop(live);
+                    weak
+                };
+
+                loop {
+                    match timeout(Duration::from_secs(5), peer_rx.next()).await {
+                        Ok(Some(peer)) => {
+                            trace!(?peer, "received peer from peer_rx");
+                            let live = match live.upgrade() {
+                                Some(live) => live,
+                                None => return Ok(()),
+                            };
+                            live.add_peer_if_not_seen(peer).context("torrent closed")?;
+                        }
+                        Ok(None) => {
+                            debug!("peer_rx closed, closing peer adder");
+                            return Ok(());
+                        }
+                        // If timeout, check if the torrent is live.
+                        Err(_) if live.strong_count() == 0 => {
+                            debug!("timed out waiting for peers, torrent isn't live, closing peer adder");
+                            return Ok(());
+                        }
+                        Err(_) => continue,
+                    }
+                }
+            }
+        },
+    );
+}