fix(agent): wait for stdout/stderr to be sent

fix(agent): add option on process_utils functions Signed-off-by: Mateo Fernandez <[email protected]>
virt-do · May 30, 2024 · 3914d4a · 3914d4a
1 parent 8b115e8
commit 3914d4a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/src/agent/src/agents/mod.rs b/src/agent/src/agents/mod.rs
@@ -91,6 +91,7 @@ mod process_utils {
     pub async fn send_stdout_to_tx(
         stdout: ChildStdout,
         tx: mpsc::Sender<AgentOutput>,
+        stage: Option<Stage>,
     ) -> JoinHandle<()> {
         tokio::spawn(async move {
             let reader = BufReader::new(stdout);
@@ -99,7 +100,7 @@ mod process_utils {
             while let Ok(Some(line)) = reader_lines.next_line().await {
                 let _ = tx
                     .send(AgentOutput {
-                        stage: Stage::Running,
+                        stage: stage.unwrap_or(Stage::Running),
                         stdout: Some(line),
                         stderr: None,
                         exit_code: None,
@@ -113,6 +114,7 @@ mod process_utils {
     pub async fn send_stderr_to_tx(
         stderr: ChildStderr,
         tx: mpsc::Sender<AgentOutput>,
+        stage: Option<Stage>,
     ) -> JoinHandle<()> {
         tokio::spawn(async move {
             let reader = BufReader::new(stderr);
@@ -121,7 +123,7 @@ mod process_utils {
             while let Ok(Some(line)) = reader_lines.next_line().await {
                 let _ = tx
                     .send(AgentOutput {
-                        stage: Stage::Running,
+                        stage: stage.unwrap_or(Stage::Running),
                         stdout: Some(line),
                         stderr: None,
                         exit_code: None,

diff --git a/src/agent/src/agents/rust.rs b/src/agent/src/agents/rust.rs
@@ -1,4 +1,5 @@
 use super::{Agent, AgentOutput};
+use crate::agent::execute_response::Stage;
 use crate::agents::process_utils;
 use crate::{workload, AgentError, AgentResult};
 use async_trait::async_trait;
@@ -117,7 +118,10 @@ impl Agent for RustAgent {
 
         let (tx, rx) = mpsc::channel(10);
         tokio::spawn(async move {
-            let _ = process_utils::send_stderr_to_tx(child.stderr.take().unwrap(), tx.clone()).await.await;
+            let stderr = child.stderr.take().unwrap();
+            let _ = process_utils::send_stderr_to_tx(stderr, tx.clone(), Some(Stage::Building))
+                .await
+                .await;
             let build_result = process_utils::send_exit_status_to_tx(child, tx, false).await;
             // if error in build, short-circuit the execution
             if build_result.is_err() {
@@ -173,12 +177,16 @@ impl Agent for RustAgent {
         let tx_stderr = tx;
 
         tokio::spawn(async move {
-            let _ = process_utils::send_stdout_to_tx(child_stdout, tx_stdout.clone()).await.await;
+            let _ = process_utils::send_stdout_to_tx(child_stdout, tx_stdout.clone(), None)
+                .await
+                .await;
             let _ = process_utils::send_exit_status_to_tx(child, tx_stdout, true).await;
         });
 
         tokio::spawn(async move {
-            process_utils::send_stderr_to_tx(child_stderr, tx_stderr).await;
+            let _ = process_utils::send_stderr_to_tx(child_stderr, tx_stderr, None)
+                .await
+                .await;
         });
 
         Ok(rx)