feat(agent): add status code, stderr and stdout to prepare and run

Co-authored-by: Matéo Fernandez <[email protected]> Signed-off-by: Martin Moreira de Jesus <[email protected]>
virt-do · Apr 23, 2024 · 9cf5fe3 · 9cf5fe3
1 parent ba094d6
commit 9cf5fe3
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 18 deletions.
diff --git a/src/agent/src/agents/mod.rs b/src/agent/src/agents/mod.rs
@@ -3,9 +3,16 @@ use serde::Deserialize;
 
 pub mod rust;
 
+#[derive(Debug, Clone)]
+pub struct AgentOutput {
+    pub exit_code: i32,
+    pub stdout: String,
+    pub stderr: String,
+}
+
 pub trait Agent {
-    fn prepare(&self) -> AgentResult<()>;
-    fn run(&self) -> AgentResult<()>;
+    fn prepare(&self) -> AgentResult<AgentOutput>;
+    fn run(&self) -> AgentResult<AgentOutput>;
 }
 
 #[derive(Debug, Clone, Deserialize)]

diff --git a/src/agent/src/agents/rust.rs b/src/agent/src/agents/rust.rs
@@ -1,5 +1,5 @@
-use super::Agent;
-use crate::{workload, AgentResult};
+use super::{Agent, AgentOutput};
+use crate::{workload, AgentError, AgentResult};
 use rand::distributions::{Alphanumeric, DistString};
 use serde::Deserialize;
 use std::{fs::create_dir_all, process::Command};
@@ -22,7 +22,7 @@ pub struct RustAgent {
 }
 
 impl RustAgent {
-    fn build(&self, function_dir: &String) -> AgentResult<String> {
+    fn build(&self, function_dir: &String) -> AgentResult<AgentOutput> {
         if self.rust_config.build.release {
             let output = Command::new("cargo")
                 .arg("build")
@@ -31,15 +31,23 @@ impl RustAgent {
                 .output()
                 .expect("Failed to build function");
 
-            Ok(std::str::from_utf8(&output.stdout).unwrap().to_string())
+            Ok(AgentOutput {
+                exit_code: output.status.code().unwrap(),
+                stdout: std::str::from_utf8(&output.stdout).unwrap().to_string(),
+                stderr: std::str::from_utf8(&output.stderr).unwrap().to_string(),
+            })
         } else {
             let output = Command::new("cargo")
                 .arg("build")
                 .current_dir(function_dir)
                 .output()
                 .expect("Failed to build function");
 
-            Ok(std::str::from_utf8(&output.stdout).unwrap().to_string())
+            Ok(AgentOutput {
+                exit_code: output.status.code().unwrap(),
+                stdout: std::str::from_utf8(&output.stdout).unwrap().to_string(),
+                stderr: std::str::from_utf8(&output.stderr).unwrap().to_string(),
+            })
         }
     }
 }
@@ -57,7 +65,7 @@ impl From<workload::config::Config> for RustAgent {
 }
 
 impl Agent for RustAgent {
-    fn prepare(&self) -> AgentResult<()> {
+    fn prepare(&self) -> AgentResult<AgentOutput> {
         let code = std::fs::read_to_string(&self.rust_config.build.source_code_path).unwrap();
 
         let function_dir = format!(
@@ -87,7 +95,14 @@ impl Agent for RustAgent {
 
         let result = self.build(&function_dir)?;
 
-        println!("{}", result);
+        if result.exit_code != 0 {
+            println!("Build failed: {:?}", result);
+            return Err(AgentError::BuildFailed(AgentOutput {
+                exit_code: result.exit_code,
+                stdout: result.stdout,
+                stderr: result.stderr,
+            }));
+        }
 
         // Copy the binary to /tmp, we could imagine a more complex scenario where we would put this in an artifact repository (like S3)
         let binary_path = match self.rust_config.build.release {
@@ -109,18 +124,29 @@ impl Agent for RustAgent {
 
         std::fs::remove_dir_all(&function_dir).expect("Unable to remove directory");
 
-        Ok(())
+        Ok(AgentOutput {
+            exit_code: result.exit_code,
+            stdout: "Build successful".to_string(),
+            stderr: "".to_string(),
+        })
     }
 
-    fn run(&self) -> AgentResult<()> {
+    fn run(&self) -> AgentResult<AgentOutput> {
         let output = Command::new(format!("/tmp/{}", self.workload_config.workload_name))
             .output()
             .expect("Failed to run function");
 
-        println!("{}", std::str::from_utf8(&output.stdout).unwrap());
+        let agent_output = AgentOutput {
+            exit_code: output.status.code().unwrap(),
+            stdout: std::str::from_utf8(&output.stdout).unwrap().to_string(),
+            stderr: std::str::from_utf8(&output.stderr).unwrap().to_string(),
+        };
 
-        println!("{}", std::str::from_utf8(&output.stderr).unwrap());
+        if !output.status.success() {
+            println!("Run failed: {:?}", agent_output);
+            return Err(AgentError::BuildFailed(agent_output));
+        }
 
-        Ok(())
+        Ok(agent_output)
     }
 }
diff --git a/src/agent/src/lib.rs b/src/agent/src/lib.rs
@@ -1,3 +1,5 @@
+use agents::AgentOutput;
+
 mod agents;
 pub mod workload {
     pub mod config;
@@ -8,6 +10,7 @@ pub mod workload {
 pub enum AgentError {
     OpenConfigFileError(std::io::Error),
     ParseConfigError(toml::de::Error),
+    BuildFailed(AgentOutput),
 }
 
 pub type AgentResult<T> = Result<T, AgentError>;
diff --git a/src/agent/src/workload/runner.rs b/src/agent/src/workload/runner.rs
@@ -21,14 +21,17 @@ impl Runner {
     }
 
     pub fn run(&self) -> AgentResult<()> {
-        match self.config.action {
+        let result = match self.config.action {
             Action::Prepare => self.agent.prepare()?,
             Action::Run => self.agent.run()?,
             Action::PrepareAndRun => {
-                self.agent.prepare()?;
-                self.agent.run()?;
+                let res = self.agent.prepare()?;
+                println!("Prepare result {:?}", res);
+                self.agent.run()?
             }
-        }
+        };
+
+        println!("Result: {:?}", result);
 
         Ok(())
     }