From 0528c780920bfe5972539a7772c20df5a0b21e5c Mon Sep 17 00:00:00 2001
From: quexeky <git@quexeky.dev>
Date: Mon, 4 Nov 2024 17:11:37 +1100
Subject: [PATCH] Functioning download progress updates

Signed-off-by: quexeky <git@quexeky.dev>
---
 pages/store/index.vue                        | 18 ++++++++++
 src-tauri/Cargo.lock                         |  7 ++++
 src-tauri/Cargo.toml                         |  1 +
 src-tauri/src/downloads/download_agent.rs    | 11 ++++--
 src-tauri/src/downloads/download_commands.rs | 12 +++++++
 src-tauri/src/downloads/download_logic.rs    | 13 +++++---
 src-tauri/src/downloads/progress.rs          | 35 ++++++++++++--------
 src-tauri/src/lib.rs                         |  5 +--
 src-tauri/src/tests/progress_tests.rs        | 14 ++++----
 9 files changed, 89 insertions(+), 27 deletions(-)

diff --git a/pages/store/index.vue b/pages/store/index.vue
index 1b3ea21..53db899 100644
--- a/pages/store/index.vue
+++ b/pages/store/index.vue
@@ -19,6 +19,12 @@
   >
     Cancel game download
   </button>
+  <button
+    class="w-full rounded-md p-4 bg-blue-600 text-white"
+    @click="getGameDownloadProgressWrapper"
+  >
+    Get game download progress
+  </button>
 </template>
 <script setup lang="ts">
 import { invoke } from "@tauri-apps/api/core";
@@ -66,4 +72,16 @@ function cancelGameDownloadWrapper() {
       console.log(e)
     })
 }
+async function getGameDownloadProgress() {
+  console.log("Getting game download status");
+  await invoke("get_game_download_progress", { gameId: gameId.value })
+}
+function getGameDownloadProgressWrapper() {
+  getGameDownloadProgress()
+    .then(() => {})
+    .catch((e) => {
+      console.log(e)
+    })
+
+}
 </script>
diff --git a/src-tauri/Cargo.lock b/src-tauri/Cargo.lock
index 520eb1a..e0d8a65 100644
--- a/src-tauri/Cargo.lock
+++ b/src-tauri/Cargo.lock
@@ -298,6 +298,12 @@ dependencies = [
  "system-deps",
 ]
 
+[[package]]
+name = "atomic-counter"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "62f447d68cfa5a9ab0c1c862a703da2a65b5ed1b7ce1153c9eb0169506d56019"
+
 [[package]]
 name = "atomic-waker"
 version = "1.1.2"
