stuart/databuild
1
1
Fork 0
Code Issues 1 Pull requests Projects Releases Packages Wiki Activity Actions

refactor job run state to separate state types

Browse source
This commit is contained in:
Stuart Axelbrooke 2025-10-16 18:19:24 -07:00
parent cfcb201285
commit 6572d4e3bd
1 changed files with 144 additions and 103 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

131
databuild/job_run.rs
View file

@ -1,5 +1,4 @@
use std::collections::HashMap;
use crate::build_event_log::{BELStorage, MemoryBELStorage};
use crate::data_build_event::Event;
use crate::{JobRunHeartbeatEventV1, JobRunStatus, JobRunStatusCode, JobRunSuccessEventV1, JobRunFailureEventV1, JobRunCancelEventV1, EventSource};
use std::error::Error;
@ -36,14 +35,30 @@ pub fn spawn_job_run(config: JobRunConfig) -> Result<Box<dyn JobRun>, Box<dyn Er
pub struct SubProcessJobRun {
pub job_run_id: Uuid,
pub process: Option<Child>,
pub storage: MemoryBELStorage,
pub running: bool,
pub entry_point: String,
pub args: Vec<String>,
pub stdout_buffer: Vec<String>, // Buffered stdout lines
pub exit_code: Option<i32>, // Cached exit code
pub cancel_source: Option<EventSource>, // Source of cancellation if canceled
pub state: JobRunState,
}
enum JobRunState {
NotStarted,
Running {
process: Child,
stdout_buffer: Vec<String>,
},
Completed {
exit_code: i32,
stdout_buffer: Vec<String>,
},
Failed {
exit_code: i32,
reason: String,
stdout_buffer: Vec<String>,
},
Canceled {
source: EventSource,
stdout_buffer: Vec<String>,
},
}
impl JobRun for SubProcessJobRun {
@ -54,36 +69,37 @@ impl JobRun for SubProcessJobRun {
fn visit(&mut self) -> Result<JobRunPollResult, Box<dyn Error>> {
let mut new_events = Vec::new();
// visit() should only be called on running jobs
if !self.running {
return Err("visit() called on non-running job".into());
}
// Process must be present if running
let child = self.process.as_mut()
.ok_or("visit() called but no process present")?;
match &mut self.state {
JobRunState::Running { process, stdout_buffer } => {
// Non-blocking check for exit status
if let Some(exit_status) = child.try_wait()? {
self.running = false;
if let Some(exit_status) = process.try_wait()? {
// Read any remaining stdout
if let Some(stdout) = child.stdout.take() {
if let Some(stdout) = process.stdout.take() {
let reader = BufReader::new(stdout);
for line in reader.lines() {
// TODO we should write lines to the job's file logs
if let Ok(line) = line {
self.stdout_buffer.push(line);
stdout_buffer.push(line);
}
}
}
// Take ownership of the current state to transition
let old_state = std::mem::replace(&mut self.state, JobRunState::NotStarted);
let stdout_buf = if let JobRunState::Running { stdout_buffer, .. } = old_state {
stdout_buffer
} else {
Vec::new()
};
// Check exit status
// Check exit status and transition to terminal state
match exit_status.code() {
Some(0) => {
// Success case
self.exit_code = Some(0);
self.state = JobRunState::Completed {
exit_code: 0,
stdout_buffer: stdout_buf,
};
new_events.push(Event::JobRunSuccessV1(JobRunSuccessEventV1 {
job_run_id: self.job_run_id.to_string(),
}));
@ -94,10 +110,15 @@ impl JobRun for SubProcessJobRun {
}
Some(code) => {
// Failed with exit code
self.exit_code = Some(code);
let reason = format!("Job failed with exit code {}", code);
self.state = JobRunState::Failed {
exit_code: code,
reason: reason.clone(),
stdout_buffer: stdout_buf,
};
new_events.push(Event::JobRunFailureV1(JobRunFailureEventV1 {
job_run_id: self.job_run_id.to_string(),
reason: format!("Job failed with exit code {}", code),
reason,
}));
return Ok(JobRunPollResult {
new_events,
@ -106,9 +127,15 @@ impl JobRun for SubProcessJobRun {
}
None => {
// Terminated by signal (Unix) - treat as failure
let reason = format!("Job terminated by signal: {}", exit_status);
self.state = JobRunState::Failed {
exit_code: -1,
reason: reason.clone(),
stdout_buffer: stdout_buf,
};
new_events.push(Event::JobRunFailureV1(JobRunFailureEventV1 {
job_run_id: self.job_run_id.to_string(),
reason: format!("Job terminated by signal: {}", exit_status),
reason,
}));
return Ok(JobRunPollResult {
new_events,
@ -124,22 +151,24 @@ impl JobRun for SubProcessJobRun {
status: JobRunStatusCode::JobRunRunning.into(),
})
}
fn cancel(&mut self, source: EventSource) -> Result<JobRunCancelEventV1, Box<dyn Error>> {
if !self.running {
return Err("cancel() called on non-running job".into());
_ => Err("visit() called on non-running job".into()),
}
}
let child = self.process.as_mut()
.ok_or("cancel() called but no process present")?;
fn cancel(&mut self, source: EventSource) -> Result<JobRunCancelEventV1, Box<dyn Error>> {
match std::mem::replace(&mut self.state, JobRunState::NotStarted) {
JobRunState::Running { mut process, stdout_buffer } => {
// Kill the process
child.kill()?;
process.kill()?;
// Wait for it to actually terminate
child.wait()?;
process.wait()?;
self.running = false;
// Transition to Canceled state
self.state = JobRunState::Canceled {
source: source.clone(),
stdout_buffer,
};
Ok(JobRunCancelEventV1 {
job_run_id: self.job_run_id.to_string(),
@ -147,33 +176,45 @@ impl JobRun for SubProcessJobRun {
comment: Some("Job was canceled".to_string()),
})
}
other_state => {
// Restore the state and error
self.state = other_state;
Err("cancel() called on non-running job".into())
}
}
}
/// Mostly for test purposes
fn run_with_env(&mut self, env: Option<HashMap<String, String>>) -> Result<JobRunHeartbeatEventV1, Box<dyn Error>> {
self.process = Some(Command::new(self.entry_point.clone())
match &self.state {
JobRunState::NotStarted => {
let process = Command::new(self.entry_point.clone())
.args(self.args.clone())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.envs(env.unwrap_or_default())
.spawn()?);
self.running = true;
.spawn()?;
self.state = JobRunState::Running {
process,
stdout_buffer: Vec::new(),
};
// TODO should this return the event now? Or enqueue it? No sense in waiting I suppose, and orchestrator should just handle it?
Ok(JobRunHeartbeatEventV1 { job_run_id: self.id().to_string() })
}
_ => Err("run() called on already-running or completed job".into())
}
}
}
impl SubProcessJobRun {
pub fn spawn(entry_point: String, args: Vec<String>) -> Result<Box<dyn JobRun>, Box<dyn Error>> {
Ok(Box::new(SubProcessJobRun {
job_run_id: Uuid::new_v4(),
process: None,
storage: MemoryBELStorage::default(),
running: false,
entry_point,
args,
stdout_buffer: Vec::new(),
exit_code: None,
cancel_source: None,
state: JobRunState::NotStarted,
}))
}
}