From 59cc48b3084001d102915c2569ceef5caf3cce22 Mon Sep 17 00:00:00 2001
From: DecDuck <declanahofmeyr@gmail.com>
Date: Thu, 26 Dec 2024 17:19:19 +1100
Subject: [PATCH] feat(process): better process management, including running
 state

---
 desktop/components/GameStatusButton.vue       |  42 ++--
 desktop/pages/library/[id]/index.vue          |   4 +-
 desktop/src-tauri/src/db.rs                   |   1 +
 .../src-tauri/src/downloads/download_agent.rs |   7 +-
 .../src/downloads/progress_object.rs          |  11 +-
 desktop/src-tauri/src/lib.rs                  |   8 +-
 .../src-tauri/src/process/process_commands.rs |   2 +-
 .../src-tauri/src/process/process_manager.rs  | 183 ++++++++++++++----
 desktop/types.ts                              |   1 +
 9 files changed, 186 insertions(+), 73 deletions(-)

diff --git a/desktop/components/GameStatusButton.vue b/desktop/components/GameStatusButton.vue
index dd94b3b0..65f3c7a1 100644
--- a/desktop/components/GameStatusButton.vue
+++ b/desktop/components/GameStatusButton.vue
@@ -1,5 +1,5 @@
 <template>
-  <div class="inline-flex">
+  <div class="inline-flex divide-x divide-zinc-900">
     <button type="button" @click="() => buttonActions[props.status.type]()" :class="[
       styles[props.status.type],
       showDropdown ? 'rounded-l-md' : 'rounded-md',
@@ -10,10 +10,11 @@
     </button>
     <Menu v-if="showDropdown" as="div" class="relative inline-block text-left grow">
       <div class="h-full">
-        <MenuButton
-          class="inline-flex w-full h-full justify-center items-center rounded-r-md bg-zinc-800 px-1 py-2 text-sm font-semibold text-zinc-100 shadow-sm ring-1 ring-inset ring-zinc-600 hover:bg-zinc-800">
-
-          <ChevronDownIcon class="size-5 text-gray-400" aria-hidden="true" />
+        <MenuButton :class="[
+          styles[props.status.type],
+          'inline-flex w-full h-full justify-center items-center rounded-r-md px-1 py-2 text-sm font-semibold shadow-sm'
+        ]">
+          <ChevronDownIcon class="size-5" aria-hidden="true" />
         </MenuButton>
       </div>
 
@@ -53,26 +54,22 @@ import { Menu, MenuButton, MenuItem, MenuItems } from '@headlessui/vue'
 const props = defineProps<{ status: GameStatus }>();
 const emit = defineEmits<{
   (e: "install"): void;
-  (e: "play"): void;
+  (e: "launch"): void;
   (e: "queue"): void;
   (e: "uninstall"): void;
 }>();
 
