From 61e539e0715a559b04359e0ef0fc7545928f5ec5 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 16 Nov 2024 04:09:56 -0800 Subject: [PATCH] Refactored to add command and diskpart script support Replaced previous background task setup to ensure tasks are run sequentially and not in parallel --- src/action.rs | 2 +- src/app.rs | 255 ++++++++++++++++++++++++++------------- src/components/footer.rs | 6 +- src/components/left.rs | 3 +- src/components/right.rs | 2 +- src/system/diskpart.rs | 2 +- 6 files changed, 178 insertions(+), 92 deletions(-) diff --git a/src/action.rs b/src/action.rs index 4285a7c..66cc37f 100644 --- a/src/action.rs +++ b/src/action.rs @@ -37,7 +37,7 @@ pub enum Action { Select(Option, Option), // indicies for (source, dest) or (boot, os) SelectDriver(Driver), SelectTableType(PartitionTableType), - UpdateDestDisk, + UpdateDestDisk(usize), // (disk_index) UpdateDiskList(Vec), // Screens DismissPopup, diff --git a/src/app.rs b/src/app.rs index 3f1e058..512dc7b 100644 --- a/src/app.rs +++ b/src/app.rs @@ -14,9 +14,11 @@ // along with Deja-vu. If not, see . // use std::{ + collections::VecDeque, env, iter::zip, path::PathBuf, + process::{Command, Stdio}, sync::{Arc, Mutex}, thread::{self, sleep, JoinHandle}, time::Duration, @@ -40,12 +42,131 @@ use crate::{ config::Config, system::{ disk::{get_disks, refresh_disk_info, Disk, PartitionTableType}, - diskpart::build_dest_format_script, + diskpart::{build_dest_format_script, run_script_raw}, drivers::{self, Driver}, }, tui::{Event, Tui}, }; +#[derive(Debug)] +pub struct Tasks { + action_tx: mpsc::UnboundedSender, + disk_list: Arc>>, + handle: Option>, + task_list: VecDeque, + task_rx: mpsc::UnboundedReceiver, + task_tx: mpsc::UnboundedSender, +} + +impl Tasks { + pub fn add(&mut self, task: Action) { + self.task_list.push_back(task); + } + + pub fn idle(&self) -> bool { + self.handle.is_none() + } + + pub fn poll(&mut self) -> Result<()> { + debug!("Polling tasks {:?} // {:?}", &self.handle, &self.task_list); + + // Forward any errors (actions) to main app + if let Ok(action) = self.task_rx.try_recv() { + let result = self.action_tx.send(action.clone()); + if result.is_err() { + panic!("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() { + info!("Starting task: {task:?}"); + let task_str = format!("{task:?}"); + let task_tx = self.task_tx.clone(); + match task { + Action::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 output = Command::new(cmd_path) + .args(cmd_args) + .stdout(Stdio::piped()) + .output(); + if output.is_err() { + if let Err(_) = task_tx + .send(Action::Error(String::from("Failed to run command"))) + { + panic!("Failed to send Action: {task_str:?}"); + } + } + })); + } else { + // Simulate task if not running under Windows + self.handle = Some(thread::spawn(|| sleep(Duration::from_secs(1)))); + } + } + Action::Diskpart(ref script) => { + if cfg!(windows) { + let script = String::from(script); + self.handle = Some(thread::spawn(move || { + let output = run_script_raw(script.as_str()); + if !output.status.success() { + if let Err(_) = task_tx.send(Action::Error(String::from( + "Diskpart script returned an error", + ))) { + panic!("Failed to send Action: {task_str:?}"); + } + } + })); + } else { + // Simulate task if not running under Windows + self.handle = Some(thread::spawn(|| sleep(Duration::from_secs(1)))); + } + } + Action::ScanDisks => { + let disk_list_arc = self.disk_list.clone(); + self.handle = Some(thread::spawn(move || { + let mut disks = disk_list_arc.lock().unwrap(); + *disks = get_disks(); + })) + } + Action::UpdateDestDisk(index) => { + let disk_list_arc = self.disk_list.clone(); + self.handle = Some(thread::spawn(move || { + let mut disks = disk_list_arc.lock().unwrap(); + refresh_disk_info(&mut disks[index]); + })); + } + _ => {} + } + } + Ok(()) + } +} + pub struct App { // TUI action_rx: mpsc::UnboundedReceiver, @@ -68,7 +189,7 @@ pub struct App { prev_mode: Mode, selections: Vec>, table_type: Option, - task_handles: Vec>, + tasks: Tasks, } #[derive(Default, Debug, Copy, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -81,6 +202,7 @@ pub enum Mode { Confirm, PreClone, Clone, + UpdateDestDisk, SelectParts, PostClone, Done, @@ -90,6 +212,17 @@ pub enum Mode { impl App { pub fn new(tick_rate: f64, frame_rate: f64) -> Result { let (action_tx, action_rx) = mpsc::unbounded_channel(); + let (task_tx, task_rx) = mpsc::unbounded_channel(); + let disk_list = Arc::new(Mutex::new(Vec::new())); + let mut tasks = Tasks { + action_tx: action_tx.clone(), + disk_list: disk_list.clone(), + handle: None, + task_list: VecDeque::new(), + task_rx, + task_tx, + }; + tasks.add(Action::ScanDisks); Ok(Self { // TUI action_rx, @@ -112,14 +245,14 @@ impl App { cur_mode: Mode::ScanDisks, disk_index_dest: None, disk_index_source: None, - disk_list: Arc::new(Mutex::new(Vec::new())), + disk_list, driver: None, part_index_boot: None, part_index_os: None, prev_mode: Mode::ScanDisks, selections: vec![None, None], table_type: None, - task_handles: Vec::new(), + tasks, }) } @@ -137,7 +270,8 @@ impl App { (Mode::SelectDisks, Mode::Confirm) => Mode::SelectTableType, (Mode::SelectTableType, Mode::Confirm) => Mode::PreClone, (_, Mode::PreClone) => Mode::Clone, - (_, Mode::Clone) => Mode::SelectParts, + (_, Mode::Clone) => Mode::UpdateDestDisk, + (_, Mode::UpdateDestDisk) => Mode::SelectParts, (Mode::SelectParts, Mode::Confirm) => Mode::PostClone, (_, Mode::PostClone) => Mode::Done, (_, Mode::Done) => Mode::Done, @@ -163,8 +297,6 @@ impl App { } pub async fn run(&mut self) -> Result<()> { - let disk_list_arc = Arc::clone(&self.disk_list); - self.task_handles.push(lazy_get_disks(disk_list_arc)); let mut tui = Tui::new()? // .mouse(true) // uncomment this line to enable mouse support .tick_rate(self.tick_rate) @@ -255,21 +387,13 @@ impl App { Action::Tick => { self.last_tick_key_events.drain(..); match self.cur_mode { - Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => { + Mode::ScanDisks + | Mode::PreClone + | Mode::Clone + | Mode::UpdateDestDisk + | Mode::PostClone => { // Check background task - if let Ok(_) = &self.disk_list.try_lock() { - if self.task_handles.is_empty() { - // All tasks complete - self.action_tx.send(Action::NextScreen)?; - } else { - // Tasks remain, check if current is complete - if let Some(handle) = self.task_handles.first() { - if handle.is_finished() { - self.task_handles.remove(0); - } - } - } - } + self.tasks.poll()?; // Once all are complete Action::NextScreen is sent } _ => {} } @@ -279,10 +403,13 @@ impl App { Action::Resume => self.should_suspend = false, Action::ClearScreen => tui.terminal.clear()?, Action::Command(ref cmd_path, ref cmd_args) => self - .task_handles - .push(lazy_command(&cmd_path, cmd_args.clone())), - Action::Diskpart(ref script) => { - self.task_handles.push(lazy_diskpart(&script.to_owned())) + .tasks + .add(Action::Command(cmd_path.clone(), cmd_args.clone())), + Action::Diskpart(ref script) => self.tasks.add(Action::Diskpart(script.clone())), + Action::Error(ref msg) => { + self.action_tx + .send(Action::DisplayPopup(popup::Type::Error, msg.clone()))?; + self.action_tx.send(Action::SetMode(Mode::Failed))?; } Action::InstallDriver => { self.action_tx.send(Action::SetMode(Mode::InstallDrivers))? @@ -307,11 +434,7 @@ impl App { self.action_tx.send(Action::SetMode(mode))?; } } - Action::ScanDisks => { - let disk_list_arc = Arc::clone(&self.disk_list); - self.task_handles.push(lazy_get_disks(disk_list_arc)); - self.action_tx.send(Action::SetMode(Mode::ScanDisks))?; - } + Action::ScanDisks => self.action_tx.send(Action::SetMode(Mode::ScanDisks))?, Action::Select(one, two) => { match self.cur_mode { Mode::SelectDisks => { @@ -333,6 +456,9 @@ impl App { match new_mode { Mode::ScanDisks => { self.prev_mode = self.cur_mode; + if self.tasks.idle() { + self.tasks.add(Action::ScanDisks); + } } Mode::SelectDisks | Mode::SelectParts => { let disk_list = self.disk_list.lock().unwrap(); @@ -365,7 +491,11 @@ impl App { self.config.clone_app_path.clone(), Vec::new(), ))?; - self.action_tx.send(Action::UpdateDestDisk)?; + } + Mode::UpdateDestDisk => { + if let Some(dest_index) = self.disk_index_dest { + self.action_tx.send(Action::UpdateDestDisk(dest_index))?; + } } Mode::PostClone => { self.action_tx.send(Action::DisplayPopup( @@ -378,11 +508,9 @@ impl App { match env::var("SYSTEMROOT") { Ok(path) => path, Err(_) => { - self.action_tx.send(Action::DisplayPopup( - popup::Type::Error, - String::from("ERROR\n\n\nFailed to find SYSTEMROOT"), - ))?; - self.action_tx.send(Action::SetMode(Mode::Failed))?; + self.action_tx.send(Action::Error(String::from( + "ERROR\n\n\nFailed to find SYSTEMROOT", + )))?; return Ok(()); } } @@ -402,13 +530,9 @@ impl App { // Safety check if letter_boot.is_empty() || letter_os.is_empty() { - self.action_tx.send(Action::DisplayPopup( - popup::Type::Error, - String::from( - "ERROR\n\n\nFailed to get drive letters for the destination", - ), - ))?; - self.action_tx.send(Action::SetMode(Mode::Failed))?; + self.action_tx.send(Action::Error(String::from( + "ERROR\n\n\nFailed to get drive letters for the destination", + )))?; return Ok(()); } @@ -489,13 +613,7 @@ impl App { _ => {} } } - Action::UpdateDestDisk => { - let disk_list_arc = Arc::clone(&self.disk_list); - if let Some(dest_index) = self.disk_index_dest { - self.task_handles - .push(lazy_update_dest_disk(disk_list_arc, dest_index)); - } - } + Action::UpdateDestDisk(index) => self.tasks.add(Action::UpdateDestDisk(index)), _ => {} } for component in self.components.iter_mut() { @@ -587,40 +705,3 @@ fn get_chunks(r: Rect) -> Vec { // Done chunks } - -fn lazy_command(cmd_path: &PathBuf, cmd_args: Vec) -> JoinHandle<()> { - info!("Running Command: {cmd_path:?} {cmd_args:?}"); - if cfg!(windows) { - // TODO: FIXME - thread::spawn(|| sleep(Duration::from_secs(1))) - } else { - thread::spawn(|| sleep(Duration::from_secs(1))) - } -} - -fn lazy_diskpart(script: &str) -> JoinHandle<()> { - if cfg!(windows) { - // TODO: FIXME - thread::spawn(|| sleep(Duration::from_secs(1))) - } else { - info!("Running (lazy) Diskpart: {:?}", &script); - thread::spawn(|| sleep(Duration::from_secs(1))) - } -} - -fn lazy_get_disks(disk_list_arc: Arc>>) -> JoinHandle<()> { - thread::spawn(move || { - let mut disks = disk_list_arc.lock().unwrap(); - *disks = get_disks(); - }) -} - -fn lazy_update_dest_disk( - disk_list_arc: Arc>>, - dest_index: usize, -) -> JoinHandle<()> { - thread::spawn(move || { - let mut disks = disk_list_arc.lock().unwrap(); - refresh_disk_info(&mut disks[dest_index]); - }) -} diff --git a/src/components/footer.rs b/src/components/footer.rs index afd5f04..0bfd5fa 100644 --- a/src/components/footer.rs +++ b/src/components/footer.rs @@ -58,7 +58,11 @@ impl Component for Footer { Action::SetMode(new_mode) => { self.footer_text = match new_mode { Mode::ScanDisks => String::from("(q) to quit"), - Mode::PreClone | Mode::Clone | Mode::PostClone | Mode::SelectParts => { + Mode::PreClone + | Mode::Clone + | Mode::UpdateDestDisk + | Mode::SelectParts + | Mode::PostClone => { String::from("(Enter) to select / (s) to start over / (q) to quit") } Mode::SelectDisks => String::from( diff --git a/src/components/left.rs b/src/components/left.rs index 0bde6b1..5b26a0a 100644 --- a/src/components/left.rs +++ b/src/components/left.rs @@ -209,7 +209,7 @@ impl Component for Left { self.selections[1] = None; self.title_text = String::from("Select Partition Table Type"); } - (_, Mode::PreClone | Mode::Clone | Mode::PostClone) => { + (_, Mode::PreClone | Mode::Clone | Mode::UpdateDestDisk | Mode::PostClone) => { self.title_text = String::from("Processing"); } (_, Mode::SelectParts) => { @@ -276,6 +276,7 @@ impl Component for Left { Mode::ScanDisks | Mode::PreClone | Mode::Clone + | Mode::UpdateDestDisk | Mode::PostClone | Mode::Done | Mode::Failed => { diff --git a/src/components/right.rs b/src/components/right.rs index 051f73b..4081d58 100644 --- a/src/components/right.rs +++ b/src/components/right.rs @@ -158,7 +158,7 @@ impl Component for Right { .areas(area); // Title - let title_text = String::from("Info?"); + let title_text = String::from("Info"); let title = Paragraph::new(Line::from(title_text).centered()) .block(Block::default().borders(Borders::NONE)); frame.render_widget(title, title_area); diff --git a/src/system/diskpart.rs b/src/system/diskpart.rs index 79fcc36..839bbcc 100644 --- a/src/system/diskpart.rs +++ b/src/system/diskpart.rs @@ -343,7 +343,7 @@ pub fn run_script_raw(script: &str) -> Output { .args(["/s", format!("{}", script_path.display()).as_str()]) .stdout(Stdio::piped()) .output() - .expect("Failed to execute command"); + .expect("Failed to execute Diskpart script"); output }