diff --git a/boot_diags/src/app.rs b/boot_diags/src/app.rs index 78d3725..c8897e6 100644 --- a/boot_diags/src/app.rs +++ b/boot_diags/src/app.rs @@ -295,9 +295,15 @@ impl App { self.last_tick_key_events.drain(..); // Check background task(s) match self.cur_mode { + Mode::BootDiags => { + // Check result of background task (if finished) + if let Some(handle) = self.tasks.poll()? { + // TODO: Impl logic + } + } Mode::ScanDisks => { - // Check background task - self.tasks.poll()?; // Once all are complete Action::NextScreen is sent + // Check background task (Once all are complete Action::NextScreen is sent) + self.tasks.poll()?; } _ => {} } diff --git a/core/src/tasks.rs b/core/src/tasks.rs index 50bca0d..09e8cb7 100644 --- a/core/src/tasks.rs +++ b/core/src/tasks.rs @@ -43,11 +43,17 @@ pub enum Task { UpdateDiskList, } +#[derive(Debug)] +pub struct TaskHandle { + task: Task, + handle: JoinHandle<()>, +} + #[derive(Debug)] pub struct Tasks { action_tx: mpsc::UnboundedSender, disk_list: Arc>>, - handle: Option>, + 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 @@ -78,7 +84,8 @@ impl Tasks { self.handle.is_none() } - pub fn poll(&mut self) -> Result<()> { + pub fn poll(&mut self) -> Result> { + let mut return_handle: Option = None; // Forward any actions to main app if let Ok(action) = self.task_rx.try_recv() { let result = self.action_tx.send(action.clone()); @@ -87,8 +94,9 @@ impl Tasks { // 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 let Some(task_handle) = self.handle.take() { + if task_handle.handle.is_finished() { + return_handle = Some(task_handle); if self.task_list.is_empty() { // No tasks remain self.task_tx.send(Action::NextScreen)?; @@ -98,78 +106,28 @@ impl Tasks { } } else { // Task not complete, return handle - self.handle = Some(handle); + self.handle = Some(task_handle); } } else if !self.task_list.is_empty() { // No current task but one is available self.start()?; } - Ok(()) + Ok(return_handle) } 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::Command(cmd_path, cmd_args) => { + self.handle = Some(run_task_command( + cmd_path.clone(), + cmd_args.clone(), + task_tx, + )); } - 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::Diskpart(script) => { + self.handle = Some(run_task_diskpart(&script, task_tx)); } Task::ScanDisks => { let disk_list_arc = self.disk_list.clone(); @@ -177,13 +135,19 @@ impl Tasks { // 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(); - })); + self.handle = Some(TaskHandle { + task: task.clone(), + handle: 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)))); + self.handle = Some(TaskHandle { + task: task.clone(), + handle: thread::spawn(|| sleep(Duration::from_millis(250))), + }); } Task::UpdateDestDisk(index) => { self.action_tx.send(Action::DisplayPopup( @@ -197,24 +161,96 @@ impl Tasks { // 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); - })); + self.handle = Some(TaskHandle { + task: task.clone(), + handle: 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:?}"); - } - })); + self.handle = Some(TaskHandle { + handle: thread::spawn(move || { + if let Err(err) = action_tx.send(Action::UpdateDiskList(disks_copy)) { + panic!("Failed to send Action: {err:?}"); + } + }), + task: task.clone(), + }); } } } Ok(()) } } + +fn run_task_command( + cmd_path: PathBuf, + cmd_args: Vec, + task_tx: mpsc::UnboundedSender, +) -> TaskHandle { + let task = Task::Command(cmd_path.clone(), cmd_args.clone()); + let handle = if cfg!(windows) { + 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 + thread::spawn(|| sleep(Duration::from_millis(250))) + }; + TaskHandle { handle, task } +} + +fn run_task_diskpart(script: &str, task_tx: mpsc::UnboundedSender) -> TaskHandle { + let task = Task::Diskpart(String::from(script)); + let task_str = format!("{:?}", &task); + let handle = if cfg!(windows) { + let script = String::from(script); + thread::spawn(move || { + let output = diskpart::run_script_raw(&script); + 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 + thread::spawn(|| sleep(Duration::from_millis(250))) + }; + TaskHandle { handle, task } +}