change: Improved follow behavior in daemon

Nukesor · Jun 6, 2023 · 9c01f8f · 9c01f8f
1 parent 15bd74b
commit 9c01f8f
Showing 1 changed file with 21 additions and 3 deletions.
diff --git a/pueue/src/daemon/network/follow_log.rs b/pueue/src/daemon/network/follow_log.rs
@@ -49,6 +49,24 @@ pub async fn handle_follow(
         }
     };
 
+    // It might be that the task is not yet running.
+    // Ensure that it exists and is started.
+    loop {
+        {
+            let state = state.lock().unwrap();
+            let Some(task) = state.tasks.get(&task_id) else {
+                return Ok(create_failure_message(
+                        "Pueue: The followed task does not exist.",
+                        ));
+            };
+            // The task is running or finished, we can start to follow.
+            if task.is_running() || task.is_done() {
+                break;
+            }
+        }
+        tokio::time::sleep(Duration::from_millis(1000)).await;
+    }
+
     let mut handle = match get_log_file_handle(task_id, pueue_directory) {
         Err(_) => {
             return Ok(create_failure_message(
@@ -103,9 +121,9 @@ pub async fn handle_follow(
         {
             let state = state.lock().unwrap();
             let Some(task) = state.tasks.get(&task_id) else {
-                return Ok(create_success_message(
-                    "Pueue: The followed task has been removed.",
-                ));
+                return Ok(create_failure_message(
+                        "Pueue: The followed task has been removed.",
+                        ));
             };
 
             // The task is done, just close the stream.