From 4e80badc921d31845ff26bffb5eb21e4d5caee29 Mon Sep 17 00:00:00 2001 From: 2Shirt <2xShirt@gmail.com> Date: Sat, 16 Nov 2024 23:53:20 -0800 Subject: [PATCH] Move fast, break things (round 3?) - Move sequential items to Tasks and leave async in Actions - Move main control of disk_list (Arc) to Tasks This allows updates to happen AFTER I/O events (clone, update dest, etc) - Finally got SelectParts to use the actual state of the disks post-clone - Broke Tasks somewhere - bcdedit is never called despite being added to the VecDeque - Probably more --- src/action.rs | 4 - src/app.rs | 193 +++++-------------------------------- src/components/footer.rs | 2 +- src/components/left.rs | 5 +- src/components/right.rs | 30 +++--- src/main.rs | 1 + src/system/disk.rs | 22 +++-- src/system/diskpart.rs | 3 +- src/tasks.rs | 203 +++++++++++++++++++++++++++++++++++++++ 9 files changed, 262 insertions(+), 201 deletions(-) create mode 100644 src/tasks.rs diff --git a/src/action.rs b/src/action.rs index 66cc37f..ba18a79 100644 --- a/src/action.rs +++ b/src/action.rs @@ -14,7 +14,6 @@ // along with Deja-vu. If not, see . // use serde::{Deserialize, Serialize}; -use std::path::PathBuf; use strum::Display; use crate::{ @@ -29,15 +28,12 @@ use crate::{ #[derive(Debug, Clone, PartialEq, Eq, Display, Serialize, Deserialize)] pub enum Action { // App - Command(PathBuf, Vec), // (command, args) - Diskpart(String), // (script_as_string) InstallDriver, Process, ScanDisks, Select(Option, Option), // indicies for (source, dest) or (boot, os) SelectDriver(Driver), SelectTableType(PartitionTableType), - UpdateDestDisk(usize), // (disk_index) UpdateDiskList(Vec), // Screens DismissPopup, diff --git a/src/app.rs b/src/app.rs index e8638ff..ad58435 100644 --- a/src/app.rs +++ b/src/app.rs @@ -14,14 +14,10 @@ // 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, }; use color_eyre::Result; @@ -41,132 +37,14 @@ use crate::{ }, config::Config, system::{ - disk::{get_disks, refresh_disk_info, Disk, PartitionTableType}, - diskpart::{build_dest_format_script, run_script_raw}, + disk::{Disk, PartitionTableType}, + diskpart::build_dest_format_script, drivers::{self, Driver}, }, + tasks::{Task, Tasks}, 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 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, @@ -202,7 +80,6 @@ pub enum Mode { Confirm, PreClone, Clone, - UpdateDestDisk, SelectParts, PostClone, Done, @@ -212,17 +89,9 @@ 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); + let disk_list_arc = Arc::new(Mutex::new(Vec::new())); + let mut tasks = Tasks::new(action_tx.clone(), disk_list_arc.clone()); + tasks.add(Task::ScanDisks); Ok(Self { // TUI action_rx, @@ -245,7 +114,7 @@ impl App { cur_mode: Mode::ScanDisks, disk_index_dest: None, disk_index_source: None, - disk_list, + disk_list: disk_list_arc, driver: None, part_index_boot: None, part_index_os: None, @@ -270,8 +139,7 @@ impl App { (Mode::SelectDisks, Mode::Confirm) => Mode::SelectTableType, (Mode::SelectTableType, Mode::Confirm) => Mode::PreClone, (_, Mode::PreClone) => Mode::Clone, - (_, Mode::Clone) => Mode::UpdateDestDisk, - (_, Mode::UpdateDestDisk) => Mode::SelectParts, + (_, Mode::Clone) => Mode::SelectParts, (Mode::SelectParts, Mode::Confirm) => Mode::PostClone, (_, Mode::PostClone) => Mode::Done, (_, Mode::Done) => Mode::Done, @@ -387,11 +255,7 @@ impl App { Action::Tick => { self.last_tick_key_events.drain(..); match self.cur_mode { - Mode::ScanDisks - | Mode::PreClone - | Mode::Clone - | Mode::UpdateDestDisk - | Mode::PostClone => { + Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => { // Check background task self.tasks.poll()?; // Once all are complete Action::NextScreen is sent } @@ -402,10 +266,6 @@ impl App { Action::Suspend => self.should_suspend = true, Action::Resume => self.should_suspend = false, Action::ClearScreen => tui.terminal.clear()?, - Action::Command(ref cmd_path, ref cmd_args) => self - .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()))?; @@ -457,14 +317,9 @@ impl App { Mode::ScanDisks => { self.prev_mode = self.cur_mode; if self.tasks.idle() { - self.tasks.add(Action::ScanDisks); + self.tasks.add(Task::ScanDisks); } } - Mode::SelectDisks | Mode::SelectParts => { - let disk_list = self.disk_list.lock().unwrap(); - self.action_tx - .send(Action::UpdateDiskList(disk_list.clone()))?; - } Mode::PreClone => { self.action_tx.send(Action::DisplayPopup( popup::Type::Info, @@ -478,7 +333,7 @@ impl App { let table_type = self.table_type.clone().unwrap(); let diskpart_script = build_dest_format_script(disk.id, &table_type); - self.action_tx.send(Action::Diskpart(diskpart_script))?; + self.tasks.add(Task::Diskpart(diskpart_script)); } } } @@ -487,14 +342,12 @@ impl App { popup::Type::Info, String::from("Running Clone Tool"), ))?; - self.action_tx.send(Action::Command( + self.tasks.add(Task::Command( self.config.clone_app_path.clone(), Vec::new(), - ))?; - } - Mode::UpdateDestDisk => { + )); if let Some(dest_index) = self.disk_index_dest { - self.action_tx.send(Action::UpdateDestDisk(dest_index))?; + self.tasks.add(Task::UpdateDestDisk(dest_index)); } } Mode::PostClone => { @@ -537,7 +390,7 @@ impl App { } // Create boot files - self.action_tx.send(Action::Command( + self.tasks.add(Task::Command( PathBuf::from(format!("{system32}/bcdboot.exe")), vec![ format!("{letter_os}:\\Windows"), @@ -549,12 +402,11 @@ impl App { PartitionTableType::Legacy => "BIOS", }), ], - ))?; + )); // Update boot sector (for legacy setups) if table_type == PartitionTableType::Legacy { - // - self.action_tx.send(Action::Command( + self.tasks.add(Task::Command( PathBuf::from(format!("{system32}/bootsect.exe")), vec![ String::from("/nt60"), @@ -562,7 +414,7 @@ impl App { String::from("/force"), String::from("/mbr"), ], - ))?; + )); } // Lock in safe mode @@ -574,7 +426,7 @@ impl App { format!("{letter_boot}:\\Boot\\BCD") } }; - self.action_tx.send(Action::Command( + self.tasks.add(Task::Command( PathBuf::from(format!("{system32}/bcdedit.exe")), vec![ String::from("/store"), @@ -584,13 +436,13 @@ impl App { String::from("safeboot"), String::from("minimal"), ], - ))?; + )); // Inject driver(s) (if selected) if let Some(driver) = &self.driver { if let Some(os_path) = driver.path.to_str() { let driver_path_str = String::from(os_path); - self.action_tx.send(Action::Command( + self.tasks.add(Task::Command( PathBuf::from(format!("{system32}/dism.exe")), vec![ format!("/image:{letter_os}:\\"), @@ -598,7 +450,7 @@ impl App { format!("/driver:\"{}\"", driver_path_str,), String::from("/recurse"), ], - ))?; + )); } } } @@ -613,7 +465,6 @@ impl App { _ => {} } } - Action::UpdateDestDisk(index) => self.tasks.add(Action::UpdateDestDisk(index)), _ => {} } for component in self.components.iter_mut() { diff --git a/src/components/footer.rs b/src/components/footer.rs index dab12f3..8258e21 100644 --- a/src/components/footer.rs +++ b/src/components/footer.rs @@ -60,7 +60,7 @@ impl Component for Footer { Mode::ScanDisks | Mode::PreClone | Mode::Clone | Mode::PostClone => { String::from("(q) to quit") } - Mode::UpdateDestDisk | Mode::SelectParts => { + Mode::SelectParts => { 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 3772576..62f4689 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::UpdateDestDisk | Mode::PostClone) => { + (_, Mode::PreClone | Mode::Clone | Mode::PostClone) => { self.title_text = String::from("Processing"); } (_, Mode::SelectParts) => { @@ -229,7 +229,7 @@ impl Component for Left { Action::UpdateDiskList(disks) => { info!("Updating disk list"); self.list_disks.set_items(disks); - if self.mode == Mode::SelectParts { + if self.mode == Mode::Clone { if let Some(index) = self.disk_id_dest { if let Some(disk) = self.list_disks.get(index) { self.list_parts.set_items(disk.get_parts()); @@ -277,7 +277,6 @@ 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 f0a7e6b..5eea48d 100644 --- a/src/components/right.rs +++ b/src/components/right.rs @@ -122,24 +122,28 @@ impl Component for Right { Action::UpdateDiskList(disks) => { info!("Updating disk list"); self.list_disks.set_items(disks); - if self.cur_mode == Mode::SelectParts { + if self.cur_mode == Mode::Clone { if let Some(index) = self.selected_disks[1] { if let Some(disk) = self.list_disks.get(index) { self.list_parts.set_items(disk.get_parts()); // Auto-select first partition and highlight likely OS partition - if let Some(table_type) = &self.table_type { - match table_type { - PartitionTableType::Guid => { - if disk.num_parts() >= 3 { - self.selections[0] = Some(0); - self.list_parts.select(2); - } - } - PartitionTableType::Legacy => { - if disk.num_parts() >= 2 { - self.selections[0] = Some(0); - self.list_parts.select(1); + if let Some(index) = self.selected_disks[1] { + if let Some(disk) = self.list_disks.get(index) { + if let Some(table_type) = &self.table_type { + match table_type { + PartitionTableType::Guid => { + if disk.num_parts() >= 3 { + self.selections[0] = Some(0); + self.list_parts.select(2); + } + } + PartitionTableType::Legacy => { + if disk.num_parts() >= 2 { + self.selections[0] = Some(0); + self.list_parts.select(1); + } + } } } } diff --git a/src/main.rs b/src/main.rs index efa0927..72a1b88 100644 --- a/src/main.rs +++ b/src/main.rs @@ -27,6 +27,7 @@ mod config; mod errors; mod logging; mod system; +mod tasks; mod tests; mod tui; diff --git a/src/system/disk.rs b/src/system/disk.rs index 29d44b2..1e55115 100644 --- a/src/system/disk.rs +++ b/src/system/disk.rs @@ -314,19 +314,26 @@ pub fn get_disk_serial_number(id: usize) -> String { serial } -pub fn refresh_disk_info(disk: &mut Disk) { +pub fn refresh_disk_info(disk: &mut Disk) -> Disk { + sleep(Duration::from_millis(5000)); // Wait for clone tool / disk locks + let mut new_disk: Disk; + info!("-- Before: {:?}", &disk.parts); if cfg!(windows) { info!("Refresh disk via Diskpart"); - diskpart::refresh_disk_info(disk); + new_disk = diskpart::refresh_disk_info(disk); } else { info!("Refresh fake disk"); - refresh_fake_disk_info(disk); - sleep(Duration::from_millis(500)); + new_disk = disk.clone(); + new_disk.parts = refresh_fake_disk_info(); + new_disk.generate_descriptions(); } + info!("-------------------------------------------------"); + info!("-- After: {:?}", &new_disk.parts); + new_disk } -fn refresh_fake_disk_info(disk: &mut Disk) { - disk.parts = vec![ +fn refresh_fake_disk_info() -> Vec { + vec![ Partition { id: 1, fs_type: String::from("FAT32"), @@ -352,8 +359,7 @@ fn refresh_fake_disk_info(disk: &mut Disk) { size: 824_340_119_552, ..Default::default() }, - ]; - disk.generate_descriptions(); + ] } /// Misc diff --git a/src/system/diskpart.rs b/src/system/diskpart.rs index 839bbcc..04612c9 100644 --- a/src/system/diskpart.rs +++ b/src/system/diskpart.rs @@ -321,11 +321,12 @@ pub fn parse_partition_details(parts: &mut Vec, contents: &str) { } } -pub fn refresh_disk_info(disk: &Disk) { +pub fn refresh_disk_info(disk: &Disk) -> Disk { info!("Refresh disk info"); let mut disk = get_disk_details(disk.id, disk.size, None); disk.parts = get_partition_details(disk.id, None, None); disk.generate_descriptions(); + disk } /// # Panics diff --git a/src/tasks.rs b/src/tasks.rs new file mode 100644 index 0000000..f43e39f --- /dev/null +++ b/src/tasks.rs @@ -0,0 +1,203 @@ +// 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<()> { + info!("Polling tasks {:?} // {:?}", &self.handle, &self.task_list); + + // Forward any 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()?; + } else { + // No current task or any queued, continue + self.task_tx.send(Action::NextScreen)?; + } + 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 { + 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 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)))); + } + } + 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() { + 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)))); + } + } + 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_secs(1)))); + } + 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(); + self.action_tx.send(Action::UpdateDiskList(disks.clone()))?; + } + } + } + Ok(()) + } +}