Skip to content

Commit

Permalink
refactor: move rust and Agent trait into agents module and AgentRunne…
Browse files Browse the repository at this point in the history
…r into workload module

Signed-off-by: Matéo Fernandez <[email protected]>
Signed-off-by: Martin Moreira de Jesus <[email protected]>
  • Loading branch information
mfernd authored and mmoreiradj committed Apr 23, 2024
1 parent 30a60b5 commit ba094d6
Show file tree
Hide file tree
Showing 6 changed files with 105 additions and 83 deletions.
23 changes: 23 additions & 0 deletions src/agent/src/agents/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
use crate::AgentResult;
use serde::Deserialize;

pub mod rust;

pub trait Agent {
fn prepare(&self) -> AgentResult<()>;
fn run(&self) -> AgentResult<()>;
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum Language {
Rust,
}

impl std::fmt::Display for Language {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Language::Rust => write!(f, "rust"),
}
}
}
30 changes: 16 additions & 14 deletions src/agent/src/rust.rs → src/agent/src/agents/rust.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
use std::{fs::create_dir_all, process::Command};

use super::Agent;
use crate::{workload, AgentResult};
use rand::distributions::{Alphanumeric, DistString};
use serde::Deserialize;

use crate::{Agent, AgentConfig, AgentResult};
use std::{fs::create_dir_all, process::Command};

#[derive(Deserialize)]
#[serde(rename_all = "kebab-case")]
Expand All @@ -18,7 +17,7 @@ struct RustAgentConfig {
}

pub struct RustAgent {
agent_config: AgentConfig,
workload_config: workload::config::Config,
rust_config: RustAgentConfig,
}

Expand All @@ -43,12 +42,15 @@ impl RustAgent {
Ok(std::str::from_utf8(&output.stdout).unwrap().to_string())
}
}
}

