use std::collections::BTreeMap;
pub type JobId = u32;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobDisposition {
Foreground,
Background,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProcessState {
Running,
Stopped,
Completed(i32),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobState {
Running,
Stopped,
Completed,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProcessRecord {
pid: u32,
state: ProcessState,
summary: String,
}
impl ProcessRecord {
pub fn new(pid: u32, summary: impl Into<String>) -> Self {
Self {
pid,
state: ProcessState::Running,
summary: summary.into(),
}
}
pub fn pid(&self) -> u32 {
self.pid
}
pub fn state(&self) -> ProcessState {
self.state
}
pub fn summary(&self) -> &str {
&self.summary
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct JobRecord {
id: JobId,
pgid: u32,
summary: String,
disposition: JobDisposition,
state: JobState,
processes: Vec<ProcessRecord>,
}
impl JobRecord {
pub fn id(&self) -> JobId {
self.id
}
pub fn pgid(&self) -> u32 {
self.pgid
}
pub fn summary(&self) -> &str {
&self.summary
}
pub fn disposition(&self) -> JobDisposition {
self.disposition
}
pub fn state(&self) -> JobState {
self.state
}
pub fn processes(&self) -> &[ProcessRecord] {
&self.processes
}
}
#[derive(Debug, Clone, Default)]
pub struct Jobs {
next_job_id: JobId,
foreground_job: Option<JobId>,
jobs: BTreeMap<JobId, JobRecord>,
}
impl Jobs {
pub fn insert(
&mut self,
pgid: u32,
summary: impl Into<String>,
disposition: JobDisposition,
processes: Vec<ProcessRecord>,
) -> JobId {
let id = self.alloc_job_id();
let record = JobRecord {
id,
pgid,
summary: summary.into(),
disposition,
state: derive_job_state(&processes),
processes,
};
if matches!(disposition, JobDisposition::Foreground) {
self.foreground_job = Some(id);
}
self.jobs.insert(id, record);
id
}
pub fn add_process(&mut self, job_id: JobId, process: ProcessRecord) -> bool {
let Some(job) = self.jobs.get_mut(&job_id) else {
return false;
};
job.processes.push(process);
job.state = derive_job_state(&job.processes);
true
}
pub fn update_process_state(&mut self, job_id: JobId, pid: u32, state: ProcessState) -> bool {
let Some(job) = self.jobs.get_mut(&job_id) else {
return false;
};
let Some(process) = job.processes.iter_mut().find(|process| process.pid == pid) else {
return false;
};
process.state = state;
job.state = derive_job_state(&job.processes);
if !matches!(job.state, JobState::Running) && self.foreground_job == Some(job_id) {
self.foreground_job = None;
}
true
}
pub fn set_disposition(&mut self, job_id: JobId, disposition: JobDisposition) -> bool {
let Some(job) = self.jobs.get_mut(&job_id) else {
return false;
};
job.disposition = disposition;
match disposition {
JobDisposition::Foreground => self.foreground_job = Some(job_id),
JobDisposition::Background if self.foreground_job == Some(job_id) => {
self.foreground_job = None;
}
JobDisposition::Background => {}
}
true
}
pub fn remove(&mut self, job_id: JobId) -> Option<JobRecord> {
if self.foreground_job == Some(job_id) {
self.foreground_job = None;
}
self.jobs.remove(&job_id)
}
pub fn get(&self, job_id: JobId) -> Option<&JobRecord> {
self.jobs.get(&job_id)
}
pub fn job_id_for_pid(&self, pid: u32) -> Option<JobId> {
self.jobs.iter().find_map(|(job_id, job)| {
job.processes()
.iter()
.any(|process| process.pid() == pid)
.then_some(*job_id)
})
}
pub fn current_job(&self) -> Option<JobId> {
self.jobs
.iter()
.rev()
.find(|(_, job)| !matches!(job.state, JobState::Completed))
.map(|(job_id, _)| *job_id)
}
pub fn foreground_job(&self) -> Option<JobId> {
self.foreground_job
}
pub fn set_all_processes_running(&mut self, job_id: JobId) -> bool {
self.update_non_completed_processes(job_id, ProcessState::Running)
}
pub fn set_all_processes_stopped(&mut self, job_id: JobId) -> bool {
self.update_non_completed_processes(job_id, ProcessState::Stopped)
}
pub fn iter(&self) -> impl Iterator<Item = &JobRecord> {
self.jobs.values()
}
pub fn len(&self) -> usize {
self.jobs.len()
}
pub fn is_empty(&self) -> bool {
self.jobs.is_empty()
}
fn alloc_job_id(&mut self) -> JobId {
self.next_job_id = self.next_job_id.saturating_add(1).max(1);
self.next_job_id
}
fn update_non_completed_processes(&mut self, job_id: JobId, state: ProcessState) -> bool {
let Some(job) = self.jobs.get_mut(&job_id) else {
return false;
};
for process in &mut job.processes {
if !matches!(process.state, ProcessState::Completed(_)) {
process.state = state;
}
}
job.state = derive_job_state(&job.processes);
if !matches!(job.state, JobState::Running) && self.foreground_job == Some(job_id) {
self.foreground_job = None;
}
true
}
}
fn derive_job_state(processes: &[ProcessRecord]) -> JobState {
if processes.is_empty() {
return JobState::Completed;
}
if processes
.iter()
.all(|process| matches!(process.state, ProcessState::Completed(_)))
{
JobState::Completed
} else if processes
.iter()
.any(|process| matches!(process.state, ProcessState::Stopped))
{
JobState::Stopped
} else {
JobState::Running
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn inserts_jobs_and_tracks_foreground() {
let mut jobs = Jobs::default();
let id = jobs.insert(
1000,
"sleep 1",
JobDisposition::Foreground,
vec![ProcessRecord::new(1000, "sleep 1")],
);
assert_eq!(id, 1);
assert_eq!(jobs.foreground_job(), Some(id));
let job = jobs.get(id).expect("job should exist");
assert_eq!(job.pgid(), 1000);
assert_eq!(job.summary(), "sleep 1");
assert_eq!(job.state(), JobState::Running);
}
#[test]
fn process_state_transitions_update_job_state() {
let mut jobs = Jobs::default();
let id = jobs.insert(
2000,
"pipeline",
JobDisposition::Background,
vec![
ProcessRecord::new(2001, "printf a"),
ProcessRecord::new(2002, "cat"),
],
);
assert!(jobs.update_process_state(id, 2001, ProcessState::Stopped));
assert_eq!(
jobs.get(id).expect("job should exist").state(),
JobState::Stopped
);
assert!(jobs.update_process_state(id, 2001, ProcessState::Completed(0)));
assert_eq!(
jobs.get(id).expect("job should exist").state(),
JobState::Running
);
assert!(jobs.update_process_state(id, 2002, ProcessState::Completed(0)));
assert_eq!(
jobs.get(id).expect("job should exist").state(),
JobState::Completed
);
}
#[test]
fn stopped_foreground_job_releases_foreground_slot() {
let mut jobs = Jobs::default();
let id = jobs.insert(
3000,
"sleep 10",
JobDisposition::Foreground,
vec![ProcessRecord::new(3000, "sleep 10")],
);
assert!(jobs.update_process_state(id, 3000, ProcessState::Stopped));
assert_eq!(jobs.foreground_job(), None);
assert_eq!(
jobs.get(id).expect("job should exist").state(),
JobState::Stopped
);
}
}