mold_server/
job_registry.rs

1//! Server-side ledger of in-flight generation jobs.
2//!
3//! The web UI tracks each "Generate" click as a card in `useGenerateStream`
4//! and relies on the SSE stream as the only signal that the work is still
5//! happening. When that stream silently drops (network blip, proxy idle
6//! timeout, server restart, browser tab suspended past keepalive), the card
7//! gets stuck `running` forever because no terminal event arrives.
8//!
9//! `JobRegistry` is the server's authoritative list of "things still owed an
10//! output" — every job between `submit()` and worker completion. The new
11//! `GET /api/queue` endpoint exposes this list so the SPA can poll it and
12//! dead-letter cards whose server-assigned `id` is no longer present.
13//!
14//! The registry deliberately doesn't track *completed* jobs — the gallery DB
15//! is the source of truth for those. Anything in here is currently queued or
16//! actively running on some worker.
17
18use serde::Serialize;
19use std::sync::{Arc, RwLock};
20use std::time::{SystemTime, UNIX_EPOCH};
21use tokio::sync::Notify;
22
23/// Wire-facing job state. Mirrors the actual lifecycle:
24///
25/// - `queued` — accepted by `submit()`, sitting in the channel awaiting a
26///   dispatcher decision OR the dispatcher is mid-retry across workers.
27/// - `running` — handed off to a GPU worker thread; flipping happens when
28///   the worker pulls the job off its channel and starts loading / inferring.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, utoipa::ToSchema)]
30#[serde(rename_all = "snake_case")]
31pub enum JobLifecycle {
32    Queued,
33    Running,
34}
35
36/// One row in the `GET /api/queue` response.
37///
38/// `position` is the 0-based FIFO index at the time of the snapshot — 0 is at
39/// the head (about to be dispatched, or already running on a worker), N-1 is
40/// the most recently submitted. Position is derived from insertion order, so
41/// it shifts as earlier jobs finish and drop out.
42#[derive(Debug, Clone, Serialize, utoipa::ToSchema)]
43pub struct JobEntry {
44    pub id: String,
45    pub model: String,
46    pub state: JobLifecycle,
47    pub started_at_unix_ms: u64,
48    pub position: usize,
49    /// GPU ordinal currently running this job (`null` for queued rows).
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub gpu: Option<usize>,
52    /// Preferred GPU ordinal for queued jobs (`None` means Auto).
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub target_gpu: Option<usize>,
55}
56
57/// Whole-queue listing returned by `GET /api/queue`. Wrapped in a struct so
58/// the response can grow extra fields (totals, byte counts, etc.) without a
59/// breaking change.
60#[derive(Debug, Clone, Serialize, utoipa::ToSchema)]
61pub struct QueueListing {
62    pub entries: Vec<JobEntry>,
63}
64
65#[derive(Debug, Clone)]
66struct EntryInternal {
67    id: String,
68    model: String,
69    state: JobLifecycle,
70    started_at_unix_ms: u64,
71    gpu: Option<usize>,
72    target_gpu: Option<usize>,
73    /// Cancellation signal for `DELETE /api/queue/:id`. The submitting
74    /// handler holds the clone returned by `register*()` and selects on
75    /// `notified()` alongside the job's result channel; `cancel_queued`
76    /// fires `notify_one()` so the permit survives even when the cancel
77    /// lands before the waiter starts awaiting.
78    cancel: Arc<Notify>,
79}
80
81#[derive(Debug, Clone, Copy, PartialEq, Eq)]
82pub enum TargetGpuUpdateError {
83    NotFound,
84    AlreadyRunning,
85}
86
87/// Why a `DELETE /api/queue/:id` cancel attempt was rejected.
88#[derive(Debug, Clone, Copy, PartialEq, Eq)]
89pub enum QueuedJobCancelError {
90    NotFound,
91    AlreadyRunning,
92}
93
94/// The registry itself. Construct via `JobRegistry::new` and share through
95/// `AppState`. All mutation is fire-and-forget — if the inner lock is
96/// poisoned (extremely unlikely in practice) we recover from the inner
97/// state rather than propagating the panic into the dispatcher hot path.
98pub struct JobRegistry {
99    inner: RwLock<Vec<EntryInternal>>,
100}
101
102/// Cheap-cloneable handle. Workers and routes pass this around by value.
103pub type SharedJobRegistry = Arc<JobRegistry>;
104
105impl JobRegistry {
106    pub fn new() -> SharedJobRegistry {
107        Arc::new(Self {
108            inner: RwLock::new(Vec::new()),
109        })
110    }
111
112    /// Insert a freshly-submitted job at the tail in `Queued` state.
113    /// Returns the job's cancellation signal — see `register_with_target_gpu`.
114    pub fn register(&self, id: impl Into<String>, model: impl Into<String>) -> Arc<Notify> {
115        self.register_with_target_gpu(id, model, None)
116    }
117
118    /// Insert a freshly-submitted job with an optional queued lane target.
119    ///
120    /// Returns the job's cancellation signal. The submitting handler must
121    /// hold it and select on `notified()` alongside the result channel —
122    /// `cancel_queued` resolves it when `DELETE /api/queue/:id` removes the
123    /// entry. Callers that never wait (tests poking the registry directly)
124    /// can drop the handle.
125    pub fn register_with_target_gpu(
126        &self,
127        id: impl Into<String>,
128        model: impl Into<String>,
129        target_gpu: Option<usize>,
130    ) -> Arc<Notify> {
131        let started_at_unix_ms = SystemTime::now()
132            .duration_since(UNIX_EPOCH)
133            .unwrap_or_default()
134            .as_millis() as u64;
135        let cancel = Arc::new(Notify::new());
136        let mut entries = self.inner.write().unwrap_or_else(|e| e.into_inner());
137        entries.push(EntryInternal {
138            id: id.into(),
139            model: model.into(),
140            state: JobLifecycle::Queued,
141            started_at_unix_ms,
142            gpu: None,
143            target_gpu,
144            cancel: cancel.clone(),
145        });
146        cancel
147    }
148
149    /// Cancel a still-queued job: remove its entry and fire the cancel
150    /// signal returned by `register*()` so the waiting request future
151    /// resolves with a cancellation error. Running jobs are not cancelable
152    /// — the GPU worker owns them and there is no safe preemption point.
153    ///
154    /// The state check and removal happen under the same write lock that
155    /// `mark_running` takes, so a job can never be both cancelled and
156    /// promoted. (A worker that already dequeued the job before the cancel
157    /// landed will observe the closed result channel and skip it.)
158    pub fn cancel_queued(&self, id: &str) -> Result<(), QueuedJobCancelError> {
159        let mut entries = self.inner.write().unwrap_or_else(|e| e.into_inner());
160        let Some(pos) = entries.iter().position(|e| e.id == id) else {
161            return Err(QueuedJobCancelError::NotFound);
162        };
163        if entries[pos].state == JobLifecycle::Running {
164            return Err(QueuedJobCancelError::AlreadyRunning);
165        }
166        let entry = entries.remove(pos);
167        entry.cancel.notify_one();
168        Ok(())
169    }
170
171    /// Promote a registry entry from `Queued` to `Running`. No-op if `id`
172    /// isn't present (the entry may have been removed concurrently).
173    pub fn mark_running(&self, id: &str, gpu: Option<usize>) {
174        let mut entries = self.inner.write().unwrap_or_else(|e| e.into_inner());
175        if let Some(e) = entries.iter_mut().find(|e| e.id == id) {
176            e.state = JobLifecycle::Running;
177            e.gpu = gpu;
178            e.target_gpu = None;
179        }
180    }
181
182    pub fn set_target_gpu(
183        &self,
184        id: &str,
185        target_gpu: Option<usize>,
186    ) -> Result<(), TargetGpuUpdateError> {
187        let mut entries = self.inner.write().unwrap_or_else(|e| e.into_inner());
188        let Some(e) = entries.iter_mut().find(|e| e.id == id) else {
189            return Err(TargetGpuUpdateError::NotFound);
190        };
191        if e.state == JobLifecycle::Running {
192            return Err(TargetGpuUpdateError::AlreadyRunning);
193        }
194        e.target_gpu = target_gpu;
195        Ok(())
196    }
197
198    pub fn target_gpu(&self, id: &str) -> Option<Option<usize>> {
199        let entries = self.inner.read().unwrap_or_else(|e| e.into_inner());
200        entries.iter().find(|e| e.id == id).map(|e| e.target_gpu)
201    }
202
203    pub fn entry(&self, id: &str) -> Option<JobEntry> {
204        let entries = self.inner.read().unwrap_or_else(|e| e.into_inner());
205        entries.iter().enumerate().find_map(|(i, e)| {
206            (e.id == id).then(|| JobEntry {
207                id: e.id.clone(),
208                model: e.model.clone(),
209                state: e.state,
210                started_at_unix_ms: e.started_at_unix_ms,
211                position: i,
212                gpu: e.gpu,
213                target_gpu: e.target_gpu,
214            })
215        })
216    }
217
218    /// Drop the entry — call once on every terminal path (success, error,
219    /// client-disconnect skip, dispatch failure). Idempotent.
220    pub fn remove(&self, id: &str) {
221        if id.is_empty() {
222            return;
223        }
224        let mut entries = self.inner.write().unwrap_or_else(|e| e.into_inner());
225        entries.retain(|e| e.id != id);
226    }
227
228    /// Snapshot the registry as a wire-shaped listing. Positions are assigned
229    /// in insertion order at snapshot time.
230    pub fn snapshot(&self) -> QueueListing {
231        let entries = self.inner.read().unwrap_or_else(|e| e.into_inner());
232        let out = entries
233            .iter()
234            .enumerate()
235            .map(|(i, e)| JobEntry {
236                id: e.id.clone(),
237                model: e.model.clone(),
238                state: e.state,
239                started_at_unix_ms: e.started_at_unix_ms,
240                position: i,
241                gpu: e.gpu,
242                target_gpu: e.target_gpu,
243            })
244            .collect();
245        QueueListing { entries: out }
246    }
247
248    /// Currently-tracked job count. Exposed for tests and metrics.
249    pub fn len(&self) -> usize {
250        self.inner.read().unwrap_or_else(|e| e.into_inner()).len()
251    }
252
253    /// Returns true when nothing is queued or running. Public so other
254    /// callers (metrics, integration tests) can check emptiness without
255    /// allocating a full snapshot.
256    pub fn is_empty(&self) -> bool {
257        self.len() == 0
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    #[test]
266    fn register_appends_in_fifo_order_with_queued_state() {
267        let reg = JobRegistry::new();
268        reg.register("a", "flux-dev:fp16");
269        reg.register("b", "sdxl:q8");
270        let snap = reg.snapshot();
271        assert_eq!(snap.entries.len(), 2);
272        assert_eq!(snap.entries[0].id, "a");
273        assert_eq!(snap.entries[0].position, 0);
274        assert_eq!(snap.entries[0].state, JobLifecycle::Queued);
275        assert_eq!(snap.entries[1].id, "b");
276        assert_eq!(snap.entries[1].position, 1);
277    }
278
279    #[test]
280    fn mark_running_flips_state_and_records_gpu_ordinal() {
281        let reg = JobRegistry::new();
282        reg.register("a", "flux-dev:fp16");
283        reg.mark_running("a", Some(1));
284        let snap = reg.snapshot();
285        assert_eq!(snap.entries[0].state, JobLifecycle::Running);
286        assert_eq!(snap.entries[0].gpu, Some(1));
287    }
288
289    #[test]
290    fn queued_entries_can_carry_target_gpu_metadata() {
291        let reg = JobRegistry::new();
292        reg.register_with_target_gpu("a", "flux-dev:fp16", Some(1));
293        let snap = reg.snapshot();
294        assert_eq!(snap.entries[0].state, JobLifecycle::Queued);
295        assert_eq!(snap.entries[0].target_gpu, Some(1));
296        assert_eq!(snap.entries[0].gpu, None);
297    }
298
299    #[test]
300    fn target_gpu_updates_only_apply_to_queued_entries() {
301        let reg = JobRegistry::new();
302        reg.register("a", "flux-dev:fp16");
303        reg.set_target_gpu("a", Some(1)).unwrap();
304        assert_eq!(reg.target_gpu("a"), Some(Some(1)));
305
306        reg.mark_running("a", Some(1));
307        let err = reg.set_target_gpu("a", None).unwrap_err();
308        assert_eq!(err, TargetGpuUpdateError::AlreadyRunning);
309        assert_eq!(reg.target_gpu("a"), Some(None));
310    }
311
312    #[test]
313    fn mark_running_is_a_noop_for_unknown_ids() {
314        let reg = JobRegistry::new();
315        reg.register("a", "flux-dev:fp16");
316        // No panic, no insertion — bogus id is ignored entirely.
317        reg.mark_running("not-here", Some(0));
318        let snap = reg.snapshot();
319        assert_eq!(snap.entries.len(), 1);
320        assert_eq!(snap.entries[0].state, JobLifecycle::Queued);
321    }
322
323    #[test]
324    fn remove_compacts_positions_for_the_survivors() {
325        let reg = JobRegistry::new();
326        reg.register("a", "flux-dev:fp16");
327        reg.register("b", "sdxl:q8");
328        reg.register("c", "ltx-video:q8");
329        reg.remove("b");
330        let snap = reg.snapshot();
331        assert_eq!(snap.entries.len(), 2);
332        assert_eq!(snap.entries[0].id, "a");
333        assert_eq!(snap.entries[0].position, 0);
334        assert_eq!(snap.entries[1].id, "c");
335        assert_eq!(snap.entries[1].position, 1);
336    }
337
338    #[test]
339    fn remove_is_idempotent_and_ignores_empty_ids() {
340        // The worker's QueueSlot drop guard removes unconditionally — if the
341        // dispatcher already removed the entry on an error path, the worker's
342        // remove must not panic. Same for jobs that bypassed the registry
343        // entirely (id == "").
344        let reg = JobRegistry::new();
345        reg.register("a", "flux-dev:fp16");
346        reg.remove("a");
347        reg.remove("a"); // second remove is a no-op
348        reg.remove("");
349        reg.remove("never-existed");
350        assert!(reg.is_empty());
351    }
352
353    #[test]
354    fn cancel_queued_removes_the_entry() {
355        let reg = JobRegistry::new();
356        reg.register("a", "flux-dev:fp16");
357        reg.register("b", "sdxl:q8");
358        reg.cancel_queued("a").unwrap();
359        let snap = reg.snapshot();
360        assert_eq!(snap.entries.len(), 1);
361        assert_eq!(snap.entries[0].id, "b");
362        assert_eq!(snap.entries[0].position, 0);
363    }
364
365    #[test]
366    fn cancel_queued_rejects_running_jobs_and_keeps_the_entry() {
367        let reg = JobRegistry::new();
368        reg.register("a", "flux-dev:fp16");
369        reg.mark_running("a", Some(0));
370        let err = reg.cancel_queued("a").unwrap_err();
371        assert_eq!(err, QueuedJobCancelError::AlreadyRunning);
372        assert_eq!(reg.len(), 1, "running entry must survive a cancel attempt");
373    }
374
375    #[test]
376    fn cancel_queued_unknown_id_is_not_found() {
377        let reg = JobRegistry::new();
378        let err = reg.cancel_queued("never-existed").unwrap_err();
379        assert_eq!(err, QueuedJobCancelError::NotFound);
380    }
381
382    #[tokio::test]
383    async fn cancel_queued_signals_the_registered_waiter() {
384        // The handle returned by register() must resolve `notified()` even
385        // when the cancel fires before the waiter starts awaiting — Notify
386        // stores the permit from notify_one().
387        let reg = JobRegistry::new();
388        let cancel = reg.register("a", "flux-dev:fp16");
389        reg.cancel_queued("a").unwrap();
390        tokio::time::timeout(std::time::Duration::from_secs(1), cancel.notified())
391            .await
392            .expect("cancel signal must resolve the waiter");
393    }
394
395    #[test]
396    fn snapshot_serializes_with_snake_case_state_and_omits_gpu_when_queued() {
397        // Wire contract: queued rows must NOT carry a `gpu` field at all
398        // (clients shouldn't see `"gpu": null` and infer GPU 0). The state
399        // tag is lowercase to match the rest of the SSE/JSON style.
400        let reg = JobRegistry::new();
401        reg.register("a", "flux-dev:fp16");
402        let snap = reg.snapshot();
403        let json = serde_json::to_string(&snap.entries[0]).unwrap();
404        assert!(json.contains(r#""state":"queued""#), "got: {json}");
405        assert!(
406            !json.contains("gpu"),
407            "queued row leaked a gpu field: {json}"
408        );
409
410        reg.mark_running("a", Some(0));
411        let snap2 = reg.snapshot();
412        let json2 = serde_json::to_string(&snap2.entries[0]).unwrap();
413        assert!(json2.contains(r#""state":"running""#));
414        assert!(json2.contains(r#""gpu":0"#));
415    }
416}
mold_server/job_registry.rs

mold_server/
job_registry.rs
