studio_worker/
job_gate.rs

1//! The worker's single-job reservation gate.
2//!
3//! A worker owns one GPU, so exactly one generation may run at a time.
4//! Three independent code paths race for it: the WS session (studio
5//! offers), the always-on local API (`POST /image`), and the
6//! auto-updater (which must not `restart_self` mid-job).  Before this
7//! gate the WS session guarded itself with a bare `Arc<AtomicBool>`
8//! CAS while the local API ignored it entirely — so a local job and a
9//! studio job could run concurrently and OOM each other, and an update
10//! could kill an in-flight job.
11//!
12//! [`JobGate`] wraps that shared flag and hands out an RAII
13//! [`JobReservation`] whose `Drop` releases the slot, so no code path
14//! can forget to clear it on an early return or a panic.  The gate is
15//! cheap to clone (an `Arc`) and the reservation is `Send`, so it
16//! moves cleanly into a spawned task for the lifetime of a job.
17
18use std::sync::atomic::{AtomicBool, Ordering};
19use std::sync::Arc;
20
21/// A cloneable handle to the worker's one-job-at-a-time flag.
22#[derive(Clone, Default)]
23pub struct JobGate {
24    busy: Arc<AtomicBool>,
25}
26
27impl JobGate {
28    pub fn new() -> Self {
29        Self::default()
30    }
31
32    /// Adopt an existing shared flag (the runtime already threads one
33    /// `Arc<AtomicBool>` through the WS session + auto-updater).
34    pub fn from_shared(busy: Arc<AtomicBool>) -> Self {
35        Self { busy }
36    }
37
38    /// The underlying flag, for readers that only need to observe
39    /// busyness (e.g. a heartbeat) without reserving.
40    pub fn shared(&self) -> Arc<AtomicBool> {
41        self.busy.clone()
42    }
43
44    /// True while a job holds the slot.
45    pub fn is_busy(&self) -> bool {
46        self.busy.load(Ordering::SeqCst)
47    }
48
49    /// Atomically claim the slot.  `Some(reservation)` means the caller
50    /// owns the worker until the reservation drops; `None` means a job
51    /// is already running and the caller must back off (reject the
52    /// offer / return 503 / skip the update).
53    pub fn try_reserve(&self) -> Option<JobReservation> {
54        self.busy
55            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
56            .ok()
57            .map(|_| JobReservation {
58                busy: self.busy.clone(),
59            })
60    }
61}
62
63/// RAII proof that the holder owns the worker's single job slot.
64/// Releasing on `Drop` means every exit path — success, error, or
65/// panic — frees the slot without an explicit store.
66pub struct JobReservation {
67    busy: Arc<AtomicBool>,
68}
69
70impl Drop for JobReservation {
71    fn drop(&mut self) {
72        self.busy.store(false, Ordering::SeqCst);
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn reserve_is_exclusive_until_dropped() {
82        let gate = JobGate::new();
83        assert!(!gate.is_busy());
84        let reservation = gate.try_reserve().expect("first reserve wins");
85        assert!(gate.is_busy());
86        assert!(
87            gate.try_reserve().is_none(),
88            "a second reserve must fail while the first is held"
89        );
90        drop(reservation);
91        assert!(!gate.is_busy(), "dropping the reservation frees the slot");
92        assert!(gate.try_reserve().is_some(), "the slot is claimable again");
93    }
94
95    #[test]
96    fn clones_share_one_slot() {
97        // The local API, session, and updater hold separate clones of
98        // the same gate — a reservation on one must block the others.
99        let gate = JobGate::new();
100        let other = gate.clone();
101        let _held = gate.try_reserve().unwrap();
102        assert!(other.is_busy());
103        assert!(other.try_reserve().is_none());
104    }
105
106    #[test]
107    fn from_shared_adopts_an_existing_flag() {
108        let flag = Arc::new(AtomicBool::new(false));
109        let gate = JobGate::from_shared(flag.clone());
110        let _held = gate.try_reserve().unwrap();
111        assert!(
112            flag.load(Ordering::SeqCst),
113            "reserving must set the adopted flag so existing readers see it"
114        );
115    }
116
117    #[test]
118    fn reservation_releases_on_panic_unwind() {
119        // A job that panics mid-flight must still free the worker.
120        let gate = JobGate::new();
121        let gate_for_thread = gate.clone();
122        let _ = std::thread::spawn(move || {
123            let _held = gate_for_thread.try_reserve().unwrap();
124            panic!("job blew up");
125        })
126        .join();
127        assert!(
128            !gate.is_busy(),
129            "the slot must be free after a panicking holder unwinds"
130        );
131    }
132}
studio_worker/job_gate.rs

studio_worker/
job_gate.rs
