refactor: refactored persistent torrent loading/saving

mickvandijke · mickvandijke · commit ab605455ddb0 · 2022-03-16T21:36:51.000+01:00
diff --git a/src/database.rs b/src/database.rs
@@ -1,11 +1,11 @@
-use crate::{InfoHash, AUTH_KEY_LENGTH, TorrentTracker};
+use std::collections::BTreeMap;
+use crate::{InfoHash, AUTH_KEY_LENGTH, TorrentEntry};
 use log::debug;
 use r2d2_sqlite::{SqliteConnectionManager, rusqlite};
 use r2d2::{Pool};
 use r2d2_sqlite::rusqlite::NO_PARAMS;
 use crate::key_manager::AuthKey;
 use std::str::FromStr;
-use std::sync::Arc;
 
 pub struct SqliteDatabase {
     pool: Pool<SqliteConnectionManager>
@@ -73,40 +73,34 @@ impl SqliteDatabase {
         }
     }
 
-    pub async fn load_persistent_torrent_data(&self, tracker: Arc<TorrentTracker>) -> Result<bool, rusqlite::Error> {
-        let tracker_copy = tracker.clone();
+    pub async fn load_persistent_torrent_data(&self) -> Result<Vec<(InfoHash, u32)>, rusqlite::Error> {
         let conn = self.pool.get().unwrap();
         let mut stmt = conn.prepare("SELECT info_hash, completed FROM torrents")?;
 
-        let info_hash_iter = stmt.query_map(NO_PARAMS, |row| {
-            let info_hash: String = row.get(0)?;
-            let info_hash_converted = InfoHash::from_str(&info_hash).unwrap();
+        let torrent_iter = stmt.query_map(NO_PARAMS, |row| {
+            let info_hash_string: String = row.get(0)?;
+            let info_hash = InfoHash::from_str(&info_hash_string).unwrap();
             let completed: u32 = row.get(1)?;
-            Ok((info_hash_converted, completed))
+            Ok((info_hash, completed))
         })?;
 
-        for info_hash_item in info_hash_iter {
-            let (info_hash, completed): (InfoHash, u32) = info_hash_item.unwrap();
-            tracker_copy.add_torrent(&info_hash, 0u32, completed, 0u32).await;
-        }
+        let torrents: Vec<(InfoHash, u32)> = torrent_iter.filter_map(|x| x.ok() ).collect();
 
-        Ok(true)
+        Ok(torrents)
     }
 
-    pub async fn save_persistent_torrent_data(&self, tracker: Arc<TorrentTracker>) -> Result<bool, rusqlite::Error> {
-        let tracker_copy = tracker.clone();
+    pub async fn save_persistent_torrent_data(&self, torrents: &BTreeMap<InfoHash, TorrentEntry>) -> Result<(), rusqlite::Error> {
         let mut conn = self.pool.get().unwrap();
-        let db = tracker_copy.get_torrents().await;
         let db_transaction = conn.transaction()?;
-        let _: Vec<_> = db
-            .iter()
-            .map(|(info_hash, torrent_entry)| {
+
+        for (info_hash, torrent_entry) in torrents {
                 let (_seeders, completed, _leechers) = torrent_entry.get_stats();
                 let _ = db_transaction.execute("INSERT OR REPLACE INTO torrents (info_hash, completed) VALUES (?, ?)", &[info_hash.to_string(), completed.to_string()]);
-            })
-            .collect();
+        }
+
         let _ = db_transaction.commit();
-        Ok(true)
+
+        Ok(())
     }
 
     pub async fn get_info_hash_from_whitelist(&self, info_hash: &str) -> Result<InfoHash, rusqlite::Error> {
diff --git a/src/main.rs b/src/main.rs
@@ -17,11 +17,17 @@ async fn main() {
     logging::setup_logging(&config);
 
     // the singleton torrent tracker that gets passed to the HTTP and UDP server
-    let tracker = Arc::new(TorrentTracker::new(config.clone()));
+    let tracker = Arc::new(TorrentTracker::new(config.clone()).unwrap_or_else(|e| {
+        panic!("{}", e)
+    }));
 
-    // Load torrents if enabled
+    // load persistent torrents if enabled
     if config.persistence {
-        load_torrents_into_memory(tracker.clone()).await;
+        info!("Loading persistent torrents into memory...");
+        if tracker.load_torrents().await.is_err() {
+            panic!("Could not load persistent torrents.")
+        };
+        info!("Persistent torrents loaded.");
     }
 
     // start torrent cleanup job (periodically removes old peers)
@@ -48,27 +54,18 @@ async fn main() {
     // handle the signals here
     let ctrl_c = tokio::signal::ctrl_c();
     tokio::select! {
-        _ = ctrl_c => { info!("Torrust shutting down..") }
-    }
-
-    // Save torrents if enabled
-    if config.persistence {
-        save_torrents_into_memory(tracker.clone()).await;
+        _ = ctrl_c => {
+            info!("Torrust shutting down..");
+            // Save torrents if enabled
+            if config.persistence {
+                info!("Saving torrents into SQL from memory...");
+                let _ = tracker.save_torrents().await;
+                info!("Torrents saved");
+            }
+        }
     }
 }
 
-async fn load_torrents_into_memory(tracker: Arc<TorrentTracker>) {
-    info!("Loading torrents from SQL into memory...");
-    let _ = tracker.load_torrents(tracker.clone()).await;
-    info!("Torrents loaded");
-}
-
-async fn save_torrents_into_memory(tracker: Arc<TorrentTracker>) {
-    info!("Saving torrents into SQL from memory...");
-    let _ = tracker.save_torrents(tracker.clone()).await;
-    info!("Torrents saved");
-}
-
 fn start_torrent_cleanup_job(config: Arc<Configuration>, tracker: Arc<TorrentTracker>) -> Option<JoinHandle<()>> {
     let weak_tracker = std::sync::Arc::downgrade(&tracker);
     let interval = config.cleanup_interval.unwrap_or(600);
diff --git a/src/tracker.rs b/src/tracker.rs
@@ -10,7 +10,7 @@ use std::collections::btree_map::Entry;
 use crate::database::SqliteDatabase;
 use std::sync::Arc;
 use aquatic_udp_protocol::{AnnounceEvent, NumberOfBytes};
-use log::debug;
+use log::{debug};
 use crate::key_manager::AuthKey;
 use r2d2_sqlite::rusqlite;
 use crate::torrust_http_tracker::AnnounceRequest;
@@ -270,12 +270,10 @@ pub struct TorrentTracker {
 }
 
 impl TorrentTracker {
-    pub fn new(config: Arc<Configuration>) -> TorrentTracker {
-        let database = SqliteDatabase::new(&config.db_path).unwrap_or_else(|error| {
-            panic!("Could not create SQLite database. Reason: {}", error)
-        });
+    pub fn new(config: Arc<Configuration>) -> Result<TorrentTracker, rusqlite::Error> {
+        let database = SqliteDatabase::new(&config.db_path)?;
 
-        TorrentTracker {
+        Ok(TorrentTracker {
             config,
             torrents: RwLock::new(std::collections::BTreeMap::new()),
             database,
@@ -293,7 +291,7 @@ impl TorrentTracker {
                 udp6_announces_handled: 0,
                 udp6_scrapes_handled: 0,
             }),
-        }
+        })
     }
 
     fn is_public(&self) -> bool {
@@ -355,13 +353,20 @@ impl TorrentTracker {
     }
 
     // Loading the torrents into memory
-    pub async fn load_torrents(&self, tracker: Arc<TorrentTracker>) -> Result<bool, rusqlite::Error> {
-        self.database.load_persistent_torrent_data(tracker).await
+    pub async fn load_torrents(&self) -> Result<(), rusqlite::Error> {
+        let torrents = self.database.load_persistent_torrent_data().await?;
+
+        for torrent in torrents {
+            self.add_torrent(torrent.0, 0, torrent.1, 0).await;
+        }
+
+        Ok(())
     }
 
     // Saving the torrents from memory
-    pub async fn save_torrents(&self, tracker: Arc<TorrentTracker>) -> Result<bool, rusqlite::Error> {
-        self.database.save_persistent_torrent_data(tracker).await
+    pub async fn save_torrents(&self) -> Result<(), rusqlite::Error> {
+        let torrents = self.torrents.read().await;
+        self.database.save_persistent_torrent_data(&*torrents).await
     }
 
     // Adding torrents is not relevant to public trackers.
@@ -419,12 +424,15 @@ impl TorrentTracker {
         }
     }
 
-    pub async fn add_torrent(&self, info_hash: &InfoHash, seeders: u32, completed: u32, leechers: u32) -> TorrentStats {
+    pub async fn add_torrent(&self, info_hash: InfoHash, seeders: u32, completed: u32, leechers: u32) -> TorrentStats {
         let mut torrents = self.torrents.write().await;
 
         if !torrents.contains_key(&info_hash) {
-            let mut torrent_entry = TorrentEntry::new();
-            torrent_entry.completed = completed;
+            let torrent_entry = TorrentEntry {
+                peers: Default::default(),
+                completed,
+                seeders
+            };
             torrents.insert(info_hash.clone(), torrent_entry);
         }
 