-const showDropdown = computed(() => props.status.type === GameStatusEnum.Installed);
+const showDropdown = computed(() => props.status.type === GameStatusEnum.Installed || props.status.type === GameStatusEnum.SetupRequired);
 
 const styles: { [key in GameStatusEnum]: string } = {
-  [GameStatusEnum.Remote]:
-    "bg-blue-600 text-white hover:bg-blue-500 focus-visible:outline-blue-600",
-  [GameStatusEnum.Queued]:
-    "bg-zinc-800 text-white hover:bg-zinc-700 focus-visible:outline-zinc-700",
-  [GameStatusEnum.Downloading]:
-    "bg-zinc-800 text-white hover:bg-zinc-700 focus-visible:outline-zinc-700",
-  [GameStatusEnum.SetupRequired]:
-    "bg-yellow-600 text-white hover:bg-yellow-500 focus-visible:outline-yellow-600",
-  [GameStatusEnum.Installed]:
-    "bg-green-600 text-white hover:bg-green-500 focus-visible:outline-green-600",
-  [GameStatusEnum.Updating]: "",
-  [GameStatusEnum.Uninstalling]: "",
+  [GameStatusEnum.Remote]: "bg-blue-600 text-white hover:bg-blue-500 focus-visible:outline-blue-600",
+  [GameStatusEnum.Queued]: "bg-zinc-800 text-white hover:bg-zinc-700 focus-visible:outline-zinc-700",
+  [GameStatusEnum.Downloading]: "bg-zinc-800 text-white hover:bg-zinc-700 focus-visible:outline-zinc-700",
+  [GameStatusEnum.SetupRequired]: "bg-yellow-600 text-white hover:bg-yellow-500 focus-visible:outline-yellow-600",
+  [GameStatusEnum.Installed]: "bg-green-600 text-white hover:bg-green-500 focus-visible:outline-green-600",
+  [GameStatusEnum.Updating]: "bg-zinc-800 text-white hover:bg-zinc-700 focus-visible:outline-zinc-700",
+  [GameStatusEnum.Uninstalling]: "bg-zinc-800 text-white hover:bg-zinc-700 focus-visible:outline-zinc-700",
+  [GameStatusEnum.Running]: "bg-zinc-800 text-white hover:bg-zinc-700 focus-visible:outline-zinc-700"
 };
 
 const buttonNames: { [key in GameStatusEnum]: string } = {
@@ -83,6 +80,7 @@ const buttonNames: { [key in GameStatusEnum]: string } = {
   [GameStatusEnum.Installed]: "Play",
   [GameStatusEnum.Updating]: "Updating",
   [GameStatusEnum.Uninstalling]: "Uninstalling",
+  [GameStatusEnum.Running]: "Running"
 };
 
 const buttonIcons: { [key in GameStatusEnum]: Component } = {
@@ -93,15 +91,17 @@ const buttonIcons: { [key in GameStatusEnum]: Component } = {
   [GameStatusEnum.Installed]: PlayIcon,
   [GameStatusEnum.Updating]: ArrowDownTrayIcon,
   [GameStatusEnum.Uninstalling]: TrashIcon,
+  [GameStatusEnum.Running]: PlayIcon
 };
 
 const buttonActions: { [key in GameStatusEnum]: () => void } = {
   [GameStatusEnum.Remote]: () => emit("install"),
   [GameStatusEnum.Queued]: () => emit("queue"),
   [GameStatusEnum.Downloading]: () => emit("queue"),
-  [GameStatusEnum.SetupRequired]: () => { },
-  [GameStatusEnum.Installed]: () => emit("play"),
+  [GameStatusEnum.SetupRequired]: () => emit("launch"),
+  [GameStatusEnum.Installed]: () => emit("launch"),
   [GameStatusEnum.Updating]: () => emit("queue"),
   [GameStatusEnum.Uninstalling]: () => { },
+  [GameStatusEnum.Running]: () => { }
 };
 </script>
diff --git a/desktop/pages/library/[id]/index.vue b/desktop/pages/library/[id]/index.vue
index 30f2bb63..4cc2e7ea 100644
--- a/desktop/pages/library/[id]/index.vue
+++ b/desktop/pages/library/[id]/index.vue
@@ -20,7 +20,7 @@
       <div class="h-full flex flex-row gap-x-4 items-stretch">
         <GameStatusButton
           @install="() => installFlow()"
-          @play="() => play()"
+          @launch="() => launch()"
           @queue="() => queue()"
           @uninstall="() => uninstall()"
           :status="status"
@@ -390,7 +390,7 @@ async function install() {
   }
 }
 
-async function play() {
+async function launch() {
   try {
     await invoke("launch_game", { gameId: game.value.id });
   } catch (e) {
diff --git a/desktop/src-tauri/src/db.rs b/desktop/src-tauri/src/db.rs
index 6e21030a..662d943b 100644
--- a/desktop/src-tauri/src/db.rs
+++ b/desktop/src-tauri/src/db.rs
@@ -47,6 +47,7 @@ pub enum GameTransientStatus {
     Downloading { version_name: String },
     Uninstalling {},
     Updating { version_name: String },
+    Running {},
 }
 
 #[derive(Serialize, Deserialize, Clone)]
diff --git a/desktop/src-tauri/src/downloads/download_agent.rs b/desktop/src-tauri/src/downloads/download_agent.rs
index b3dd3c02..21ac7deb 100644
--- a/desktop/src-tauri/src/downloads/download_agent.rs
+++ b/desktop/src-tauri/src/downloads/download_agent.rs
@@ -221,11 +221,6 @@ impl GameDownloadAgent {
 
         *self.completed_contexts.lock().unwrap() = self.stored_manifest.get_completed_contexts();
 
-        info!(
-            "Completed contexts: {:?}",
-            *self.completed_contexts.lock().unwrap()
-        );
-
         for (raw_path, chunk) in manifest {
             let path = base_path.join(Path::new(&raw_path));
 
@@ -275,6 +270,8 @@ impl GameDownloadAgent {
         pool.scope(move |scope| {
             let completed_lock = self.completed_contexts.lock().unwrap();
 
+            let count = self.contexts.len();
+
             for (index, context) in self.contexts.iter().enumerate() {
                 let progress = self.progress.get(index); // Clone arcs
                 let progress_handle = ProgressHandle::new(progress, self.progress.clone());
diff --git a/desktop/src-tauri/src/downloads/progress_object.rs b/desktop/src-tauri/src/downloads/progress_object.rs
index 143656dd..22077027 100644
--- a/desktop/src-tauri/src/downloads/progress_object.rs
+++ b/desktop/src-tauri/src/downloads/progress_object.rs
@@ -19,7 +19,7 @@ pub struct ProgressObject {
     sender: Sender<DownloadManagerSignal>,
 
     points_towards_update: Arc<AtomicUsize>,
-    points_to_push_update: Arc<Mutex<usize>>,
+    points_to_push_update: Arc<AtomicUsize>,
 }
 
 pub struct ProgressHandle {
@@ -58,7 +58,7 @@ impl ProgressObject {
             sender,
 
             points_towards_update: Arc::new(AtomicUsize::new(0)),
-            points_to_push_update: Arc::new(Mutex::new(points_to_push_update)),
+            points_to_push_update: Arc::new(AtomicUsize::new(points_to_push_update)),
         }
     }
 
@@ -67,9 +67,7 @@ impl ProgressObject {
             .points_towards_update
             .fetch_add(amount_added, Ordering::Relaxed);
 
-        let to_update_handle = self.points_to_push_update.lock().unwrap();
-        let to_update = *to_update_handle;
-        drop(to_update_handle);
+        let to_update = self.points_to_push_update.fetch_add(0, Ordering::Relaxed);
 
         if current_amount < to_update {
             return;
@@ -95,7 +93,8 @@ impl ProgressObject {
     }
     pub fn set_max(&self, new_max: usize) {
         *self.max.lock().unwrap() = new_max;
-        *self.points_to_push_update.lock().unwrap() = new_max / PROGRESS_UPDATES;
+        self.points_to_push_update
+            .store(new_max / PROGRESS_UPDATES, Ordering::Relaxed);
         info!("points to push update: {}", new_max / PROGRESS_UPDATES);
     }
     pub fn set_size(&self, length: usize) {
diff --git a/desktop/src-tauri/src/lib.rs b/desktop/src-tauri/src/lib.rs
index 20aa24d1..e3ff447c 100644
--- a/desktop/src-tauri/src/lib.rs
+++ b/desktop/src-tauri/src/lib.rs
@@ -22,7 +22,9 @@ use downloads::download_manager::DownloadManager;
 use downloads::download_manager_builder::DownloadManagerBuilder;
 use http::Response;
 use http::{header::*, response::Builder as ResponseBuilder};
-use library::{fetch_game, fetch_game_status, fetch_game_verion_options, fetch_library, uninstall_game, Game};
+use library::{
+    fetch_game, fetch_game_status, fetch_game_verion_options, fetch_library, uninstall_game, Game,
+};
 use log::{debug, info, warn, LevelFilter};
 use log4rs::append::console::ConsoleAppender;
 use log4rs::append::file::FileAppender;
@@ -115,8 +117,8 @@ fn setup(handle: AppHandle) -> AppState<'static> {
     log4rs::init_config(config).unwrap();
 
     let games = HashMap::new();
-    let download_manager = Arc::new(DownloadManagerBuilder::build(handle));
-    let process_manager = Arc::new(Mutex::new(ProcessManager::new()));
+    let download_manager = Arc::new(DownloadManagerBuilder::build(handle.clone()));
+    let process_manager = Arc::new(Mutex::new(ProcessManager::new(handle.clone())));
     let compat_manager = Arc::new(Mutex::new(CompatibilityManager::new()));
 
     debug!("Checking if database is set up");
diff --git a/desktop/src-tauri/src/process/process_commands.rs b/desktop/src-tauri/src/process/process_commands.rs
index ae100405..6b4aba5a 100644
--- a/desktop/src-tauri/src/process/process_commands.rs
+++ b/desktop/src-tauri/src/process/process_commands.rs
@@ -10,7 +10,7 @@ pub fn launch_game(
     let state_lock = state.lock().unwrap();
     let mut process_manager_lock = state_lock.process_manager.lock().unwrap();
 
-    process_manager_lock.launch_game(game_id)?;
+    process_manager_lock.launch_process(game_id)?;
 
     drop(process_manager_lock);
     drop(state_lock);
diff --git a/desktop/src-tauri/src/process/process_manager.rs b/desktop/src-tauri/src/process/process_manager.rs
index a854e7bd..a170f766 100644
--- a/desktop/src-tauri/src/process/process_manager.rs
+++ b/desktop/src-tauri/src/process/process_manager.rs
@@ -2,27 +2,34 @@ use std::{
     collections::HashMap,
     fs::{File, OpenOptions},
     path::{Path, PathBuf},
-    process::{Child, Command},
-    sync::LazyLock,
+    process::{Child, Command, ExitStatus},
+    sync::{Arc, LazyLock, Mutex},
+    thread::spawn,
 };
 
-use log::info;
+use http::version;
+use log::{info, warn};
 use serde::{Deserialize, Serialize};
+use tauri::{AppHandle, Emitter, Manager};
 
 use crate::{
-    db::{GameStatus, DATA_ROOT_DIR},
-    DB,
+    db::{GameStatus, GameTransientStatus, DATA_ROOT_DIR},
+    library::push_game_update,
+    process::process_manager,
+    state::GameStatusManager,
+    AppState, DB,
 };
 
 pub struct ProcessManager<'a> {
     current_platform: Platform,
     log_output_dir: PathBuf,
-    processes: HashMap<String, Child>,
+    processes: HashMap<String, Arc<Mutex<Child>>>,
+    app_handle: AppHandle,
     game_launchers: HashMap<(Platform, Platform), &'a (dyn ProcessHandler + Sync + Send + 'static)>,
 }
 
 impl ProcessManager<'_> {
-    pub fn new() -> Self {
+    pub fn new(app_handle: AppHandle) -> Self {
         let root_dir_lock = DATA_ROOT_DIR.lock().unwrap();
         let log_output_dir = root_dir_lock.join("logs");
         drop(root_dir_lock);
@@ -34,6 +41,7 @@ impl ProcessManager<'_> {
                 Platform::Linux
             },
 
+            app_handle,
             processes: HashMap::new(),
             log_output_dir,
             game_launchers: HashMap::from([
@@ -48,13 +56,13 @@ impl ProcessManager<'_> {
                 ),
                 (
                     (Platform::Linux, Platform::Windows),
-                    &UMULauncher {} as &(dyn ProcessHandler + Sync + Send + 'static)
-                )
+                    &UMULauncher {} as &(dyn ProcessHandler + Sync + Send + 'static),
+                ),
             ]),
         }
     }
 
-    fn process_command(&self, install_dir: &String, raw_command: String) -> (String, Vec<String>) {
+    fn process_command(&self, install_dir: &String, raw_command: String) -> (PathBuf, Vec<String>) {
         let command_components = raw_command.split(" ").collect::<Vec<&str>>();
         let root = command_components[0].to_string();
 
@@ -65,7 +73,51 @@ impl ProcessManager<'_> {
             .into_iter()
             .map(|v| v.to_string())
             .collect();
-        (absolute_exe.to_str().unwrap().to_owned(), args)
+        (absolute_exe, args)
+    }
+
+    fn on_process_finish(&mut self, game_id: String, result: Result<ExitStatus, std::io::Error>) {
+        if !self.processes.contains_key(&game_id) {
+            warn!("process on_finish was called, but game_id is no longer valid. finished with result: {:?}", result);
+            return;
+        }
+
+        info!("process for {} exited with {:?}", game_id, result);
+
+        self.processes.remove(&game_id);
+
+        let mut db_handle = DB.borrow_data_mut().unwrap();
+        db_handle.games.transient_statuses.remove(&game_id);
+
+        let current_state = db_handle.games.statuses.get(&game_id).cloned();
+        if let Some(saved_state) = current_state {
+            match saved_state {
+                GameStatus::SetupRequired {
+                    version_name,
+                    install_dir,
+                } => {
+                    if let Some(exit_code) = result.ok() {
+                        if exit_code.success() {
+                            db_handle.games.statuses.insert(
+                                game_id.clone(),
+                                GameStatus::Installed {
+                                    version_name: version_name.to_string(),
+                                    install_dir: install_dir.to_string(),
+                                },
+                            );
+                        }
+                    }
+                }
+                _ => (),
+            }
+        }
+        drop(db_handle);
+
+        let status = GameStatusManager::fetch_state(&game_id);
+
+        push_game_update(&self.app_handle, game_id.clone(), status);
+
+        // TODO better management
     }
 
     pub fn valid_platform(&self, platform: &Platform) -> Result<bool, String> {
@@ -75,26 +127,36 @@ impl ProcessManager<'_> {
             .contains_key(&(current.clone(), platform.clone())))
     }
 
-    pub fn launch_game(&mut self, game_id: String) -> Result<(), String> {
+    pub fn launch_process(&mut self, game_id: String) -> Result<(), String> {
         if self.processes.contains_key(&game_id) {
             return Err("Game or setup is already running.".to_owned());
         }
 
-        let db_lock = DB.borrow_data().unwrap();
+        let mut db_lock = DB.borrow_data_mut().unwrap();
         let game_status = db_lock
             .games
             .statuses
             .get(&game_id)
             .ok_or("Game not installed")?;
 
-        let GameStatus::Installed {
-            version_name,
-            install_dir,
-        } = game_status
-        else {
-            return Err("Game not installed.".to_owned());
+        let status_metadata: Option<(&String, &String)> = match game_status {
+            GameStatus::Installed {
+                version_name,
+                install_dir,
+            } => Some((version_name, install_dir)),
+            GameStatus::SetupRequired {
+                version_name,
+                install_dir,
+            } => Some((version_name, install_dir)),
+            _ => None,
         };
 
+        if status_metadata.is_none() {
+            return Err("Game has not been downloaded.".to_owned());
+        }
+
+        let (version_name, install_dir) = status_metadata.unwrap();
+
         let game_version = db_lock
             .games
             .versions
@@ -103,10 +165,27 @@ impl ProcessManager<'_> {
             .get(version_name)
             .ok_or("Invalid version name".to_owned())?;
 
-        let (command, args) =
-            self.process_command(install_dir, game_version.launch_command.clone());
+        let raw_command: String = match game_status {
+            GameStatus::Installed {
+                version_name: _,
+                install_dir: _,
+            } => game_version.launch_command.clone(),
+            GameStatus::SetupRequired {
+                version_name: _,
+                install_dir: _,
+            } => game_version.setup_command.clone(),
+            _ => panic!("unreachable code"),
+        };
 
-        info!("launching process {} in {}", command, install_dir);
+        let (command, args) = self.process_command(install_dir, raw_command);
+
+        let target_current_dir = command.parent().unwrap().to_str().unwrap();
+
+        info!(
+            "launching process {} in {}",
+            command.to_str().unwrap(),
+            target_current_dir
+        );
 
         let current_time = chrono::offset::Local::now();
         let mut log_file = OpenOptions::new()
@@ -132,8 +211,6 @@ impl ProcessManager<'_> {
             )))
             .map_err(|v| v.to_string())?;
 
-        info!("opened log file for {}", command);
-
         let current_platform = self.current_platform.clone();
         let target_platform = game_version.platform.clone();
 
@@ -143,17 +220,53 @@ impl ProcessManager<'_> {
             .ok_or("Invalid version for this platform.")
             .map_err(|e| e.to_string())?;
 
-        let launch_process = game_launcher.launch_game(
+        let launch_process = game_launcher.launch_process(
             &game_id,
             version_name,
-            command,
+            command.to_str().unwrap().to_owned(),
             args,
-            install_dir,
+            &target_current_dir.to_string(),
             log_file,
             error_file,
         )?;
 
-        self.processes.insert(game_id, launch_process);
+        let launch_process_handle = Arc::new(Mutex::new(launch_process));
+
+        db_lock
+            .games
+            .transient_statuses
+            .insert(game_id.clone(), GameTransientStatus::Running {});
+
+        push_game_update(
+            &self.app_handle,
+            game_id.clone(),
+            (None, Some(GameTransientStatus::Running {})),
+        );
+
+        let wait_thread_handle = launch_process_handle.clone();
+        let wait_thread_apphandle = self.app_handle.clone();
+        let wait_thread_game_id = game_id.clone();
+
+        spawn(move || {
+            let mut child_handle = wait_thread_handle.lock().unwrap();
+            let result: Result<ExitStatus, std::io::Error> = child_handle.wait();
+
+            let app_state = wait_thread_apphandle.state::<Mutex<AppState>>();
+            let app_state_handle = app_state.lock().unwrap();
+
+            let mut process_manager_handle = app_state_handle.process_manager.lock().unwrap();
+            process_manager_handle.on_process_finish(wait_thread_game_id, result);
+
+            // As everything goes out of scope, they should get dropped
+            // But just to explicit about it
+            drop(process_manager_handle);
+            drop(app_state_handle);
+            drop(child_handle);
+        });
+
+        self.processes.insert(game_id, launch_process_handle);
+
+        info!("finished spawning process");
 
         Ok(())
     }
@@ -166,13 +279,13 @@ pub enum Platform {
 }
 
 pub trait ProcessHandler: Send + 'static {
-    fn launch_game(
+    fn launch_process(
         &self,
         game_id: &String,
         version_name: &String,
         command: String,
         args: Vec<String>,
-        install_dir: &String,
+        current_dir: &String,
         log_file: File,
         error_file: File,
     ) -> Result<Child, String>;
@@ -180,18 +293,18 @@ pub trait ProcessHandler: Send + 'static {
 
 struct NativeGameLauncher;
 impl ProcessHandler for NativeGameLauncher {
-    fn launch_game(
+    fn launch_process(
         &self,
         game_id: &String,
         version_name: &String,
         command: String,
         args: Vec<String>,
-        install_dir: &String,
+        current_dir: &String,
         log_file: File,
         error_file: File,
     ) -> Result<Child, String> {
         Command::new(command)
-            .current_dir(install_dir)
+            .current_dir(current_dir)
             .stdout(log_file)
             .stderr(error_file)
             .args(args)
@@ -202,13 +315,13 @@ impl ProcessHandler for NativeGameLauncher {
 
 struct UMULauncher;
 impl ProcessHandler for UMULauncher {
-    fn launch_game(
+    fn launch_process(
         &self,
         game_id: &String,
         version_name: &String,
         command: String,
         args: Vec<String>,
-        install_dir: &String,
+        current_dir: &String,
         log_file: File,
         error_file: File,
     ) -> Result<Child, String> {
diff --git a/desktop/types.ts b/desktop/types.ts
index 5702594c..63e2f598 100644
--- a/desktop/types.ts
+++ b/desktop/types.ts
@@ -52,6 +52,7 @@ export enum GameStatusEnum {
   Updating = "Updating",
   Uninstalling = "Uninstalling",
   SetupRequired = "SetupRequired",
+  Running = "Running"
 }
 
 export type GameStatus = {