Refactor tasks

Retain info about tasks to use in app logic

boot_diags/src/app.rs

@@ -144,7 +144,11 @@ impl App {
     }
 
     pub fn set_boot_mode(&mut self, boot_mode: SafeMode) {
-        info!("Setting boot mode to: {:?}", boot_mode);
+        let new_mode = match boot_mode {
+            SafeMode::Disable => "Normal",
+            SafeMode::Enable => "Safe Mode (minimal)",
+        };
+        info!("Setting boot mode to: {new_mode}");
         let disk_list = self.clone.disk_list.lock().unwrap();
         if let Some(disk_index) = self.clone.disk_index_dest {
             if let Some(disk) = disk_list.get(disk_index) {
@@ -291,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()?;
                         }
                         _ => {}
                     }

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<Action>,
     disk_list: Arc<Mutex<Vec<disk::Disk>>>,
-    handle: Option<JoinHandle<()>>,
+    handle: Option<TaskHandle>,
     task_list: VecDeque<Task>,
     task_rx: mpsc::UnboundedReceiver<Action>, // Used to forward Actions from Tasks to App
     task_tx: mpsc::UnboundedSender<Action>,   // 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<Option<TaskHandle>> {
+        let mut return_handle: Option<TaskHandle> = 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<String>,
+    task_tx: mpsc::UnboundedSender<Action>,
+) -> 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<Action>) -> 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 }
+}