icebox/core/
job.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::{Duration, Instant};
5
6use crate::core::module::ModuleResult;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
9pub struct JobId(pub u64);
10
11impl JobId {
12    pub fn as_u64(&self) -> u64 {
13        self.0
14    }
15}
16
17impl std::fmt::Display for JobId {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        write!(f, "{}", self.0)
20    }
21}
22
23static NEXT_JOB_ID: AtomicU64 = AtomicU64::new(1);
24
25fn next_job_id() -> JobId {
26    JobId(NEXT_JOB_ID.fetch_add(1, Ordering::Relaxed))
27}
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
30pub enum JobStatus {
31    Running,
32    Completed,
33    Failed,
34    Cancelled,
35}
36
37impl JobStatus {
38    pub fn as_str(&self) -> &'static str {
39        match self {
40            JobStatus::Running => "running",
41            JobStatus::Completed => "completed",
42            JobStatus::Failed => "failed",
43            JobStatus::Cancelled => "cancelled",
44        }
45    }
46}
47
48#[derive(Debug, Clone)]
49pub struct Job {
50    pub id: JobId,
51    pub module_name: String,
52    pub target: String,
53    pub status: JobStatus,
54    pub started_at: Instant,
55    pub finished_at: Option<Instant>,
56    pub result: Option<ModuleResult>,
57}
58
59impl Job {
60    pub fn new(module_name: impl Into<String>, target: impl Into<String>) -> Self {
61        Job {
62            id: next_job_id(),
63            module_name: module_name.into(),
64            target: target.into(),
65            status: JobStatus::Running,
66            started_at: Instant::now(),
67            finished_at: None,
68            result: None,
69        }
70    }
71
72    pub fn elapsed(&self) -> Duration {
73        let end = self.finished_at.unwrap_or(Instant::now());
74        end.duration_since(self.started_at)
75    }
76}
77
78#[derive(Debug, Default)]
79pub struct JobManager {
80    jobs: HashMap<JobId, Job>,
81}
82
83impl JobManager {
84    pub fn new() -> Self {
85        JobManager {
86            jobs: HashMap::new(),
87        }
88    }
89
90    pub fn register(&mut self, job: Job) -> JobId {
91        let id = job.id;
92        self.jobs.insert(id, job);
93        id
94    }
95
96    pub fn get(&self, id: JobId) -> Option<&Job> {
97        self.jobs.get(&id)
98    }
99
100    pub fn get_mut(&mut self, id: JobId) -> Option<&mut Job> {
101        self.jobs.get_mut(&id)
102    }
103
104    pub fn list(&self) -> Vec<&Job> {
105        let mut v: Vec<&Job> = self
106            .jobs
107            .values()
108            .filter(|j| matches!(j.status, JobStatus::Running))
109            .collect();
110        v.sort_by_key(|j| j.id);
111        v
112    }
113
114    pub fn list_recent(&self, n: usize) -> Vec<&Job> {
115        let mut v: Vec<&Job> = self.jobs.values().collect();
116        v.sort_by_key(|j| j.id);
117        v.into_iter().rev().take(n).collect()
118    }
119
120    pub fn complete(&mut self, id: JobId, result: ModuleResult) -> bool {
121        if let Some(job) = self.jobs.get_mut(&id) {
122            job.status = if result.success {
123                JobStatus::Completed
124            } else {
125                JobStatus::Failed
126            };
127            job.finished_at = Some(Instant::now());
128            job.result = Some(result);
129            true
130        } else {
131            false
132        }
133    }
134
135    pub fn advance_counter(min: u64) {
136        let mut prev = NEXT_JOB_ID.load(Ordering::Relaxed);
137        while prev <= min {
138            match NEXT_JOB_ID.compare_exchange(prev, min + 1, Ordering::Relaxed, Ordering::Relaxed)
139            {
140                Ok(_) => break,
141                Err(current) => prev = current,
142            }
143        }
144    }
145
146    pub fn cancel(&mut self, id: JobId) -> bool {
147        if let Some(job) = self.jobs.get_mut(&id) {
148            job.status = JobStatus::Cancelled;
149            job.finished_at = Some(Instant::now());
150            true
151        } else {
152            false
153        }
154    }
155}
icebox/core/job.rs

icebox/core/
job.rs
