From 0f0210f6b81531196d5a2670b1b8335d2f7f2a77 Mon Sep 17 00:00:00 2001
From: Naim A <227396+naim94a@users.noreply.github.com>
Date: Sun, 9 Dec 2018 01:19:04 +0200
Subject: [PATCH] loading and saving databases

---
 src/main.rs    | 46 +++++++++++++++++++++++++++++++++-------------
 src/tracker.rs |  8 ++++----
 udpt.toml      |  1 +
 3 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 95239f4..819fd7f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -94,19 +94,22 @@ fn main() {
             let file_path = std::path::Path::new(path);
             if !file_path.exists() {
                 warn!("database file \"{}\" doesn't exist.", path);
+                tracker::TorrentTracker::new(cfg.get_mode().clone())
             }
-            let mut input_file = match std::fs::File::open(file_path) {
-                Ok(v) => v,
-                Err(err) => {
-                    error!("failed to open \"{}\". error: {}", path.as_str(), err);
-                    panic!("error opening file. check logs.");
-                }
-            };
-            match tracker::TorrentTracker::load_database(cfg.get_mode().clone(), &mut input_file) {
-                Ok(v) => v,
-                Err(err) => {
-                    error!("failed to load database. error: {}", err);
-                    panic!("failed to load database. check logs.");
+            else {
+                let mut input_file = match std::fs::File::open(file_path) {
+                    Ok(v) => v,
+                    Err(err) => {
+                        error!("failed to open \"{}\". error: {}", path.as_str(), err);
+                        panic!("error opening file. check logs.");
+                    }
+                };
+                match tracker::TorrentTracker::load_database(cfg.get_mode().clone(), &mut input_file) {
+                    Ok(v) => v,
+                    Err(err) => {
+                        error!("failed to load database. error: {}", err);
+                        panic!("failed to load database. check logs.");
+                    }
                 }
             }
         }
@@ -124,7 +127,7 @@ fn main() {
         });
     }
 
-    let udp_server = std::sync::Arc::new(server::UDPTracker::new(cfg, tracker.clone()).unwrap());
+    let udp_server = std::sync::Arc::new(server::UDPTracker::new(cfg.clone(), tracker.clone()).unwrap());
 
     trace!("Waiting for UDP packets");
     let logical_cpus = num_cpus::get();
@@ -142,6 +145,23 @@ fn main() {
         }));
     }
 
+    match cfg.get_db_path() {
+        Some(db_path) => {
+            let db_p = db_path.clone();
+            let tracker_clone = tracker.clone();
+
+            std::thread::spawn(move || {
+                loop {
+                    std::thread::sleep_ms(1000 * 120);
+                    debug!("periodically saving database.");
+                    tracker_clone.periodic_task(db_p.as_str());
+                    debug!("database saved.");
+                }
+            });
+        },
+        None => {}
+    }
+
     while !threads.is_empty() {
         if let Some(thread) = threads.pop() {
             let _ = thread.join();
diff --git a/src/tracker.rs b/src/tracker.rs
index 57b1443..32bc640 100644
--- a/src/tracker.rs
+++ b/src/tracker.rs
@@ -282,7 +282,7 @@ impl TorrentTracker {
 
     /// flagged torrents will result in a tracking error. This is to allow enforcement against piracy.
     pub fn set_torrent_flag(&self, info_hash: &InfoHash, is_flagged: bool) {
-        if let Some(mut entry) = self
+        if let Some(entry) = self
             .database
             .torrent_peers
             .write()
@@ -369,7 +369,7 @@ impl TorrentTracker {
         serde_json::to_writer(compressor, &db)
     }
 
-    fn cleanup(&mut self) {
+    fn cleanup(&self) {
         use std::ops::Add;
 
         let now = std::time::SystemTime::now();
@@ -414,12 +414,12 @@ impl TorrentTracker {
         }
     }
 
-    pub fn periodic_task(&mut self, db_path: &str) {
+    pub fn periodic_task(&self, db_path: &str) {
         // cleanup db
         self.cleanup();
 
         // save db.
-        match std::fs::File::open(db_path) {
+        match std::fs::File::create(db_path) {
             Err(err) => {
                 error!("failed to open file '{}': {}", db_path, err);
                 return;
diff --git a/udpt.toml b/udpt.toml
index bda9232..44c0598 100644
--- a/udpt.toml
+++ b/udpt.toml
@@ -1,4 +1,5 @@
 mode = "dynamic"
+db_path = "database.json.bz2"
 
 [udp]
 announce_interval = 120         # Two minutes