pub fn new_from_config(agent_config: AgentConfig) -> Self {
let rust_config: RustAgentConfig = toml::from_str(&agent_config.config_string).unwrap();
// TODO should change with a TryFrom
impl From<workload::config::Config> for RustAgent {
fn from(workload_config: workload::config::Config) -> Self {
let rust_config: RustAgentConfig = toml::from_str(&workload_config.config_string).unwrap();

RustAgent {
agent_config,
Self {
workload_config,
rust_config,
}
}
Expand Down Expand Up @@ -77,7 +79,7 @@ impl Agent for RustAgent {
version = "0.1.0"
edition = "2018"
"#,
self.agent_config.workload_name
self.workload_config.workload_name
);

std::fs::write(format!("{}/Cargo.toml", &function_dir), cargo_toml)
Expand All @@ -91,17 +93,17 @@ impl Agent for RustAgent {
let binary_path = match self.rust_config.build.release {
true => format!(
"{}/target/release/{}",
&function_dir, self.agent_config.workload_name
&function_dir, self.workload_config.workload_name
),
false => format!(
"{}/target/debug/{}",
&function_dir, self.agent_config.workload_name
&function_dir, self.workload_config.workload_name
),
};

std::fs::copy(
binary_path,
format!("/tmp/{}", self.agent_config.workload_name),
format!("/tmp/{}", self.workload_config.workload_name),
)
.expect("Unable to copy binary");

Expand All @@ -111,7 +113,7 @@ impl Agent for RustAgent {
}

fn run(&self) -> AgentResult<()> {
let output = Command::new(format!("/tmp/{}", self.agent_config.workload_name))
let output = Command::new(format!("/tmp/{}", self.workload_config.workload_name))
.output()
.expect("Failed to run function");

Expand Down
44 changes: 10 additions & 34 deletions src/agent/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,37 +1,13 @@
mod agent;
mod rust;

use std::path::PathBuf;

use agent::{Action, Agent, AgentConfig, AgentResult};

pub struct AgentRunner {
config: AgentConfig,
agent: Box<dyn Agent>,
mod agents;
pub mod workload {
pub mod config;
pub mod runner;
}

impl AgentRunner {
pub fn new(config_path: String) -> Self {
let config = AgentConfig::new_from_file(&PathBuf::from(config_path)).unwrap();

let agent: Box<dyn Agent> = match config.language.as_str() {
"rust" => Box::new(rust::RustAgent::new_from_config(config.clone())),
_ => panic!("Unsupported language: {}", config.language),
};

AgentRunner { config, agent }
}

pub fn run(&self) -> AgentResult<()> {
match self.config.action {
Action::Prepare => self.agent.prepare()?,
Action::Run => self.agent.run()?,
Action::PrepareAndRun => {
self.agent.prepare()?;
self.agent.run()?;
}
}

Ok(())
}
#[derive(Debug)]
pub enum AgentError {
OpenConfigFileError(std::io::Error),
ParseConfigError(toml::de::Error),
}

pub type AgentResult<T> = Result<T, AgentError>;
10 changes: 6 additions & 4 deletions src/agent/src/main.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,18 @@
use agent::AgentRunner;
use agent::workload::{config::Config, runner::Runner};
use clap::Parser;
use std::path::PathBuf;

#[derive(Debug, Parser)]
struct Args {
#[clap(short, long, default_value = "/etc/cloudlet/agent/config.toml")]
config: String,
config: PathBuf,
}

fn main() {
let args = Args::parse();

let agent_runner = AgentRunner::new(args.config);
let config = Config::from_file(&args.config).unwrap();
let runner = Runner::new(config);

agent_runner.run().unwrap();
runner.run().unwrap();
}
46 changes: 15 additions & 31 deletions src/agent/src/agent.rs → src/agent/src/workload/config.rs
Original file line number Diff line number Diff line change
@@ -1,50 +1,26 @@
use std::path::PathBuf;

use crate::{agents::Language, AgentError, AgentResult};
use serde::Deserialize;

#[derive(Debug)]
pub enum AgentError {
OpenConfigFileError(std::io::Error),
ParseConfigError(toml::de::Error),
}

pub type AgentResult<T> = std::result::Result<T, AgentError>;

pub trait Agent {
fn prepare(&self) -> AgentResult<()>;
fn run(&self) -> AgentResult<()>;
}

#[derive(Debug, Clone, Deserialize)]
pub enum Action {
#[serde(rename = "prepare")]
Prepare,
#[serde(rename = "run")]
Run,
#[serde(rename = "prepare-and-run")]
PrepareAndRun,
}
use std::path::PathBuf;

/// Generic agent configuration.
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub struct AgentConfig {
pub struct Config {
/// Name of the worklod, used to identify the workload.
pub workload_name: String,
/// Language of the workload.
pub language: String,
pub language: Language,
/// Action to perform.
pub action: Action,
/// Rest of the configuration as a string.
#[serde(skip)]
pub config_string: String,
}

impl AgentConfig {
pub fn new_from_file(file_path: &PathBuf) -> AgentResult<Self> {
impl Config {
pub fn from_file(file_path: &PathBuf) -> AgentResult<Self> {
let config = std::fs::read_to_string(file_path).map_err(AgentError::OpenConfigFileError)?;
let mut config: AgentConfig =
toml::from_str(&config).map_err(AgentError::ParseConfigError)?;
let mut config: Config = toml::from_str(&config).map_err(AgentError::ParseConfigError)?;

let config_string =
std::fs::read_to_string(file_path).map_err(AgentError::OpenConfigFileError)?;
Expand All @@ -54,3 +30,11 @@ impl AgentConfig {
Ok(config)
}
}

#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum Action {
Prepare,
Run,
PrepareAndRun,
}
35 changes: 35 additions & 0 deletions src/agent/src/workload/runner.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
use super::config::{Action, Config};
use crate::{
agents::{rust, Agent, Language},
AgentResult,
};

/// Runner for a workload.
/// Will execute the workload based on the inner agent (language).
pub struct Runner {
config: Config,
agent: Box<dyn Agent>,
}

impl Runner {
pub fn new(config: Config) -> Self {
let agent: Box<dyn Agent> = match config.language {
Language::Rust => Box::new(rust::RustAgent::from(config.clone())),
};

Runner { config, agent }
}

pub fn run(&self) -> AgentResult<()> {
match self.config.action {
Action::Prepare => self.agent.prepare()?,
Action::Run => self.agent.run()?,
Action::PrepareAndRun => {
self.agent.prepare()?;
self.agent.run()?;
}
}

Ok(())
}
}

0 comments on commit ba094d6

Please sign in to comment.