kitchen_shell/
jobs.rs

1//! What runs in the background, in one place (from lumen's job registry): each job has a kind
2//! (`histology`, `model`, `export`…), a name saying what it is doing, how far it is, and a cancel
3//! flag. Only one job that drives the GPU runs at a time (a second is refused, saying what
4//! runs). The status bar lists them with a cancel button each; the work itself (off the UI
5//! thread) updates its job through its [`JobHandle`], which ends the job when dropped (finished,
6//! failed or cancelled alike).
7
8use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
9use std::sync::{Arc, Mutex, OnceLock};
10
11/// A job running.
12#[derive(Debug, Clone)]
13pub struct Job {
14    /// Its id (unique in this run of the app).
15    pub id: u64,
16    /// What kind of work it is (`model`, `export`).
17    pub kind: String,
18    /// What it is doing now ("Finding cells: tile 3 of 12").
19    pub name: String,
20    /// How far (0..1), when known.
21    pub progress: Option<f32>,
22    /// Whether it drives the GPU (one such job at a time).
23    pub gpu: bool,
24    /// Set to ask it to stop.
25    pub cancel: Arc<AtomicBool>,
26}
27
28#[derive(Default)]
29struct Inner {
30    list: Vec<Job>,
31    next: u64,
32}
33
34fn inner() -> &'static Mutex<Inner> {
35    static J: OnceLock<Mutex<Inner>> = OnceLock::new();
36    J.get_or_init(Default::default)
37}
38
39static GENERATION: AtomicU64 = AtomicU64::new(0);
40
41fn with<T>(f: impl FnOnce(&mut Inner) -> T) -> T {
42    let mut g = inner().lock().unwrap_or_else(|e| e.into_inner());
43    let out = f(&mut g);
44    GENERATION.fetch_add(1, Ordering::Relaxed);
45    out
46}
47
48/// Starts a job; refused while another GPU job runs (when this one is one too).
49pub fn start(kind: impl Into<String>, name: &str, gpu: bool) -> Result<JobHandle, String> {
50    let kind = kind.into();
51    with(|j| {
52        if gpu && let Some(other) = j.list.iter().find(|o| o.gpu) {
53            return Err(format!("{} is still running", other.name));
54        }
55        j.next += 1;
56        let cancel = Arc::new(AtomicBool::new(false));
57        j.list.push(Job { id: j.next, kind, name: name.to_string(), progress: None, gpu, cancel: cancel.clone() });
58        Ok(JobHandle { id: j.next, cancel })
59    })
60}
61
62/// The jobs running, oldest first.
63pub fn running() -> Vec<Job> {
64    let g = inner().lock().unwrap_or_else(|e| e.into_inner());
65    g.list.clone()
66}
67
68/// The latest job of `kind`.
69pub fn of(kind: &str) -> Option<Job> {
70    running().into_iter().rev().find(|o| o.kind == kind)
71}
72
73/// Asks every job of `kind` to stop (when its work next checks).
74pub fn cancel(kind: &str) {
75    with(|j| j.list.iter().filter(|o| o.kind == kind).for_each(|o| o.cancel.store(true, Ordering::Relaxed)));
76}
77
78/// Asks the job `id` to stop.
79pub fn cancel_id(id: u64) {
80    with(|j| j.list.iter().filter(|o| o.id == id).for_each(|o| o.cancel.store(true, Ordering::Relaxed)));
81}
82
83/// A number that changes whenever a job starts, moves or ends (the status bar redraws then).
84pub fn generation() -> u64 {
85    GENERATION.load(Ordering::Relaxed)
86}
87
88/// A running job, held by the work doing it.
89pub struct JobHandle {
90    /// The job's id.
91    pub id: u64,
92    /// Set when someone asked it to stop.
93    pub cancel: Arc<AtomicBool>,
94}
95
96impl JobHandle {
97    /// Says what it is doing now, and how far.
98    pub fn set(&self, name: &str, progress: Option<f32>) {
99        with(|j| {
100            if let Some(o) = j.list.iter_mut().find(|o| o.id == self.id) {
101                o.name = name.to_string();
102                o.progress = progress;
103            }
104        });
105    }
106
107    /// Whether someone asked it to stop.
108    pub fn cancelled(&self) -> bool {
109        self.cancel.load(Ordering::Relaxed)
110    }
111}
112
113impl Drop for JobHandle {
114    fn drop(&mut self) {
115        let id = self.id;
116        with(|j| j.list.retain(|o| o.id != id));
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    #[test]
123    fn one_gpu_job_at_a_time_and_a_job_ends_with_its_handle() {
124        let g0 = super::generation();
125        let a = super::start("test-a", "first", true).unwrap();
126        assert!(super::start("test-b", "second", true).is_err(), "a second GPU job is refused");
127        let cpu = super::start("test-c", "on the CPU", false).unwrap();
128        a.set("first: half", Some(0.5));
129        assert_eq!(super::of("test-a").unwrap().progress, Some(0.5));
130        super::cancel("test-a");
131        assert!(a.cancelled());
132        super::cancel_id(cpu.id);
133        assert!(cpu.cancelled());
134        drop(a);
135        assert!(super::of("test-a").is_none());
136        assert!(super::start("test-b", "second", true).is_ok(), "free again");
137        drop(cpu);
138        assert!(super::generation() > g0);
139    }
140}
kitchen_shell/jobs.rs

kitchen_shell/
jobs.rs
