// This file is part of Deja-vu. // // Deja-vu is free software: you can redistribute it and/or modify it // under the terms of the GNU General Public License as published by // the Free Software Foundation, either version 3 of the License, or // (at your option) any later version. // // Deja-vu is distributed in the hope that it will be useful, but // WITHOUT ANY WARRANTY; without even the implied warranty of // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. // See the GNU General Public License for more details. // // You should have received a copy of the GNU General Public License // along with Deja-vu. If not, see . // use std::{ collections::VecDeque, path::PathBuf, process::{Command, Stdio}, sync::{Arc, Mutex}, thread::{self, sleep, JoinHandle}, time::Duration, }; use color_eyre::Result; use tokio::sync::mpsc; use tracing::info; use crate::{ action::Action, components::popup, system::{disk, diskpart}, }; #[derive(Clone, Debug)] pub enum Task { Command(PathBuf, Vec), // (command, args) Diskpart(String), // (script_as_string) ScanDisks, Sleep, UpdateDestDisk(usize), // (disk_index) UpdateDiskList, } #[derive(Debug)] pub struct Tasks { action_tx: mpsc::UnboundedSender, disk_list: Arc>>, handle: Option>, task_list: VecDeque, task_rx: mpsc::UnboundedReceiver, // Used to forward Actions from Tasks to App task_tx: mpsc::UnboundedSender, // Used to forward Actions from Tasks to App } impl Tasks { pub fn new( action_tx: mpsc::UnboundedSender, disk_list_arc: Arc>>, ) -> Self { let (task_tx, task_rx) = mpsc::unbounded_channel(); Tasks { action_tx, disk_list: disk_list_arc, handle: None, task_list: VecDeque::new(), task_rx, task_tx, } } pub fn add(&mut self, task: Task) { info!("Adding task: {:?}", &task); self.task_list.push_back(task); } pub fn idle(&self) -> bool { self.handle.is_none() } pub fn poll(&mut self) -> Result<()> { // Forward any actions to main app if let Ok(action) = self.task_rx.try_recv() { let result = self.action_tx.send(action.clone()); assert!(result.is_ok(), "Failed to send Action: {action:?}"); } // Check status of current task (if one is running). // NOTE: Action::NextScreen is sent once all tasks are complete if let Some(handle) = self.handle.take() { if handle.is_finished() { if self.task_list.is_empty() { // No tasks remain self.task_tx.send(Action::NextScreen)?; } else { // Start next task self.start()?; } } else { // Task not complete, return handle self.handle = Some(handle); } } else if !self.task_list.is_empty() { // No current task but one is available self.start()?; } Ok(()) } pub fn start(&mut self) -> Result<()> { if let Some(task) = self.task_list.pop_front() { let task_str = format!("{task:?}"); let task_tx = self.task_tx.clone(); match task { Task::Command(ref cmd_path, ref cmd_args) => { let cmd_path = cmd_path.clone(); let cmd_args = cmd_args.clone(); if cfg!(windows) { self.handle = Some(thread::spawn(move || { let result = Command::new(cmd_path) .args(cmd_args) .stdout(Stdio::piped()) .output(); if let Some(action) = match result { Ok(output) => { if output.status.success() { None } else { // Command returned an error status let mut msg = String::new(); if let Ok(stdout) = String::from_utf8(output.stdout) { msg = String::from(stdout.trim()); } if msg.is_empty() { msg = String::from("Generic error"); } Some(Action::Error(format!("Command failed: {msg}",))) } } Err(err) => { Some(Action::Error(format!("Failed to run command: {err:?}"))) } } { let msg = format!("{:?}", &action); let result = task_tx.send(action); assert!(result.is_ok(), "Failed to send Action: {msg}"); } })); } else { // Simulate task if not running under Windows self.handle = Some(thread::spawn(|| sleep(Duration::from_millis(250)))); } } Task::Diskpart(ref script) => { if cfg!(windows) { let script = String::from(script); self.handle = Some(thread::spawn(move || { let output = diskpart::run_script_raw(script.as_str()); if !output.status.success() && task_tx .send(Action::Error(String::from( "Diskpart script returned an error", ))) .is_err() { panic!("Failed to send Action: {task_str:?}"); } })); } else { // Simulate task if not running under Windows self.handle = Some(thread::spawn(|| sleep(Duration::from_millis(250)))); } } Task::ScanDisks => { let disk_list_arc = self.disk_list.clone(); // Queue UpdateDiskList for various components self.add(Task::UpdateDiskList); self.handle = Some(thread::spawn(move || { let mut disks = disk_list_arc.lock().unwrap(); *disks = disk::get_disks(); })); } Task::Sleep => { self.handle = Some(thread::spawn(|| sleep(Duration::from_millis(250)))); } Task::UpdateDestDisk(index) => { self.action_tx.send(Action::DisplayPopup( popup::Type::Info, String::from("Refreshing disk info"), ))?; // Queue UpdateDiskList for various components self.add(Task::Sleep); self.add(Task::UpdateDiskList); // Update destination disk ~in-place let disk_list_arc = self.disk_list.clone(); self.handle = Some(thread::spawn(move || { let mut disks = disk_list_arc.lock().unwrap(); let old_disk = &mut disks[index]; disks[index] = disk::refresh_disk_info(old_disk); })); } Task::UpdateDiskList => { let disks = self.disk_list.lock().unwrap(); let disks_copy = disks.clone(); let action_tx = self.action_tx.clone(); self.handle = Some(thread::spawn(move || { if let Err(err) = action_tx.send(Action::UpdateDiskList(disks_copy)) { panic!("Failed to send Action: {err:?}"); } })); } } } Ok(()) } }