|
|
|
|
@ -1,97 +1,136 @@ |
|
|
|
|
use crate::{ |
|
|
|
|
cache::JobCache, |
|
|
|
|
executor::{DynFut, Waiter}, |
|
|
|
|
messaging::Reportable, |
|
|
|
|
models::{Agent, AssignedJob, JobMeta, JobType, RawJobMeta}, |
|
|
|
|
executor::{ExecResult, Waiter}, |
|
|
|
|
models::{Agent, AssignedJob, AssignedJobById, JobMeta, JobType}, |
|
|
|
|
utils::{CombinedResult, OneOrVec, Platform}, |
|
|
|
|
utils::{ProcOutput, TempFile}, |
|
|
|
|
UError, UResult, |
|
|
|
|
}; |
|
|
|
|
use std::collections::HashMap; |
|
|
|
|
use std::process::exit; |
|
|
|
|
use tokio::process::Command; |
|
|
|
|
|
|
|
|
|
pub struct JobBuilder { |
|
|
|
|
pub struct JobRunner { |
|
|
|
|
waiter: Waiter, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl JobBuilder { |
|
|
|
|
pub fn from_request(job_requests: impl OneOrVec<AssignedJob>) -> CombinedResult<Self> { |
|
|
|
|
let job_requests = job_requests.into_vec(); |
|
|
|
|
let mut prepared: Vec<DynFut> = vec![]; |
|
|
|
|
let mut result = CombinedResult::<JobBuilder, _>::new(); |
|
|
|
|
for req in job_requests { |
|
|
|
|
let job_meta = JobCache::get(req.job_id); |
|
|
|
|
if job_meta.is_none() { |
|
|
|
|
result.err(UError::NoJob(req.job_id)); |
|
|
|
|
continue; |
|
|
|
|
} |
|
|
|
|
let job_meta = job_meta.unwrap(); |
|
|
|
|
impl JobRunner { |
|
|
|
|
pub fn from_jobs(jobs: impl OneOrVec<AssignedJobById>) -> CombinedResult<Self> { |
|
|
|
|
let jobs = jobs.into_vec(); |
|
|
|
|
let mut waiter = Waiter::new(); |
|
|
|
|
let mut result = CombinedResult::<Self, _>::new(); |
|
|
|
|
for job in jobs { |
|
|
|
|
//waiting for try-blocks stabilization
|
|
|
|
|
let built_req = (|| -> UResult<()> { |
|
|
|
|
Ok(match job_meta.exec_type { |
|
|
|
|
JobType::Shell => { |
|
|
|
|
let meta = JobCache::get(req.job_id).ok_or(UError::NoJob(req.job_id))?; |
|
|
|
|
let curr_platform = Platform::current(); |
|
|
|
|
if !curr_platform.matches(&meta.platform) { |
|
|
|
|
return Err(UError::InsuitablePlatform( |
|
|
|
|
meta.platform.clone(), |
|
|
|
|
curr_platform.into_string(), |
|
|
|
|
) |
|
|
|
|
.into()); |
|
|
|
|
} |
|
|
|
|
let job = AssignedJob::new(req.job_id, Some(&req)); |
|
|
|
|
prepared.push(Box::pin(job.run())) |
|
|
|
|
} |
|
|
|
|
JobType::Manage => prepared.push(Box::pin(Agent::run())), |
|
|
|
|
_ => todo!(), |
|
|
|
|
}) |
|
|
|
|
let built_job = (|| -> UResult<()> { |
|
|
|
|
let meta = JobCache::get(job.job_id).ok_or(UError::NoJob(job.job_id))?; |
|
|
|
|
let curr_platform = Platform::current(); |
|
|
|
|
if !curr_platform.matches(&meta.platform) { |
|
|
|
|
return Err(UError::InsuitablePlatform( |
|
|
|
|
meta.platform.clone(), |
|
|
|
|
curr_platform.into_string(), |
|
|
|
|
) |
|
|
|
|
.into()); |
|
|
|
|
} |
|
|
|
|
let job = AssignedJob::from((&*meta, job)); |
|
|
|
|
waiter.push(run_assigned_job(job)); |
|
|
|
|
Ok(()) |
|
|
|
|
})(); |
|
|
|
|
if let Err(e) = built_req { |
|
|
|
|
if let Err(e) = built_job { |
|
|
|
|
result.err(e) |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
result.ok(Self { |
|
|
|
|
waiter: Waiter::new(prepared), |
|
|
|
|
}); |
|
|
|
|
result.ok(Self { waiter }); |
|
|
|
|
result |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pub fn from_meta(job_metas: impl OneOrVec<JobMeta>) -> CombinedResult<Self> { |
|
|
|
|
let job_requests = job_metas |
|
|
|
|
pub fn from_meta(metas: impl OneOrVec<JobMeta>) -> CombinedResult<Self> { |
|
|
|
|
let jobs = metas |
|
|
|
|
.into_vec() |
|
|
|
|
.into_iter() |
|
|
|
|
.map(|jm| { |
|
|
|
|
let j_uid = jm.id; |
|
|
|
|
let job_uid = jm.id; |
|
|
|
|
JobCache::insert(jm); |
|
|
|
|
AssignedJob::new(j_uid, None) |
|
|
|
|
AssignedJobById { |
|
|
|
|
job_id: job_uid, |
|
|
|
|
..Default::default() |
|
|
|
|
} |
|
|
|
|
}) |
|
|
|
|
.collect::<Vec<AssignedJob>>(); |
|
|
|
|
JobBuilder::from_request(job_requests) |
|
|
|
|
.collect::<Vec<AssignedJobById>>(); |
|
|
|
|
JobRunner::from_jobs(jobs) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Spawn jobs and pop results later
|
|
|
|
|
/// Spawn jobs
|
|
|
|
|
pub async fn spawn(mut self) -> Self { |
|
|
|
|
self.waiter = self.waiter.spawn().await; |
|
|
|
|
self |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Spawn jobs and wait for result
|
|
|
|
|
pub async fn wait(self) -> Vec<Reportable> { |
|
|
|
|
pub async fn wait(self) -> Vec<ExecResult> { |
|
|
|
|
self.waiter.spawn().await.wait().await |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Spawn one job and wait for result
|
|
|
|
|
pub async fn wait_one(self) -> Reportable { |
|
|
|
|
pub async fn wait_one(self) -> ExecResult { |
|
|
|
|
self.waiter.spawn().await.wait().await.pop().unwrap() |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pub async fn run_assigned_job(mut job: AssignedJob) -> ExecResult { |
|
|
|
|
match job.exec_type { |
|
|
|
|
JobType::Shell => { |
|
|
|
|
let (argv, _payload) = { |
|
|
|
|
let meta = JobCache::get(job.job_id).unwrap(); |
|
|
|
|
if let Some(ref payload) = meta.payload { |
|
|
|
|
let extracted_payload = match TempFile::write_exec(payload) { |
|
|
|
|
Ok(p) => p, |
|
|
|
|
Err(e) => return Err(UError::Runtime(e.to_string())), |
|
|
|
|
}; |
|
|
|
|
( |
|
|
|
|
meta.argv.replace("{}", &extracted_payload.get_path()), |
|
|
|
|
Some(extracted_payload), |
|
|
|
|
) |
|
|
|
|
} else { |
|
|
|
|
(meta.argv.clone(), None) |
|
|
|
|
} |
|
|
|
|
}; |
|
|
|
|
let mut split_cmd = shlex::split(&argv).unwrap().into_iter(); |
|
|
|
|
let cmd = split_cmd.nth(0).unwrap(); |
|
|
|
|
let args = split_cmd.collect::<Vec<String>>(); |
|
|
|
|
let cmd_result = Command::new(cmd).args(args).output().await; |
|
|
|
|
let (data, retcode) = match cmd_result { |
|
|
|
|
Ok(output) => ( |
|
|
|
|
ProcOutput::from_output(&output).into_combined(), |
|
|
|
|
output.status.code(), |
|
|
|
|
), |
|
|
|
|
Err(e) => ( |
|
|
|
|
ProcOutput::new() |
|
|
|
|
.stderr(e.to_string().into_bytes()) |
|
|
|
|
.into_combined(), |
|
|
|
|
None, |
|
|
|
|
), |
|
|
|
|
}; |
|
|
|
|
job.result = Some(data); |
|
|
|
|
job.retcode = retcode; |
|
|
|
|
} |
|
|
|
|
JobType::Init => { |
|
|
|
|
job.set_result(&Agent::run().await); |
|
|
|
|
job.retcode = Some(0); |
|
|
|
|
} |
|
|
|
|
JobType::Update => todo!(), |
|
|
|
|
JobType::Terminate => exit(0), |
|
|
|
|
}; |
|
|
|
|
Ok(job) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
/// Store jobs and get results by name
|
|
|
|
|
pub struct NamedJobBuilder { |
|
|
|
|
builder: Option<JobBuilder>, |
|
|
|
|
pub struct NamedJobRunner { |
|
|
|
|
builder: Option<JobRunner>, |
|
|
|
|
job_names: Vec<&'static str>, |
|
|
|
|
results: HashMap<&'static str, Reportable>, |
|
|
|
|
results: HashMap<&'static str, ExecResult>, |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
impl NamedJobBuilder { |
|
|
|
|
impl NamedJobRunner { |
|
|
|
|
pub fn from_shell( |
|
|
|
|
named_jobs: impl OneOrVec<(&'static str, &'static str)>, |
|
|
|
|
) -> CombinedResult<Self> { |
|
|
|
|
@ -100,7 +139,7 @@ impl NamedJobBuilder { |
|
|
|
|
.into_vec() |
|
|
|
|
.into_iter() |
|
|
|
|
.filter_map( |
|
|
|
|
|(alias, cmd)| match RawJobMeta::builder().with_shell(cmd).build() { |
|
|
|
|
|(alias, cmd)| match JobMeta::builder().with_shell(cmd).build() { |
|
|
|
|
Ok(meta) => Some((alias, meta)), |
|
|
|
|
Err(e) => { |
|
|
|
|
result.err(e); |
|
|
|
|
@ -124,7 +163,7 @@ impl NamedJobBuilder { |
|
|
|
|
}) |
|
|
|
|
.collect(); |
|
|
|
|
Self { |
|
|
|
|
builder: Some(JobBuilder::from_meta(job_metas).unwrap_one()), |
|
|
|
|
builder: Some(JobRunner::from_meta(job_metas).unwrap_one()), |
|
|
|
|
job_names, |
|
|
|
|
results: HashMap::new(), |
|
|
|
|
} |
|
|
|
|
@ -138,11 +177,11 @@ impl NamedJobBuilder { |
|
|
|
|
self |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pub fn pop_opt(&mut self, name: &'static str) -> Option<Reportable> { |
|
|
|
|
pub fn pop_opt(&mut self, name: &'static str) -> Option<ExecResult> { |
|
|
|
|
self.results.remove(name) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
pub fn pop(&mut self, name: &'static str) -> Reportable { |
|
|
|
|
pub fn pop(&mut self, name: &'static str) -> ExecResult { |
|
|
|
|
self.pop_opt(name).unwrap() |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
@ -151,8 +190,8 @@ impl NamedJobBuilder { |
|
|
|
|
mod tests { |
|
|
|
|
use super::*; |
|
|
|
|
use crate::{ |
|
|
|
|
builder::{JobBuilder, NamedJobBuilder}, |
|
|
|
|
models::{misc::JobType, JobMeta}, |
|
|
|
|
runner::{JobRunner, NamedJobRunner}, |
|
|
|
|
unwrap_enum, |
|
|
|
|
}; |
|
|
|
|
use std::time::SystemTime; |
|
|
|
|
@ -162,10 +201,10 @@ mod tests { |
|
|
|
|
#[tokio::test] |
|
|
|
|
async fn test_is_really_async() { |
|
|
|
|
const SLEEP_SECS: u64 = 1; |
|
|
|
|
let job = RawJobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
|
|
|
let job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
|
|
|
let sleep_jobs: Vec<JobMeta> = (0..50).map(|_| job.clone()).collect(); |
|
|
|
|
let now = SystemTime::now(); |
|
|
|
|
JobBuilder::from_meta(sleep_jobs).unwrap_one().wait().await; |
|
|
|
|
JobRunner::from_meta(sleep_jobs).unwrap_one().wait().await; |
|
|
|
|
assert!(now.elapsed().unwrap().as_secs() < SLEEP_SECS + 2) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
@ -199,13 +238,16 @@ mod tests { |
|
|
|
|
#[case] payload: Option<&[u8]>, |
|
|
|
|
#[case] expected_result: &str, |
|
|
|
|
) -> TestResult { |
|
|
|
|
let mut job = RawJobMeta::builder().with_shell(cmd); |
|
|
|
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
|
|
|
if let Some(p) = payload { |
|
|
|
|
job = job.with_payload(p); |
|
|
|
|
} |
|
|
|
|
let job = job.build().unwrap(); |
|
|
|
|
let job_result = JobBuilder::from_meta(job).unwrap_one().wait_one().await; |
|
|
|
|
let result = unwrap_enum!(job_result, Reportable::Assigned); |
|
|
|
|
let result = JobRunner::from_meta(job) |
|
|
|
|
.unwrap_one() |
|
|
|
|
.wait_one() |
|
|
|
|
.await |
|
|
|
|
.unwrap(); |
|
|
|
|
let result = result.to_string_result(); |
|
|
|
|
assert_eq!(result.trim(), expected_result); |
|
|
|
|
Ok(()) |
|
|
|
|
@ -215,29 +257,25 @@ mod tests { |
|
|
|
|
async fn test_complex_load() -> TestResult { |
|
|
|
|
const SLEEP_SECS: u64 = 1; |
|
|
|
|
let now = SystemTime::now(); |
|
|
|
|
let longest_job = RawJobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
|
|
|
let longest_job = JobBuilder::from_meta(longest_job) |
|
|
|
|
.unwrap_one() |
|
|
|
|
.spawn() |
|
|
|
|
.await; |
|
|
|
|
let ls = JobBuilder::from_meta(RawJobMeta::from_shell("ls")?) |
|
|
|
|
let longest_job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap(); |
|
|
|
|
let longest_job = JobRunner::from_meta(longest_job).unwrap_one().spawn().await; |
|
|
|
|
let ls = JobRunner::from_meta(JobMeta::from_shell("ls")?) |
|
|
|
|
.unwrap_one() |
|
|
|
|
.wait_one() |
|
|
|
|
.await; |
|
|
|
|
let ls = unwrap_enum!(ls, Reportable::Assigned); |
|
|
|
|
.await |
|
|
|
|
.unwrap(); |
|
|
|
|
assert_eq!(ls.retcode.unwrap(), 0); |
|
|
|
|
let folders = ls.to_string_result(); |
|
|
|
|
let subfolders_jobs: Vec<JobMeta> = folders |
|
|
|
|
.lines() |
|
|
|
|
.map(|f| RawJobMeta::from_shell(format!("ls {}", f)).unwrap()) |
|
|
|
|
.map(|f| JobMeta::from_shell(format!("ls {}", f)).unwrap()) |
|
|
|
|
.collect(); |
|
|
|
|
let ls_subfolders = JobBuilder::from_meta(subfolders_jobs) |
|
|
|
|
let ls_subfolders = JobRunner::from_meta(subfolders_jobs) |
|
|
|
|
.unwrap_one() |
|
|
|
|
.wait() |
|
|
|
|
.await; |
|
|
|
|
for result in ls_subfolders { |
|
|
|
|
let result = unwrap_enum!(result, Reportable::Assigned); |
|
|
|
|
assert_eq!(result.retcode.unwrap(), 0); |
|
|
|
|
assert_eq!(result.unwrap().retcode.unwrap(), 0); |
|
|
|
|
} |
|
|
|
|
longest_job.wait().await; |
|
|
|
|
assert_eq!(now.elapsed().unwrap().as_secs(), SLEEP_SECS); |
|
|
|
|
@ -263,9 +301,12 @@ mod tests { |
|
|
|
|
*/ |
|
|
|
|
#[tokio::test] |
|
|
|
|
async fn test_failing_shell_job() -> TestResult { |
|
|
|
|
let job = RawJobMeta::from_shell("lol_kek_puk")?; |
|
|
|
|
let job_result = JobBuilder::from_meta(job).unwrap_one().wait_one().await; |
|
|
|
|
let job_result = unwrap_enum!(job_result, Reportable::Assigned); |
|
|
|
|
let job = JobMeta::from_shell("lol_kek_puk")?; |
|
|
|
|
let job_result = JobRunner::from_meta(job) |
|
|
|
|
.unwrap_one() |
|
|
|
|
.wait_one() |
|
|
|
|
.await |
|
|
|
|
.unwrap(); |
|
|
|
|
let output = job_result.to_string_result(); |
|
|
|
|
assert!(output.contains("No such file")); |
|
|
|
|
assert!(job_result.retcode.is_none()); |
|
|
|
|
@ -281,7 +322,7 @@ mod tests { |
|
|
|
|
#[case] payload: Option<&[u8]>, |
|
|
|
|
#[case] err_str: &str, |
|
|
|
|
) -> TestResult { |
|
|
|
|
let mut job = RawJobMeta::builder().with_shell(cmd); |
|
|
|
|
let mut job = JobMeta::builder().with_shell(cmd); |
|
|
|
|
if let Some(p) = payload { |
|
|
|
|
job = job.with_payload(p); |
|
|
|
|
} |
|
|
|
|
@ -293,17 +334,17 @@ mod tests { |
|
|
|
|
|
|
|
|
|
#[tokio::test] |
|
|
|
|
async fn test_different_job_types() -> TestResult { |
|
|
|
|
let mut jobs = NamedJobBuilder::from_meta(vec![ |
|
|
|
|
("sleeper", RawJobMeta::from_shell("sleep 3")?), |
|
|
|
|
let mut jobs = NamedJobRunner::from_meta(vec![ |
|
|
|
|
("sleeper", JobMeta::from_shell("sleep 3")?), |
|
|
|
|
( |
|
|
|
|
"gatherer", |
|
|
|
|
RawJobMeta::builder().with_type(JobType::Manage).build()?, |
|
|
|
|
JobMeta::builder().with_type(JobType::Init).build()?, |
|
|
|
|
), |
|
|
|
|
]) |
|
|
|
|
.wait() |
|
|
|
|
.await; |
|
|
|
|
let gathered = jobs.pop("gatherer"); |
|
|
|
|
assert_eq!(unwrap_enum!(gathered, Reportable::Agent).alias, None); |
|
|
|
|
assert_eq!(gathered.unwrap().alias, None); |
|
|
|
|
Ok(()) |
|
|
|
|
} |
|
|
|
|
} |