@@ -1018,6 +1024,7 @@ dependencies = [
 name = "drop-app"
 version = "0.1.0"
 dependencies = [
+ "atomic-counter",
  "ciborium",
  "directories",
  "env_logger",
diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml
index 79a3d0a..aca911f 100644
--- a/src-tauri/Cargo.toml
+++ b/src-tauri/Cargo.toml
@@ -46,6 +46,7 @@ versions = { version = "6.3.2", features = ["serde"] }
 urlencoding = "2.1.3"
 rustix = "0.38.37"
 md5 = "0.7.0"
+atomic-counter = "1.0.1"
 
 [dependencies.uuid]
 version = "1.10.0"
diff --git a/src-tauri/src/downloads/download_agent.rs b/src-tauri/src/downloads/download_agent.rs
index 792c3b3..43f6322 100644
--- a/src-tauri/src/downloads/download_agent.rs
+++ b/src-tauri/src/downloads/download_agent.rs
@@ -4,6 +4,7 @@ use crate::downloads::download_logic;
 use crate::downloads::manifest::{DropDownloadContext, DropManifest};
 use crate::downloads::progress::ProgressChecker;
 use crate::DB;
+use atomic_counter::RelaxedCounter;
 use log::info;
 use rustix::fs::{fallocate, FallocateFlags};
 use serde::{Deserialize, Serialize};
@@ -18,7 +19,7 @@ pub struct GameDownloadAgent {
     pub version: String,
     state: Mutex<GameDownloadState>,
     contexts: Mutex<Vec<DropDownloadContext>>,
-    progress: ProgressChecker<DropDownloadContext>,
+    pub progress: ProgressChecker<DropDownloadContext>,
     pub manifest: Mutex<Option<DropManifest>>,
     pub callback: Arc<AtomicBool>,
 }
@@ -57,8 +58,9 @@ impl GameDownloadAgent {
             callback: callback.clone(),
             progress: ProgressChecker::new(
                 Box::new(download_logic::download_game_chunk),
-                Arc::new(AtomicUsize::new(0)),
+                Arc::new(RelaxedCounter::new(0)),
                 callback,
+                0
             ),
             contexts: Mutex::new(Vec::new()),
         }
@@ -119,11 +121,16 @@ impl GameDownloadAgent {
         }
 
         let manifest_download = response.json::<DropManifest>().unwrap();
+        let length = manifest_download.iter().map(|(_, chunk)| {
+            return chunk.lengths.iter().sum::<usize>();
+        }).sum::<usize>();
+        self.progress.set_capacity(length);
         if let Ok(mut manifest) = self.manifest.lock() {
             *manifest = Some(manifest_download)
         } else {
             return Err(GameDownloadError::System(SystemError::MutexLockFailed));
         }
+        
 
         Ok(())
     }
diff --git a/src-tauri/src/downloads/download_commands.rs b/src-tauri/src/downloads/download_commands.rs
index 2c93285..cb49237 100644
--- a/src-tauri/src/downloads/download_commands.rs
+++ b/src-tauri/src/downloads/download_commands.rs
@@ -106,3 +106,15 @@ pub async fn stop_specific_game_download(
 
     Ok(())
 }
+
+#[tauri::command]
+pub async fn get_game_download_progress(
+    state: tauri::State<'_, Mutex<AppState>>,
+    game_id: String
+) -> Result<f64, String> {
+    let lock = state.lock().unwrap();
+    let download_agent = lock.game_downloads.get(&game_id).unwrap();
+    let progress = download_agent.progress.get_progress_percentage();
+    info!("{}", progress);
+    return Ok(progress)
+}
\ No newline at end of file
diff --git a/src-tauri/src/downloads/download_logic.rs b/src-tauri/src/downloads/download_logic.rs
index c92205e..fee6c9b 100644
--- a/src-tauri/src/downloads/download_logic.rs
+++ b/src-tauri/src/downloads/download_logic.rs
@@ -2,6 +2,7 @@ use crate::auth::generate_authorization_header;
 use crate::db::DatabaseImpls;
 use crate::downloads::manifest::DropDownloadContext;
 use crate::DB;
+use atomic_counter::{AtomicCounter, RelaxedCounter};
 use log::info;
 use md5::{Context, Digest};
 use std::{
@@ -9,7 +10,7 @@ use std::{
     io::{self, Error, ErrorKind, Seek, SeekFrom, Write},
     path::PathBuf,
     sync::{
-        atomic::{AtomicBool, Ordering},
+        atomic::{AtomicBool, AtomicUsize, Ordering},
         Arc,
     },
 };
@@ -19,13 +20,15 @@ pub struct DropFileWriter {
     file: File,
     hasher: Context,
     callback: Arc<AtomicBool>,
+    progress: Arc<RelaxedCounter>
 }
 impl DropFileWriter {
-    fn new(path: PathBuf, callback: Arc<AtomicBool>) -> Self {
+    fn new(path: PathBuf, callback: Arc<AtomicBool>, progress: Arc<RelaxedCounter>) -> Self {
         Self {
             file: OpenOptions::new().write(true).open(path).unwrap(),
             hasher: Context::new(),
             callback,
+            progress
         }
     }
     fn finish(mut self) -> io::Result<Digest> {
@@ -42,6 +45,8 @@ impl Write for DropFileWriter {
                 "Interrupt command recieved",
             ));
         }
+        let len = buf.len();
+        self.progress.add(len);
 
         //info!("Writing data to writer");
         self.hasher.write_all(buf).unwrap();
@@ -58,7 +63,7 @@ impl Seek for DropFileWriter {
         self.file.seek(pos)
     }
 }
-pub fn download_game_chunk(ctx: DropDownloadContext, callback: Arc<AtomicBool>) {
+pub fn download_game_chunk(ctx: DropDownloadContext, callback: Arc<AtomicBool>, progress: Arc<RelaxedCounter>) {
     if callback.load(Ordering::Acquire) {
         info!("Callback stopped download at start");
         return;
@@ -85,7 +90,7 @@ pub fn download_game_chunk(ctx: DropDownloadContext, callback: Arc<AtomicBool>)
         .send()
         .unwrap();
 
-    let mut file: DropFileWriter = DropFileWriter::new(ctx.path, callback);
+    let mut file: DropFileWriter = DropFileWriter::new(ctx.path, callback, progress);
 
     if ctx.offset != 0 {
         file.seek(SeekFrom::Start(ctx.offset))
diff --git a/src-tauri/src/downloads/progress.rs b/src-tauri/src/downloads/progress.rs
index 18ea179..e217f56 100644
--- a/src-tauri/src/downloads/progress.rs
+++ b/src-tauri/src/downloads/progress.rs
@@ -1,15 +1,17 @@
+use atomic_counter::{AtomicCounter, RelaxedCounter};
 use log::info;
 use rayon::ThreadPoolBuilder;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
-use std::sync::Arc;
+use std::sync::{Arc, Mutex};
 
 pub struct ProgressChecker<T>
 where
     T: 'static + Send + Sync,
 {
-    counter: Arc<AtomicUsize>,
-    f: Arc<Box<dyn Fn(T, Arc<AtomicBool>) + Send + Sync + 'static>>,
+    counter: Arc<RelaxedCounter>,
+    f: Arc<Box<dyn Fn(T, Arc<AtomicBool>, Arc<RelaxedCounter>) + Send + Sync + 'static>>,
     callback: Arc<AtomicBool>,
+    capacity: Mutex<usize>
 }
 
 impl<T> ProgressChecker<T>
@@ -17,20 +19,21 @@ where
     T: Send + Sync,
 {
     pub fn new(
-        f: Box<dyn Fn(T, Arc<AtomicBool>) + Send + Sync + 'static>,
-        counter_reference: Arc<AtomicUsize>,
+        f: Box<dyn Fn(T, Arc<AtomicBool>, Arc<RelaxedCounter>) + Send + Sync + 'static>,
+        counter: Arc<RelaxedCounter>,
         callback: Arc<AtomicBool>,
+        capacity: usize
     ) -> Self {
         Self {
             f: f.into(),
-            counter: counter_reference,
+            counter,
             callback,
+            capacity: capacity.into()
         }
     }
     pub fn run_contexts_sequentially(&self, contexts: Vec<T>) {
         for context in contexts {
-            (self.f)(context, self.callback.clone());
-            self.counter.fetch_add(1, Ordering::Release);
+            (self.f)(context, self.callback.clone(), self.counter.clone());
         }
     }
     pub fn run_contexts_parallel_background(&self, contexts: Vec<T>, max_threads: usize) {
@@ -43,8 +46,9 @@ where
 
         for context in contexts {
             let callback = self.callback.clone();
+            let counter = self.counter.clone();
             let f = self.f.clone();
-            threads.spawn(move || f(context, callback));
+            threads.spawn(move || f(context, callback, counter));
         }
     }
     pub fn run_context_parallel(&self, contexts: Vec<T>, max_threads: usize) {
@@ -56,20 +60,25 @@ where
         threads.scope(|s| {
             for context in contexts {
                 let callback = self.callback.clone();
+                let counter = self.counter.clone();
                 let f = self.f.clone();
                 s.spawn(move |_| {
                     info!("Running thread");
-                    f(context, callback)
+                    f(context, callback, counter)
                 });
             }
         });
         info!("Concluded scope");
     }
+    pub fn set_capacity(&self, capacity: usize) {
+        let mut lock = self.capacity.lock().unwrap();
+        *lock = capacity;
+    }
     pub fn get_progress(&self) -> usize {
-        self.counter.load(Ordering::Relaxed)
+        self.counter.get()
     }
     // I strongly dislike type casting in my own code, so I've shovelled it into here
-    pub fn get_progress_percentage<C: Into<f64>>(&self, capacity: C) -> f64 {
-        (self.get_progress() as f64) / (capacity.into())
+    pub fn get_progress_percentage(&self) -> f64 {
+        (self.get_progress() as f64) / (*self.capacity.lock().unwrap() as f64)
     }
 }
diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs
index 575fcf6..a64595a 100644
--- a/src-tauri/src/lib.rs
+++ b/src-tauri/src/lib.rs
@@ -11,7 +11,7 @@ use crate::downloads::download_agent::GameDownloadAgent;
 use auth::{auth_initiate, generate_authorization_header, recieve_handshake};
 use db::{DatabaseInterface, DATA_ROOT_DIR};
 use downloads::download_commands::{
-    queue_game_download, start_game_downloads, stop_specific_game_download,
+    get_game_download_progress, queue_game_download, start_game_downloads, stop_specific_game_download
 };
 use env_logger::Env;
 use http::{header::*, response::Builder as ResponseBuilder};
@@ -117,7 +117,8 @@ pub fn run() {
             // Downloads
             queue_game_download,
             start_game_downloads,
-            stop_specific_game_download
+            stop_specific_game_download,
+            get_game_download_progress
         ])
         .plugin(tauri_plugin_shell::init())
         .setup(|app| {
diff --git a/src-tauri/src/tests/progress_tests.rs b/src-tauri/src/tests/progress_tests.rs
index a3c8e5f..b855c9c 100644
--- a/src-tauri/src/tests/progress_tests.rs
+++ b/src-tauri/src/tests/progress_tests.rs
@@ -1,23 +1,25 @@
+use atomic_counter::RelaxedCounter;
+
 use crate::downloads::progress::ProgressChecker;
 use std::sync::atomic::{AtomicBool, AtomicUsize};
 use std::sync::Arc;
 
 #[test]
 fn test_progress_sequentially() {
-    let counter = Arc::new(AtomicUsize::new(0));
+    let counter = Arc::new(RelaxedCounter::new(0));
     let callback = Arc::new(AtomicBool::new(false));
-    let p = ProgressChecker::new(Box::new(test_fn), counter.clone(), callback);
+    let p = ProgressChecker::new(Box::new(test_fn), counter.clone(), callback, 100);
     p.run_contexts_sequentially((1..100).collect());
-    println!("Progress: {}", p.get_progress_percentage(100));
+    println!("Progress: {}", p.get_progress_percentage());
 }
 #[test]
 fn test_progress_parallel() {
-    let counter = Arc::new(AtomicUsize::new(0));
+    let counter = Arc::new(RelaxedCounter::new(0));
     let callback = Arc::new(AtomicBool::new(false));
-    let p = ProgressChecker::new(Box::new(test_fn), counter.clone(), callback);
+    let p = ProgressChecker::new(Box::new(test_fn), counter.clone(), callback, 100);
     p.run_contexts_parallel_background((1..100).collect(), 10);
 }
 
-fn test_fn(int: usize, callback: Arc<AtomicBool>) {
+fn test_fn(int: usize, callback: Arc<AtomicBool>, counter: Arc<RelaxedCounter>) {
     println!("{}", int);
 }