feat: state management and pause option

Author: fraidev

Cargo.lock

@@ -20,7 +20,7 @@ use crate::{
     peer::PeerAddr,
     peer_state::{PeerState, PeerStates},
     protocol::{Protocol, ProtocolError},
-    session::PieceWork,
+    session::{DownloadState, PieceWork},
     utils,
 };
 
@@ -226,6 +226,7 @@ pub struct PeerHandler {
     requests_sem: Semaphore,
     peer: PeerAddr,
     torrent_downloaded_state: Arc<TorrentDownloadedState>,
+    download_state: Arc<Mutex<DownloadState>>,
 }
 
 impl PeerHandler {
@@ -237,6 +238,7 @@ impl PeerHandler {
         peers_state: Arc<PeerStates>,
         //pieces: Vec<PieceWork>,
         torrent_downloaded_state: Arc<TorrentDownloadedState>,
+        download_state: Arc<Mutex<DownloadState>>,
     ) -> Self {
         Self {
             unchoke_notify: unchoked_notify,
@@ -249,6 +251,7 @@ impl PeerHandler {
             peer_writer_tx,
             peer,
             torrent_downloaded_state,
+            download_state,
             //torrent_downloaded_state: Arc::new(TorrentDownloadedState {
             //
             //    semaphore: Semaphore::new(1),
@@ -278,6 +281,14 @@ impl PeerHandler {
         }
     }
 
+    pub fn get_download_state(&self) -> DownloadState {
+        *self.download_state.lock().unwrap()
+    }
+
+    pub fn is_downloading(&self) -> bool {
+        self.get_download_state() == DownloadState::Downloading
+    }
+
     // The job of this is to request chunks and also to keep peer alive.
     // The moment this ends, the peer is disconnected.
     pub async fn task_peer_chunk_requester(&self) -> Result<(), anyhow::Error> {
@@ -307,6 +318,11 @@ impl PeerHandler {
         };
 
         loop {
+            // Wait while not downloading
+            while !self.is_downloading() {
+                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+            }
+
             update_interest(self, true)?;
 
             trace!("waiting for unchoke");
@@ -335,6 +351,14 @@ impl PeerHandler {
 
             let mut offset: u32 = 0;
             while offset < piece.length {
+                // Check download state before requesting each block
+                if !self.is_downloading() {
+                    // Wait while not downloading
+                    while !self.is_downloading() {
+                        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+                    }
+                }
+
                 loop {
                     match (tokio::time::timeout(
                         Duration::from_secs(5),

crates/bit_rev/src/peer_connection.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 use crate::file::{self, TorrentMeta};
 use crate::peer_state::PeerStates;
@@ -8,6 +8,13 @@ use crate::utils;
 use dashmap::DashMap;
 use flume::Receiver;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum DownloadState {
+    Init,
+    Downloading,
+    Paused,
+}
+
 #[derive(Debug, Clone, Copy)]
 pub struct PieceWork {
     pub index: u32,
@@ -31,6 +38,7 @@ pub struct State {
 
 pub struct Session {
     pub streams: DashMap<[u8; 20], TrackerPeers>,
+    pub download_state: Arc<Mutex<DownloadState>>,
 }
 
 pub struct AddTorrentOptions {
@@ -70,9 +78,56 @@ impl Session {
     pub fn new() -> Self {
         Self {
             streams: DashMap::new(),
+            download_state: Arc::new(Mutex::new(DownloadState::Init)),
+        }
+    }
+
+    pub fn start_downloading(&self) {
+        {
+            let mut state = self.download_state.lock().unwrap();
+            *state = DownloadState::Downloading;
+        }
+        for entry in self.streams.iter() {
+            entry.value().set_download_state(DownloadState::Downloading);
         }
     }
 
+    pub fn pause(&self) {
+        {
+            let mut state = self.download_state.lock().unwrap();
+            *state = DownloadState::Paused;
+        }
+        for entry in self.streams.iter() {
+            entry.value().set_download_state(DownloadState::Paused);
+        }
+    }
+
+    pub fn resume(&self) {
+        {
+            let mut state = self.download_state.lock().unwrap();
+            *state = DownloadState::Downloading;
+        }
+        for entry in self.streams.iter() {
+            entry.value().set_download_state(DownloadState::Downloading);
+        }
+    }
+
+    pub fn get_download_state(&self) -> DownloadState {
+        *self.download_state.lock().unwrap()
+    }
+
+    pub fn is_paused(&self) -> bool {
+        self.get_download_state() == DownloadState::Paused
+    }
+
+    pub fn is_downloading(&self) -> bool {
+        self.get_download_state() == DownloadState::Downloading
+    }
+
+    pub fn is_init(&self) -> bool {
+        self.get_download_state() == DownloadState::Init
+    }
+
     pub async fn add_torrent(
         &self,
         add_torrent: AddTorrentOptions,
@@ -91,6 +146,7 @@ impl Session {
             peer_states,
             have_broadcast.clone(),
             pr_rx.clone(),
+            self.download_state.clone(),
         );
 
         let pieces_of_work = (0..(torrent.piece_hashes.len()) as u64)
@@ -128,6 +184,9 @@ impl Session {
         self.streams
             .insert(torrent.info_hash, tracker_stream.clone());
 
+        // Start downloading
+        self.start_downloading();
+
         Ok(AddTorrentResult {
             torrent,
             torrent_meta,