2Shirt/deja-vu
1
0
Fork 0
Code Issues 1 Pull requests Projects Releases Packages Wiki Activity Actions

Refactored to add command and diskpart script support

Browse source 
Replaced previous background task setup to ensure tasks are run
sequentially and not in parallel
This commit is contained in:
2Shirt 2024-11-16 04:09:56 -08:00
parent c9884d5b6d
commit 61e539e071
Signed by: 2Shirt
GPG key ID: 152FAC923B0E132C
6 changed files with 178 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
src/action.rs
View file

@ -37,7 +37,7 @@ pub enum Action {
Select(Option<usize>, Option<usize>), // indicies for (source, dest) or (boot, os)
SelectDriver(Driver),
SelectTableType(PartitionTableType),
UpdateDestDisk,
UpdateDestDisk(usize), // (disk_index)
UpdateDiskList(Vec<Disk>),
// Screens
DismissPopup,

255
src/app.rs
View file

@ -14,9 +14,11 @@
// along with Deja-vu. If not, see <https://www.gnu.org/licenses/>.
//
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<Action>,
disk_list: Arc<Mutex<Vec<Disk>>>,
handle: Option<JoinHandle<()>>,
task_list: VecDeque<Action>,
task_rx: mpsc::UnboundedReceiver<Action>,
task_tx: mpsc::UnboundedSender<Action>,
}
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<Action>,
@ -68,7 +189,7 @@ pub struct App {
prev_mode: Mode,
selections: Vec<Option<usize>>,
table_type: Option<PartitionTableType>,
task_handles: Vec<JoinHandle<()>>,
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<Self> {
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<Rect> {
// Done
chunks
}
fn lazy_command(cmd_path: &PathBuf, cmd_args: Vec<String>) -> 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<Mutex<Vec<Disk>>>) -> JoinHandle<()> {
thread::spawn(move || {
let mut disks = disk_list_arc.lock().unwrap();
*disks = get_disks();
})
}
fn lazy_update_dest_disk(
disk_list_arc: Arc<Mutex<Vec<Disk>>>,
dest_index: usize,
) -> JoinHandle<()> {
thread::spawn(move || {
let mut disks = disk_list_arc.lock().unwrap();
refresh_disk_info(&mut disks[dest_index]);
})
}

6
src/components/footer.rs
View file

@ -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(

3
src/components/left.rs
View file

@ -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 => {

2
src/components/right.rs
View file

@ -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);

2
src/system/diskpart.rs
View file

@ -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
}