viewport_lib/resources/upload_jobs/mod.rs
1//! Internal job runner used by the async upload entry points.
2//!
3//! Every long-running upload on `DeviceResources` routes through this
4//! runner. A submitted job runs its CPU work on a background thread, may
5//! optionally submit GPU commands, and reports completion on the main thread
6//! during `process_uploads`. Callers query progress through `upload_status`
7//! and learn about completion either by polling or by attaching a callback.
8//!
9//! No upload entry points use the runner yet. Real submitters will land
10//! alongside the async variants of each existing `upload_*` method.
11
12use std::any::Any;
13use std::collections::{HashMap, VecDeque};
14use std::panic::{AssertUnwindSafe, catch_unwind};
15use std::sync::Arc;
16use std::sync::atomic::{AtomicU32, Ordering};
17use std::sync::mpsc;
18use std::time::{Duration, Instant};
19
20use crate::error::ViewportError;
21
22/// Opaque handle for a submitted job. Returned by every `begin_upload_*`
23/// entry and accepted by status queries, completion callbacks, and the
24/// per-type result accessors.
25#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
26pub struct JobId(u64);
27
28/// Typed result slots for every async upload path, keyed by job id.
29///
30/// Each `begin_upload_*` fills the matching map from its apply closure; the
31/// paired `upload_result_*` drains it. Grouping these here keeps the async
32/// bookkeeping off `DeviceResources` as one field rather than a score of
33/// flat ones. All maps start empty, so the whole struct is `Default`.
34#[derive(Default)]
35pub(crate) struct JobResults {
36 /// Async mesh uploads (`begin_upload_mesh_data` / `upload_result_mesh`).
37 pub mesh: std::sync::Mutex<
38 std::collections::HashMap<JobId, ResultSlot<crate::resources::mesh::mesh_store::MeshId>>,
39 >,
40 /// Async texture uploads (albedo + normal map).
41 pub texture:
42 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<crate::resources::TextureId>>>,
43 /// Jobs submitted through the plugin facade; drained by `Jobs::take<T>`.
44 pub plugin: std::sync::Mutex<
45 std::collections::HashMap<JobId, ResultSlot<Box<dyn std::any::Any + Send>>>,
46 >,
47 /// Async polyline uploads.
48 pub polyline: std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::PolylineId>>>,
49 /// Async streamtube uploads.
50 pub streamtube:
51 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::StreamtubeId>>>,
52 /// Async tube uploads.
53 pub tube: std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::TubeId>>>,
54 /// Async ribbon uploads.
55 pub ribbon: std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::RibbonId>>>,
56 /// Async point cloud uploads.
57 pub point_cloud:
58 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::PointCloudId>>>,
59 /// Async glyph set uploads.
60 pub glyph_set:
61 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::GlyphSetId>>>,
62 /// Async tensor glyph set uploads.
63 pub tensor_glyph_set:
64 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::TensorGlyphSetId>>>,
65 /// Async volume texture uploads.
66 pub volume: std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::VolumeId>>>,
67 /// Async marching-cubes-ready volume uploads.
68 pub volume_mc:
69 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::McVolumeId>>>,
70 /// Async volume-mesh uploads: mesh id plus face-to-cell map.
71 pub volume_mesh: std::sync::Mutex<
72 std::collections::HashMap<
73 JobId,
74 ResultSlot<(crate::resources::mesh::mesh_store::MeshId, Vec<u32>)>,
75 >,
76 >,
77 /// Async clipped-volume-mesh uploads.
78 pub clipped_volume_mesh: std::sync::Mutex<
79 std::collections::HashMap<
80 JobId,
81 ResultSlot<(crate::resources::mesh::mesh_store::MeshId, Vec<u32>)>,
82 >,
83 >,
84 /// Async sparse-volume-grid uploads.
85 pub sparse_volume_grid: std::sync::Mutex<
86 std::collections::HashMap<JobId, ResultSlot<crate::resources::mesh::mesh_store::MeshId>>,
87 >,
88 /// Async projected-tet-mesh uploads: tet id plus packed scalar range.
89 pub projected_tet: std::sync::Mutex<
90 std::collections::HashMap<JobId, ResultSlot<(super::ProjectedTetId, f32, f32)>>,
91 >,
92 /// Async gaussian splat uploads.
93 pub gaussian_splat: std::sync::Mutex<
94 std::collections::HashMap<JobId, ResultSlot<crate::renderer::GaussianSplatId>>,
95 >,
96 /// Async overlay texture uploads.
97 pub overlay_texture: std::sync::Mutex<
98 std::collections::HashMap<JobId, ResultSlot<crate::renderer::OverlayTextureId>>,
99 >,
100 /// Async sprite set uploads.
101 pub sprite_set:
102 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::SpriteSetId>>>,
103 /// Async sprite instance set uploads.
104 pub sprite_instance_set:
105 std::sync::Mutex<std::collections::HashMap<JobId, ResultSlot<super::SpriteInstanceSetId>>>,
106}
107
108/// Current state of a submitted job.
109#[derive(Debug, Clone)]
110pub enum UploadStatus {
111 /// The job is still running. `progress` is a hint in the range 0.0 to
112 /// 1.0. Workers that do not report progress leave the value at zero.
113 Pending {
114 /// Reported completion fraction, between 0.0 and 1.0.
115 progress: f32,
116 },
117 /// The job finished successfully. The caller may take any typed result
118 /// it expects through the matching `upload_result_*` accessor.
119 Ready,
120 /// The worker returned an error, panicked, or dropped its channel
121 /// without sending. The job is not retried.
122 Failed(ViewportError),
123 /// The id has never been issued, has already been reaped, or its result
124 /// was already taken. Treat as "nothing in flight under that id".
125 Unknown,
126}
127
128/// Cheap progress counter shared between the worker thread and the runner.
129///
130/// Workers call `set` with a fraction in 0.0 to 1.0; the runner reads the
131/// value during `process_uploads` to populate `UploadStatus::Pending`.
132#[derive(Clone)]
133pub struct ProgressHandle {
134 inner: Arc<AtomicU32>,
135}
136
137impl ProgressHandle {
138 fn new() -> Self {
139 Self {
140 inner: Arc::new(AtomicU32::new(0)),
141 }
142 }
143
144 /// Record current progress. Values are clamped to 0.0 to 1.0.
145 pub fn set(&self, fraction: f32) {
146 let clamped = fraction.clamp(0.0, 1.0);
147 self.inner.store(clamped.to_bits(), Ordering::Relaxed);
148 }
149
150 fn read(&self) -> f32 {
151 f32::from_bits(self.inner.load(Ordering::Relaxed))
152 }
153}
154
155/// Closure run on the caller's thread once a job's worker (and any GPU
156/// submission) has completed. Real upload types use it to insert their newly
157/// built textures, buffers, and bind groups into `DeviceResources`.
158pub type ApplyFn = Box<dyn FnOnce(&mut super::DeviceResources) + Send>;
159
160/// Boxed GPU-submitting work for a `submit_with_gpu` job.
161///
162/// Unlike CPU jobs, this runs on the main thread inside `process`, not on a
163/// rayon worker: some drivers (notably the NVIDIA Linux Vulkan driver)
164/// corrupt the command pushbuffer (NVRM Xid 32, surfacing as a lost device)
165/// when GPU commands are submitted from a thread other than the one driving
166/// the device. CPU-side preparation still belongs on worker threads via
167/// `submit_cpu`; only the GPU calls are funnelled here.
168type GpuWorkFn = Box<
169 dyn FnOnce(&wgpu::Device, &wgpu::Queue, &ProgressHandle) -> Result<JobProduct, ViewportError>
170 + Send,
171>;
172
173/// A `submit_with_gpu` job awaiting execution on the main thread.
174struct DeferredGpuJob {
175 work: GpuWorkFn,
176 progress: ProgressHandle,
177 tx: mpsc::Sender<WorkerOutcome>,
178}
179
180/// Maximum number of deferred GPU jobs run per `process` call. The cap spreads
181/// a large batch (e.g. a scene's worth of streamed textures) across frames so
182/// a single `process` does not stall the main thread uploading everything at
183/// once.
184const MAX_GPU_JOBS_PER_PROCESS: usize = 16;
185
186/// Per-job result holder shared between a worker's apply closure and the
187/// matching `upload_result_*` accessor.
188///
189/// `ResultSlot<T>` is constructed at submit time on the main thread, cloned
190/// into the apply closure, and used to publish the upload's typed result.
191/// The accessor calls `take` to claim the value once the job reaches
192/// `Ready`.
193pub struct ResultSlot<T> {
194 inner: Arc<std::sync::Mutex<Option<T>>>,
195}
196
197impl<T> Clone for ResultSlot<T> {
198 fn clone(&self) -> Self {
199 Self {
200 inner: Arc::clone(&self.inner),
201 }
202 }
203}
204
205impl<T> ResultSlot<T> {
206 /// Build an empty slot. The apply closure fills it; the accessor takes.
207 pub fn new() -> Self {
208 Self {
209 inner: Arc::new(std::sync::Mutex::new(None)),
210 }
211 }
212
213 /// Store the result. Called from the apply closure on the main thread.
214 pub fn set(&self, value: T) {
215 let mut guard = self.inner.lock().expect("result slot poisoned");
216 *guard = Some(value);
217 }
218
219 /// Take the stored result if one is present, leaving the slot empty.
220 pub fn take(&self) -> Option<T> {
221 let mut guard = self.inner.lock().expect("result slot poisoned");
222 guard.take()
223 }
224}
225
226impl<T> Default for ResultSlot<T> {
227 fn default() -> Self {
228 Self::new()
229 }
230}
231
232/// What the worker hands back to the runner. Bundles whatever GPU
233/// completion the runner should wait on with whatever main-thread mutation
234/// the apply step needs to perform.
235///
236/// Workers that finish purely on the CPU and need no main-thread mutation
237/// return `JobProduct::default()`. Workers that submit GPU commands fill
238/// `gpu`; workers that need to store results on `DeviceResources`
239/// fill `apply`.
240pub struct JobProduct {
241 /// `Some` when the worker has submitted commands that must complete
242 /// before the job can be reported `Ready`. The runner gates on this
243 /// submission via `device.poll`.
244 pub gpu: Option<wgpu::SubmissionIndex>,
245 /// `Some` when the worker has built state that must be folded into
246 /// `DeviceResources` from the main thread.
247 pub apply: Option<ApplyFn>,
248}
249
250impl Default for JobProduct {
251 fn default() -> Self {
252 Self {
253 gpu: None,
254 apply: None,
255 }
256 }
257}
258
259impl JobProduct {
260 /// No GPU gating, no apply step. Convenient for tests and CPU-only
261 /// jobs whose effects are entirely captured in the channel send.
262 pub fn empty() -> Self {
263 Self::default()
264 }
265
266 /// Gate on a single GPU submission; no main-thread apply.
267 pub fn with_gpu(gpu: wgpu::SubmissionIndex) -> Self {
268 Self {
269 gpu: Some(gpu),
270 apply: None,
271 }
272 }
273
274 /// Run an apply step on the main thread; no GPU gating.
275 pub fn with_apply(apply: ApplyFn) -> Self {
276 Self {
277 gpu: None,
278 apply: Some(apply),
279 }
280 }
281
282 /// Gate on a GPU submission, then run the apply step.
283 pub fn with_gpu_and_apply(gpu: wgpu::SubmissionIndex, apply: ApplyFn) -> Self {
284 Self {
285 gpu: Some(gpu),
286 apply: Some(apply),
287 }
288 }
289}
290
291/// Cap on how long `process_uploads_with_budget` is allowed to spend running
292/// completed jobs' apply steps in a single call.
293///
294/// The budget measures only the main-thread apply work (and the wgpu device
295/// poll that precedes it); it does not bound worker-thread time, which runs
296/// concurrently. When the budget is exhausted, any apply closures that have
297/// not yet run are held inside the runner and processed on the next
298/// `process_uploads*` call. Their owning jobs continue to report
299/// `UploadStatus::Pending` until their apply runs, so the typed result is
300/// never observably available before it is in place.
301#[derive(Clone, Copy, Debug)]
302pub struct FrameBudget {
303 deadline: Option<Instant>,
304}
305
306impl FrameBudget {
307 /// No bound. `process_uploads_with_budget` behaves like the unbounded
308 /// `process_uploads`.
309 pub fn unbounded() -> Self {
310 Self { deadline: None }
311 }
312
313 /// Cap apply-step work for the current call to roughly `duration` from
314 /// now. The check happens between applies, so a single long-running
315 /// apply may push past the deadline once it starts.
316 pub fn from_now(duration: Duration) -> Self {
317 Self {
318 deadline: Some(Instant::now() + duration),
319 }
320 }
321
322 /// True when the budget has elapsed.
323 fn exhausted(&self) -> bool {
324 match self.deadline {
325 Some(t) => Instant::now() >= t,
326 None => false,
327 }
328 }
329}
330
331/// Outcome the worker sends through its channel.
332///
333/// The `Duration` captures the wall-clock time the worker spent on the
334/// background thread, measured from the rayon::spawn closure entry to its
335/// return. It excludes both the time the job spent in the rayon queue and
336/// any later GPU/apply-step work; the runner adds the apply-step duration
337/// on top.
338#[allow(dead_code)]
339enum WorkerOutcome {
340 Done(JobProduct, Duration),
341 Failed(ViewportError, Duration),
342}
343
344type CompletionCallback = Box<dyn FnOnce(&UploadStatus) + Send>;
345
346/// What `process` hands back for a single completed job. The caller is
347/// responsible for running `apply` (if present and the status is `Ready`)
348/// against the live `DeviceResources`, then invoking `callback`.
349pub struct Completion {
350 /// Id of the completed job; the caller uses it to record the apply
351 /// duration back on the runner.
352 pub id: JobId,
353 /// Final status the runner observed.
354 pub status: UploadStatus,
355 /// Apply closure produced by the worker. Run only when `status` is
356 /// `Ready`.
357 pub apply: Option<ApplyFn>,
358 /// Completion callback registered via `on_complete`. Fires for both
359 /// `Ready` and `Failed` so the consumer can branch.
360 pub callback: Option<CompletionCallback>,
361}
362
363struct JobSlot {
364 progress: ProgressHandle,
365 rx: mpsc::Receiver<WorkerOutcome>,
366 /// Once the worker has reported, the GPU submission to gate on (if any)
367 /// and the apply closure to run when the GPU side finishes.
368 awaiting: Option<(Option<wgpu::SubmissionIndex>, Option<ApplyFn>)>,
369 callback: Option<CompletionCallback>,
370}
371
372/// Background worker pool plus the table of in-flight jobs.
373///
374/// Owned by `DeviceResources`; reached from there via
375/// `process_uploads`, `upload_status`, and friends. `next_id` and the
376/// `submit_*` helpers are unused until upload entry points are wired
377/// through the runner.
378#[allow(dead_code)]
379pub struct JobRunner {
380 next_id: u64,
381 slots: HashMap<u64, JobSlot>,
382 /// Jobs whose worker and GPU work have finished but whose main-thread
383 /// apply closure has not yet been run.
384 ///
385 /// Populated by `process` for successful jobs that carry an apply step;
386 /// drained by the caller of `process_uploads` (and friends). Entries
387 /// remain visible as `UploadStatus::Pending { progress: 1.0 }` until
388 /// the apply runs, so the typed result is never reported as `Ready`
389 /// before it is materialized in the resource state.
390 pending_apply: VecDeque<PendingApply>,
391 /// Recently finished jobs, kept for one drain cycle so callers can still
392 /// see `Ready` or `Failed` after the completion frame.
393 finished: HashMap<u64, UploadStatus>,
394 /// Time the actual work took. Worker thread time is recorded when the
395 /// worker reports back; apply-step time is added by the caller via
396 /// `add_apply_duration` after running the apply closure. Retained
397 /// until `drop_duration` is called so consumers have at least one frame
398 /// to read the result.
399 durations: HashMap<u64, Duration>,
400 /// GPU-submitting jobs deferred to run on the main thread during
401 /// `process`. See [`GpuWorkFn`] for why these must not run on a worker.
402 deferred_gpu: VecDeque<DeferredGpuJob>,
403}
404
405/// Entry on the `pending_apply` queue.
406///
407/// Holds everything the caller needs to fold a successful job into resource
408/// state on the main thread: the id (for status updates and callback
409/// registration), the final status, the apply closure, and any completion
410/// callback that was registered on the job. Failed jobs do not produce
411/// these; they go straight to `finished` and surface in the `Completion`
412/// vec returned by `process`.
413pub struct PendingApply {
414 /// Id of the job whose apply is pending.
415 pub id: JobId,
416 /// Terminal status to record once the apply finishes; always
417 /// `UploadStatus::Ready` for entries on the queue.
418 pub status: UploadStatus,
419 /// Closure that mutates `DeviceResources` and fills any typed
420 /// result slot.
421 pub apply: ApplyFn,
422 /// Completion callback registered via `on_complete`.
423 pub callback: Option<CompletionCallback>,
424}
425
426impl Default for JobRunner {
427 fn default() -> Self {
428 Self::new()
429 }
430}
431
432#[allow(dead_code)]
433impl JobRunner {
434 /// Construct an empty runner. The `DeviceResources` initializer
435 /// holds the single instance; callers do not construct this directly.
436 pub fn new() -> Self {
437 Self {
438 next_id: 1,
439 slots: HashMap::new(),
440 pending_apply: VecDeque::new(),
441 finished: HashMap::new(),
442 durations: HashMap::new(),
443 deferred_gpu: VecDeque::new(),
444 }
445 }
446
447 /// Total work duration recorded for a job, or `None` if the job is
448 /// still in flight (or its duration record has aged out).
449 pub fn duration(&self, id: JobId) -> Option<Duration> {
450 self.durations.get(&id.0).copied()
451 }
452
453 /// Add the apply-step elapsed time to a job's running total. Called by
454 /// the caller of `process` immediately after running the apply closure
455 /// on the main thread.
456 pub fn add_apply_duration(&mut self, id: JobId, apply: Duration) {
457 let entry = self.durations.entry(id.0).or_insert(Duration::ZERO);
458 *entry = entry.saturating_add(apply);
459 }
460
461 /// Drop the recorded duration for a job. Consumers call this after they
462 /// have read the duration via `duration`; otherwise the runner keeps it
463 /// across drain cycles so a single-frame retention is not enough.
464 pub fn drop_duration(&mut self, id: JobId) {
465 self.durations.remove(&id.0);
466 }
467
468 fn issue_id(&mut self) -> JobId {
469 let id = self.next_id;
470 self.next_id = self
471 .next_id
472 .checked_add(1)
473 .expect("upload job id space exhausted");
474 JobId(id)
475 }
476
477 /// Schedule a CPU-only job on the background pool.
478 ///
479 /// The worker receives a `ProgressHandle` it can use to publish progress.
480 /// Returning `Err` or panicking marks the job as `Failed`.
481 pub(crate) fn submit_cpu<F>(&mut self, work: F) -> JobId
482 where
483 F: FnOnce(&ProgressHandle) -> Result<JobProduct, ViewportError> + Send + 'static,
484 {
485 let id = self.issue_id();
486 let progress = ProgressHandle::new();
487 let worker_progress = progress.clone();
488 let (tx, rx) = mpsc::channel();
489
490 rayon::spawn(move || {
491 let t0 = Instant::now();
492 let outcome = match catch_unwind(AssertUnwindSafe(|| work(&worker_progress))) {
493 Ok(Ok(product)) => WorkerOutcome::Done(product, t0.elapsed()),
494 Ok(Err(e)) => WorkerOutcome::Failed(e, t0.elapsed()),
495 Err(_) => WorkerOutcome::Failed(
496 ViewportError::JobWorkerLost {
497 reason: "worker panicked",
498 },
499 t0.elapsed(),
500 ),
501 };
502 // Receiver going away is fine; the runner was probably dropped.
503 let _ = tx.send(outcome);
504 });
505
506 self.slots.insert(
507 id.0,
508 JobSlot {
509 progress,
510 rx,
511 awaiting: None,
512 callback: None,
513 },
514 );
515 id
516 }
517
518 /// Schedule a CPU job that also submits GPU commands.
519 ///
520 /// The worker is handed cloned `Device` and `Queue` handles. It may
521 /// submit any number of command buffers and bundles the final
522 /// `SubmissionIndex` into the `JobProduct` it returns; the runner waits
523 /// on that submission before reporting `Ready`.
524 pub(crate) fn submit_with_gpu<F>(
525 &mut self,
526 device: &wgpu::Device,
527 queue: &wgpu::Queue,
528 work: F,
529 ) -> JobId
530 where
531 F: FnOnce(
532 &wgpu::Device,
533 &wgpu::Queue,
534 &ProgressHandle,
535 ) -> Result<JobProduct, ViewportError>
536 + Send
537 + 'static,
538 {
539 let id = self.issue_id();
540 let progress = ProgressHandle::new();
541 let (tx, rx) = mpsc::channel();
542 // GPU submission is deferred to the main thread (run in `process`)
543 // rather than spawned on a rayon worker: submitting from a non-device
544 // thread corrupts the pushbuffer on some drivers. The `device` and
545 // `queue` arguments are unused here for the same reason; `process`
546 // supplies them when the work runs.
547 let _ = (device, queue);
548 self.deferred_gpu.push_back(DeferredGpuJob {
549 work: Box::new(work),
550 progress: progress.clone(),
551 tx,
552 });
553
554 self.slots.insert(
555 id.0,
556 JobSlot {
557 progress,
558 rx,
559 awaiting: None,
560 callback: None,
561 },
562 );
563 id
564 }
565
566 /// Attach a callback to fire on completion. The callback runs on the
567 /// main thread during the same `process_uploads` call that marks the job
568 /// done. If the job has already completed and is still in the
569 /// short-retention window, the callback fires immediately.
570 pub fn on_complete<F>(&mut self, id: JobId, cb: F)
571 where
572 F: FnOnce(&UploadStatus) + Send + 'static,
573 {
574 if let Some(slot) = self.slots.get_mut(&id.0) {
575 slot.callback = Some(Box::new(cb));
576 return;
577 }
578 if let Some(status) = self.finished.get(&id.0) {
579 cb(status);
580 }
581 }
582
583 /// Look up current state. Returns `Unknown` for ids that have never been
584 /// issued or have been reaped past the retention window.
585 pub fn status(&self, id: JobId) -> UploadStatus {
586 if let Some(slot) = self.slots.get(&id.0) {
587 return UploadStatus::Pending {
588 progress: slot.progress.read(),
589 };
590 }
591 // A job whose worker is done but whose apply step has not yet
592 // run is reported as Pending at full progress. The typed result
593 // is only available after apply runs, so we keep callers in the
594 // Pending arm until then.
595 if self.pending_apply.iter().any(|pa| pa.id.0 == id.0) {
596 return UploadStatus::Pending { progress: 1.0 };
597 }
598 if let Some(status) = self.finished.get(&id.0) {
599 return status.clone();
600 }
601 UploadStatus::Unknown
602 }
603
604 /// Count of jobs still in flight, ignoring the retention window.
605 /// Includes jobs whose worker has finished but whose apply step has
606 /// not yet been drained.
607 pub fn pending(&self) -> usize {
608 self.slots.len() + self.pending_apply.len()
609 }
610
611 /// True when no jobs are in flight.
612 pub fn all_complete(&self) -> bool {
613 self.slots.is_empty() && self.pending_apply.is_empty()
614 }
615
616 /// Pop the next apply closure off the queue. Returns `None` when the
617 /// queue is empty. The caller is expected to run the closure against
618 /// `DeviceResources` and then call `mark_applied` so the job
619 /// transitions from Pending to Ready.
620 pub fn pop_pending_apply(&mut self) -> Option<PendingApply> {
621 self.pending_apply.pop_front()
622 }
623
624 /// Push an entry back onto the front of the queue. Used by
625 /// `process_uploads_with_budget` when the time budget runs out
626 /// mid-drain so the next call picks up where this one stopped.
627 pub fn requeue_pending_apply(&mut self, pa: PendingApply) {
628 self.pending_apply.push_front(pa);
629 }
630
631 /// Record that a pending-apply entry's closure has finished running.
632 /// Moves the job into the short-retention `finished` table so the
633 /// next `status` query reports `Ready`.
634 pub fn mark_applied(&mut self, id: JobId, status: UploadStatus) {
635 self.finished.insert(id.0, status);
636 }
637
638 /// Count of jobs sitting on the apply queue. Exposed for tests and
639 /// metrics; `pending` already aggregates it into the in-flight total.
640 pub fn pending_apply_len(&self) -> usize {
641 self.pending_apply.len()
642 }
643
644 /// Walk the job table, advance any worker results, and check pending
645 /// GPU submissions for completion.
646 ///
647 /// Returns one `Completion` per job that just transitioned to `Ready` or
648 /// `Failed`. The caller is expected to run any `apply` closure first
649 /// (only when `status` is `Ready`), then invoke the registered
650 /// `callback` if present. Splitting these out lets the caller drop any
651 /// external lock around the runner before mutating renderer state.
652 pub fn process(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) -> Vec<Completion> {
653 // Drop the previous frame's retention window. Callers that needed
654 // those results have already taken them.
655 self.finished.clear();
656
657 // Run deferred GPU jobs on this (the device-owning) thread, bounded so
658 // a large batch spreads across frames. Each result is sent into the
659 // job's channel, picked up by the slot loop below in this same call.
660 let n = self.deferred_gpu.len().min(MAX_GPU_JOBS_PER_PROCESS);
661 for _ in 0..n {
662 let Some(job) = self.deferred_gpu.pop_front() else {
663 break;
664 };
665 let t0 = Instant::now();
666 let outcome = match catch_unwind(AssertUnwindSafe(|| {
667 (job.work)(device, queue, &job.progress)
668 })) {
669 Ok(Ok(product)) => WorkerOutcome::Done(product, t0.elapsed()),
670 Ok(Err(e)) => WorkerOutcome::Failed(e, t0.elapsed()),
671 Err(_) => WorkerOutcome::Failed(
672 ViewportError::JobWorkerLost {
673 reason: "gpu job panicked",
674 },
675 t0.elapsed(),
676 ),
677 };
678 let _ = job.tx.send(outcome);
679 }
680
681 // Advance internal wgpu state so completed submissions are visible
682 // to the per-submission wait below.
683 let _ = device.poll(wgpu::PollType::Poll);
684
685 let mut completions = Vec::new();
686 let ids: Vec<u64> = self.slots.keys().copied().collect();
687 for id in ids {
688 // Stage 1: pick up the worker result if we have not already.
689 if self.slots.get(&id).is_some_and(|s| s.awaiting.is_none()) {
690 let outcome = self
691 .slots
692 .get(&id)
693 .map(|s| s.rx.try_recv())
694 .expect("slot existed");
695 match outcome {
696 Ok(WorkerOutcome::Done(product, worker_dur)) => {
697 self.durations.insert(id, worker_dur);
698 let JobProduct { gpu, apply } = product;
699 match gpu {
700 None => {
701 self.finish(id, UploadStatus::Ready, apply, &mut completions);
702 continue;
703 }
704 Some(sub) => {
705 if let Some(slot) = self.slots.get_mut(&id) {
706 slot.awaiting = Some((Some(sub), apply));
707 }
708 }
709 }
710 }
711 Ok(WorkerOutcome::Failed(e, worker_dur)) => {
712 self.durations.insert(id, worker_dur);
713 self.finish(id, UploadStatus::Failed(e), None, &mut completions);
714 continue;
715 }
716 Err(mpsc::TryRecvError::Empty) => continue,
717 Err(mpsc::TryRecvError::Disconnected) => {
718 // Sender dropped without sending. Catch-unwind in
719 // the spawn closure already covers panics, so this
720 // is an unexpected drop path.
721 self.finish(
722 id,
723 UploadStatus::Failed(ViewportError::JobWorkerLost {
724 reason: "worker channel closed without result",
725 }),
726 None,
727 &mut completions,
728 );
729 continue;
730 }
731 }
732 }
733
734 // Stage 2: worker reported a GPU submission; poll for it.
735 let pending_sub = self
736 .slots
737 .get(&id)
738 .and_then(|s| s.awaiting.as_ref())
739 .and_then(|(g, _)| g.clone());
740 if let Some(sub) = pending_sub {
741 let result = device.poll(wgpu::PollType::Wait {
742 submission_index: Some(sub),
743 timeout: Some(Duration::from_millis(0)),
744 });
745 match result {
746 Ok(wgpu::PollStatus::QueueEmpty) | Ok(wgpu::PollStatus::WaitSucceeded) => {
747 let apply = self
748 .slots
749 .get_mut(&id)
750 .and_then(|s| s.awaiting.take())
751 .and_then(|(_, a)| a);
752 self.finish(id, UploadStatus::Ready, apply, &mut completions);
753 }
754 Ok(wgpu::PollStatus::Poll) => {
755 // Backend still working; check again next frame.
756 }
757 Err(_) => {
758 // Timeout or device error. Leave the slot pending
759 // and try again on the next call.
760 }
761 }
762 }
763 }
764 completions
765 }
766
767 fn finish(
768 &mut self,
769 id: u64,
770 status: UploadStatus,
771 apply: Option<ApplyFn>,
772 completions: &mut Vec<Completion>,
773 ) {
774 let Some(mut slot) = self.slots.remove(&id) else {
775 return;
776 };
777 let callback = slot.callback.take();
778 // Successful jobs that carry an apply step are held on the
779 // pending_apply queue. They keep reporting Pending until the
780 // caller runs the apply (via process_uploads or
781 // process_uploads_with_budget); only then do they move into
782 // `finished` and start reporting Ready. Failures and no-apply
783 // successes go through the standard Completion path so callbacks
784 // fire immediately.
785 match (status, apply) {
786 (UploadStatus::Ready, Some(apply_fn)) => {
787 self.pending_apply.push_back(PendingApply {
788 id: JobId(id),
789 status: UploadStatus::Ready,
790 apply: apply_fn,
791 callback,
792 });
793 return;
794 }
795 (status, _) => {
796 completions.push(Completion {
797 id: JobId(id),
798 status: status.clone(),
799 apply: None,
800 callback,
801 });
802 self.finished.insert(id, status);
803 return;
804 }
805 }
806 }
807}
808
809impl super::DeviceResources {
810 /// Advance the upload-job runner. Worker results received since the
811 /// previous call are observed, GPU submissions are polled, completed
812 /// jobs are folded into renderer state, and any completion callbacks
813 /// fire on the caller's thread.
814 ///
815 /// Apply closures and callbacks both run after the runner's mutex is
816 /// released, so they are free to query the runner or submit a fresh job
817 /// without risk of deadlock.
818 pub fn process_uploads(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
819 self.process_uploads_with_budget(device, queue, FrameBudget::unbounded());
820 }
821
822 /// Advance the upload-job runner with a cap on per-call apply-step
823 /// work.
824 ///
825 /// Behaves the same as `process_uploads` except that, after the
826 /// runner has been advanced and any failure callbacks fired, the
827 /// caller stops popping apply closures off the queue once `budget`
828 /// elapses. Remaining apply closures stay on the queue and are
829 /// picked up by the next call. Their owning jobs continue to report
830 /// `UploadStatus::Pending` until their apply runs, so the typed
831 /// result is never observably available before it is in place.
832 ///
833 /// The budget covers only the main-thread apply work and the
834 /// preceding device poll. Worker-thread time is independent. The
835 /// check happens between applies, so a single long-running apply
836 /// may push past the deadline once it has started; the budget is a
837 /// soft cap, not a hard deadline.
838 pub fn process_uploads_with_budget(
839 &mut self,
840 device: &wgpu::Device,
841 queue: &wgpu::Queue,
842 budget: FrameBudget,
843 ) {
844 // Stage 1: advance the runner and drain immediate (failure /
845 // no-apply) completions. Their callbacks fire here regardless of
846 // budget: they do no main-thread work and dropping them would
847 // hide errors from consumers.
848 let completions = {
849 let mut runner = self.jobs.lock().expect("upload job runner poisoned");
850 runner.process(device, queue)
851 };
852 for Completion {
853 id: _,
854 status,
855 apply: _,
856 callback,
857 } in completions
858 {
859 if let Some(cb) = callback {
860 cb(&status);
861 }
862 }
863
864 // Stage 2: drain the apply queue under the budget. Each apply
865 // mutates `self`, so we pop one at a time and re-check the
866 // budget between iterations.
867 loop {
868 if budget.exhausted() {
869 break;
870 }
871 let next = {
872 let mut runner = self.jobs.lock().expect("upload job runner poisoned");
873 runner.pop_pending_apply()
874 };
875 let Some(PendingApply {
876 id,
877 status,
878 apply,
879 callback,
880 }) = next
881 else {
882 break;
883 };
884 let t = Instant::now();
885 apply(self);
886 let apply_d = t.elapsed();
887 {
888 let mut runner = self.jobs.lock().expect("upload job runner poisoned");
889 runner.add_apply_duration(id, apply_d);
890 runner.mark_applied(id, status.clone());
891 }
892 if let Some(cb) = callback {
893 cb(&status);
894 }
895 }
896 }
897
898 /// Total wall-clock work duration for a completed job.
899 ///
900 /// The value is the sum of the worker thread's elapsed time and the
901 /// main-thread apply-step elapsed time. It excludes frame-pacing
902 /// delays (the time the apply step sat waiting for `process_uploads`
903 /// to be called). For sync uploads this method returns `None`; sync
904 /// callers measure their own wall-clock around the call.
905 ///
906 /// Returns `None` for jobs that are still in flight, were never issued,
907 /// or whose duration record has already been dropped via
908 /// `drop_job_duration`. The runner retains durations until the consumer
909 /// drops them so single-frame retention is not enough.
910 pub fn job_duration(&self, id: JobId) -> Option<Duration> {
911 let runner = self.jobs.lock().expect("upload job runner poisoned");
912 runner.duration(id)
913 }
914
915 /// Drop the recorded duration for `id`. Call this after reading the
916 /// value via `job_duration`; otherwise the duration table grows over
917 /// time.
918 pub fn drop_job_duration(&mut self, id: JobId) {
919 let mut runner = self.jobs.lock().expect("upload job runner poisoned");
920 runner.drop_duration(id);
921 }
922
923 /// Look up the current state of a submitted upload job.
924 pub fn upload_status(&self, id: JobId) -> UploadStatus {
925 let runner = self.jobs.lock().expect("upload job runner poisoned");
926 runner.status(id)
927 }
928
929 /// Number of upload jobs still in flight.
930 pub fn uploads_pending(&self) -> usize {
931 let runner = self.jobs.lock().expect("upload job runner poisoned");
932 runner.pending()
933 }
934
935 /// True when no upload jobs are in flight.
936 pub fn all_uploads_complete(&self) -> bool {
937 let runner = self.jobs.lock().expect("upload job runner poisoned");
938 runner.all_complete()
939 }
940
941 /// Register a callback to fire when a job finishes. The callback runs on
942 /// the caller's thread during the next `process_uploads` call. If the
943 /// job has already finished and is still in the short retention window,
944 /// the callback fires immediately on the calling thread.
945 pub fn on_upload_complete<F>(&mut self, id: JobId, cb: F)
946 where
947 F: FnOnce(&UploadStatus) + Send + 'static,
948 {
949 let mut runner = self.jobs.lock().expect("upload job runner poisoned");
950 runner.on_complete(id, cb);
951 }
952
953 /// Block the calling thread, driving `process_uploads` until `id` reaches
954 /// a terminal state.
955 ///
956 /// Returns `Ok(())` when the job's worker (and any GPU submission it
957 /// queued) completes successfully. Returns the worker error when the job
958 /// fails. Used internally by the synchronous `upload_*` entries to wrap
959 /// their `begin_upload_*` counterparts in a single round-trip; consumers
960 /// who want to wait on a specific async upload can call it directly.
961 ///
962 /// The loop sleeps for a short interval between polls so it does not pin
963 /// a CPU core while the worker is running. It does not call back into the
964 /// caller's event loop; if you have other work to interleave, drive
965 /// `process_uploads` yourself instead.
966 ///
967 /// # Errors
968 ///
969 /// Returns [`ViewportError::JobResultMissing`](crate::error::ViewportError::JobResultMissing)
970 /// if `id` has already been reaped or was never issued, and the worker's
971 /// error verbatim when the job ends in `Failed`.
972 pub fn drain_until(
973 &mut self,
974 device: &wgpu::Device,
975 queue: &wgpu::Queue,
976 id: JobId,
977 ) -> crate::error::ViewportResult<()> {
978 loop {
979 self.process_uploads(device, queue);
980 match self.upload_status(id) {
981 UploadStatus::Ready => return Ok(()),
982 UploadStatus::Failed(e) => return Err(e),
983 UploadStatus::Pending { .. } => {
984 std::thread::sleep(Duration::from_micros(200));
985 }
986 UploadStatus::Unknown => {
987 return Err(crate::error::ViewportError::JobResultMissing {
988 reason: "drain target was already reaped",
989 });
990 }
991 }
992 }
993 }
994}
995
996// ---------------------------------------------------------------------------
997// Optional `future` feature: a thin Future wrapper around a JobId.
998// ---------------------------------------------------------------------------
999
1000/// Future returned by [`DeviceResources::upload_handle`].
1001///
1002/// Polling drives the wrapped job to completion using the standard
1003/// `process_uploads` machinery; the consumer's main loop must keep calling
1004/// `process_uploads` so the runner can deliver completion callbacks. Once
1005/// the future resolves, take the typed result with the matching
1006/// `upload_result_*` accessor.
1007#[cfg(feature = "future")]
1008pub struct JobHandle {
1009 id: JobId,
1010 state: Arc<std::sync::Mutex<JobHandleState>>,
1011}
1012
1013#[cfg(feature = "future")]
1014struct JobHandleState {
1015 done: Option<crate::error::ViewportResult<()>>,
1016 waker: Option<std::task::Waker>,
1017}
1018
1019#[cfg(feature = "future")]
1020impl JobHandle {
1021 /// Id of the wrapped job. Pass it to `upload_result_*` after the future
1022 /// resolves.
1023 pub fn id(&self) -> JobId {
1024 self.id
1025 }
1026}
1027
1028#[cfg(feature = "future")]
1029impl std::future::Future for JobHandle {
1030 type Output = crate::error::ViewportResult<()>;
1031
1032 fn poll(
1033 self: std::pin::Pin<&mut Self>,
1034 cx: &mut std::task::Context<'_>,
1035 ) -> std::task::Poll<Self::Output> {
1036 let mut guard = self.state.lock().expect("job handle poisoned");
1037 if let Some(result) = guard.done.take() {
1038 return std::task::Poll::Ready(result);
1039 }
1040 guard.waker = Some(cx.waker().clone());
1041 std::task::Poll::Pending
1042 }
1043}
1044
1045#[cfg(feature = "future")]
1046impl super::DeviceResources {
1047 /// Wrap a `JobId` in a future that resolves when the job completes.
1048 ///
1049 /// The future is driven by completion callbacks fired during
1050 /// `process_uploads`, so the consumer's main loop must keep calling
1051 /// `process_uploads` for the future to make progress. The resolved
1052 /// value is `Ok(())` on success and the worker error on failure; the
1053 /// caller takes the typed result through the matching
1054 /// `upload_result_*` accessor after `.await` returns.
1055 pub fn upload_handle(&mut self, id: JobId) -> JobHandle {
1056 let state = Arc::new(std::sync::Mutex::new(JobHandleState {
1057 done: None,
1058 waker: None,
1059 }));
1060 let state_for_cb = state.clone();
1061 self.on_upload_complete(id, move |status| {
1062 let result = match status {
1063 UploadStatus::Ready => Ok(()),
1064 UploadStatus::Failed(e) => Err(e.clone()),
1065 UploadStatus::Pending { .. } => {
1066 // Callbacks only fire on terminal transitions; this
1067 // arm is unreachable but we report it cleanly rather
1068 // than panic if a future runner change relaxes that.
1069 Err(crate::error::ViewportError::JobNotReady)
1070 }
1071 UploadStatus::Unknown => Err(crate::error::ViewportError::JobResultMissing {
1072 reason: "job vanished before completion callback fired",
1073 }),
1074 };
1075 let mut guard = state_for_cb.lock().expect("job handle poisoned");
1076 guard.done = Some(result);
1077 if let Some(waker) = guard.waker.take() {
1078 waker.wake();
1079 }
1080 });
1081 JobHandle { id, state }
1082 }
1083}
1084
1085// ---------------------------------------------------------------------------
1086// Plugin-facing facade
1087// ---------------------------------------------------------------------------
1088
1089/// Result slot type used by the plugin facade. Values are boxed because the
1090/// runner has no compile-time knowledge of the closure's return type.
1091type PluginResultSlot = ResultSlot<Box<dyn Any + Send>>;
1092
1093/// Plugin-facing handle to the upload-job runner. Exposed to
1094/// `ItemTypePlugin::prepare` via `ItemFrameContext::jobs`.
1095///
1096/// Plugins use it the same way built-in uploads do, but with a typed
1097/// generic return: submit a CPU job that produces a value of type `T`,
1098/// poll the returned `JobId` next frame, and `take::<T>()` the result
1099/// once the status is `Ready`.
1100///
1101/// The handle is `Copy`-like in spirit (it just holds a reference); both
1102/// `submit_cpu` and the readers use `&self` so the plugin does not need to
1103/// thread mutability through its own state.
1104pub struct Jobs<'a> {
1105 resources: &'a super::DeviceResources,
1106}
1107
1108impl<'a> Jobs<'a> {
1109 pub(crate) fn new(resources: &'a super::DeviceResources) -> Self {
1110 Self { resources }
1111 }
1112
1113 /// Schedule a CPU job whose result is delivered through `take<T>`.
1114 ///
1115 /// `work` runs on a background worker. The closure must own its
1116 /// inputs because the `&Device` and `&Queue` references passed to
1117 /// `ItemTypePlugin::prepare` are not `'static`. Panics inside `work`
1118 /// surface as `UploadStatus::Failed(JobWorkerLost)`.
1119 pub fn submit_cpu<T, F>(&self, work: F) -> JobId
1120 where
1121 T: Send + 'static,
1122 F: FnOnce() -> T + Send + 'static,
1123 {
1124 let slot: PluginResultSlot = ResultSlot::new();
1125 let slot_for_apply = slot.clone();
1126
1127 let id = {
1128 let mut runner = self
1129 .resources
1130 .jobs
1131 .lock()
1132 .expect("upload job runner poisoned");
1133 runner.submit_cpu(move |_progress| {
1134 let value: T = work();
1135 let boxed: Box<dyn Any + Send> = Box::new(value);
1136 Ok(JobProduct::with_apply(Box::new(
1137 move |_resources: &mut super::DeviceResources| {
1138 slot_for_apply.set(boxed);
1139 },
1140 )))
1141 })
1142 };
1143
1144 self.resources
1145 .job_results
1146 .plugin
1147 .lock()
1148 .expect("plugin job result map poisoned")
1149 .insert(id, slot);
1150 id
1151 }
1152
1153 /// Current state of a submitted plugin job. Same shape as the
1154 /// `upload_status` reported by built-in uploads.
1155 pub fn status(&self, id: JobId) -> UploadStatus {
1156 self.resources.upload_status(id)
1157 }
1158
1159 /// Try to take the typed result produced by a completed job.
1160 ///
1161 /// Returns `None` while the job is still in flight, when the id has
1162 /// already been taken, when it never existed, or when `T` does not
1163 /// match the type stored by the worker (plugin author error). On a
1164 /// successful take, the slot is removed.
1165 pub fn take<T: Any + Send + 'static>(&self, id: JobId) -> Option<T> {
1166 let mut map = self
1167 .resources
1168 .job_results
1169 .plugin
1170 .lock()
1171 .expect("plugin job result map poisoned");
1172 let slot = map.get(&id)?.clone();
1173 let boxed = slot.take()?;
1174 match boxed.downcast::<T>() {
1175 Ok(value) => {
1176 map.remove(&id);
1177 Some(*value)
1178 }
1179 Err(_boxed_back) => {
1180 // Type mismatch. Re-insert the box by recreating the
1181 // slot with the same value so the plugin can retry with
1182 // the right type or just leak the slot in this dropped
1183 // state. To keep things simple here, we drop on
1184 // mismatch -- plugins should not be calling take with
1185 // the wrong type.
1186 map.remove(&id);
1187 None
1188 }
1189 }
1190 }
1191}
1192
1193#[cfg(test)]
1194mod tests {
1195 use std::sync::Mutex;
1196 use std::sync::atomic::AtomicBool;
1197 use std::time::Duration;
1198
1199 use super::*;
1200
1201 /// Drive the runner until `predicate` is true or the deadline expires.
1202 /// Necessary because parallel-running test threads contend for the
1203 /// rayon pool, so a single sleep + process cycle is not enough.
1204 fn drain_until<F>(
1205 runner: &mut JobRunner,
1206 device: &wgpu::Device,
1207 queue: &wgpu::Queue,
1208 max_iterations: usize,
1209 mut predicate: F,
1210 ) where
1211 F: FnMut(&JobRunner) -> bool,
1212 {
1213 for _ in 0..max_iterations {
1214 let _ = runner.process(device, queue);
1215 if predicate(runner) {
1216 return;
1217 }
1218 std::thread::sleep(Duration::from_millis(5));
1219 }
1220 }
1221
1222 #[test]
1223 fn cpu_job_reports_ready_after_drain() {
1224 let mut runner = JobRunner::new();
1225 let id = runner.submit_cpu(|_p| Ok(JobProduct::empty()));
1226
1227 assert_eq!(runner.pending(), 1);
1228 with_test_gpu(|device, queue| {
1229 drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
1230 });
1231
1232 assert!(matches!(runner.status(id), UploadStatus::Ready));
1233 assert_eq!(runner.pending(), 0);
1234 assert!(runner.all_complete());
1235 }
1236
1237 #[test]
1238 fn cpu_job_progress_is_observable() {
1239 let mut runner = JobRunner::new();
1240 let gate = Arc::new(AtomicBool::new(false));
1241 let gate_for_worker = gate.clone();
1242
1243 let id = runner.submit_cpu(move |p| {
1244 p.set(0.25);
1245 while !gate_for_worker.load(Ordering::Relaxed) {
1246 std::thread::sleep(Duration::from_millis(1));
1247 }
1248 p.set(1.0);
1249 Ok(JobProduct::empty())
1250 });
1251
1252 // Poll until the worker publishes its first progress sample. Under
1253 // contention with other tests, the worker may not start for some
1254 // time.
1255 let mut observed = None;
1256 for _ in 0..200 {
1257 if let UploadStatus::Pending { progress } = runner.status(id) {
1258 if progress > 0.0 {
1259 observed = Some(progress);
1260 break;
1261 }
1262 }
1263 std::thread::sleep(Duration::from_millis(5));
1264 }
1265 let progress = observed.expect("worker never published progress");
1266 assert!(
1267 (0.2..=0.3).contains(&progress),
1268 "expected ~0.25, got {progress}"
1269 );
1270
1271 gate.store(true, Ordering::Relaxed);
1272 with_test_gpu(|device, queue| {
1273 drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
1274 });
1275 assert!(matches!(runner.status(id), UploadStatus::Ready));
1276 }
1277
1278 #[test]
1279 fn worker_error_surfaces_as_failed() {
1280 let mut runner = JobRunner::new();
1281 let id = runner.submit_cpu(|_| {
1282 Err(ViewportError::InvalidGaussianSplatData {
1283 reason: "test error",
1284 })
1285 });
1286
1287 with_test_gpu(|device, queue| {
1288 drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
1289 });
1290
1291 match runner.status(id) {
1292 UploadStatus::Failed(ViewportError::InvalidGaussianSplatData { reason }) => {
1293 assert_eq!(reason, "test error");
1294 }
1295 other => panic!("expected Failed, got {other:?}"),
1296 }
1297 }
1298
1299 #[test]
1300 fn worker_panic_surfaces_as_failed() {
1301 let mut runner = JobRunner::new();
1302 let id = runner.submit_cpu(|_| panic!("worker exploded"));
1303
1304 with_test_gpu(|device, queue| {
1305 drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
1306 });
1307
1308 match runner.status(id) {
1309 UploadStatus::Failed(ViewportError::JobWorkerLost { reason }) => {
1310 assert_eq!(reason, "worker panicked");
1311 }
1312 other => panic!("expected Failed(JobWorkerLost), got {other:?}"),
1313 }
1314 }
1315
1316 #[test]
1317 fn callback_fires_on_completion() {
1318 let mut runner = JobRunner::new();
1319 let seen = Arc::new(Mutex::new(None));
1320 let seen_clone = seen.clone();
1321
1322 let id = runner.submit_cpu(|_| Ok(JobProduct::empty()));
1323 runner.on_complete(id, move |status| {
1324 *seen_clone.lock().unwrap() = Some(matches!(status, UploadStatus::Ready));
1325 });
1326
1327 with_test_gpu(|device, queue| {
1328 // process() returns Completion entries so the caller can run
1329 // apply + callback after dropping any external lock. The
1330 // integration on DeviceResources does this automatically;
1331 // the test drives it by hand.
1332 for _ in 0..200 {
1333 for c in runner.process(device, queue) {
1334 if let Some(cb) = c.callback {
1335 cb(&c.status);
1336 }
1337 }
1338 if matches!(runner.status(id), UploadStatus::Ready) {
1339 break;
1340 }
1341 std::thread::sleep(Duration::from_millis(5));
1342 }
1343 });
1344
1345 let observed = seen.lock().unwrap().clone();
1346 assert_eq!(observed, Some(true));
1347 }
1348
1349 #[test]
1350 fn unknown_id_returns_unknown() {
1351 let runner = JobRunner::new();
1352 let made_up = JobId(99_999);
1353 assert!(matches!(runner.status(made_up), UploadStatus::Unknown));
1354 }
1355
1356 #[test]
1357 fn many_concurrent_jobs_all_complete() {
1358 let mut runner = JobRunner::new();
1359 let mut ids = Vec::with_capacity(256);
1360 for _ in 0..256 {
1361 ids.push(runner.submit_cpu(|p| {
1362 p.set(0.5);
1363 std::thread::sleep(Duration::from_millis(2));
1364 Ok(JobProduct::empty())
1365 }));
1366 }
1367 assert_eq!(runner.pending(), 256);
1368
1369 // Observe each id transitioning to a terminal state. The retention
1370 // window only spans one drain cycle, so we cannot query every id
1371 // after the loop -- we have to collect the observation as we go.
1372 let mut seen_ready = std::collections::HashSet::new();
1373 with_test_gpu(|device, queue| {
1374 for _ in 0..400 {
1375 let _ = runner.process(device, queue);
1376 for id in &ids {
1377 if seen_ready.contains(id) {
1378 continue;
1379 }
1380 if let UploadStatus::Ready = runner.status(*id) {
1381 seen_ready.insert(*id);
1382 }
1383 }
1384 if seen_ready.len() == ids.len() {
1385 break;
1386 }
1387 std::thread::sleep(Duration::from_millis(5));
1388 }
1389 });
1390
1391 assert_eq!(
1392 seen_ready.len(),
1393 ids.len(),
1394 "stragglers: {}",
1395 runner.pending()
1396 );
1397 assert!(runner.all_complete());
1398 }
1399
1400 #[test]
1401 fn gpu_job_waits_for_submission() {
1402 let mut runner = JobRunner::new();
1403 with_test_gpu(|device, queue| {
1404 let id = runner.submit_with_gpu(device, queue, |device, queue, _p| {
1405 let buf = device.create_buffer(&wgpu::BufferDescriptor {
1406 label: Some("upload_jobs_test_buf"),
1407 size: 16,
1408 usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC,
1409 mapped_at_creation: false,
1410 });
1411 queue.write_buffer(
1412 &buf,
1413 0,
1414 &[1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
1415 );
1416 let mut enc =
1417 device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
1418 // Force a non-trivial command buffer so the submission has
1419 // something to flush.
1420 let dst = device.create_buffer(&wgpu::BufferDescriptor {
1421 label: Some("upload_jobs_test_dst"),
1422 size: 16,
1423 usage: wgpu::BufferUsages::COPY_DST,
1424 mapped_at_creation: false,
1425 });
1426 enc.copy_buffer_to_buffer(&buf, 0, &dst, 0, 16);
1427 let sub = queue.submit(std::iter::once(enc.finish()));
1428 Ok(JobProduct::with_gpu(sub))
1429 });
1430
1431 // Poll until the GPU submission completes. Generous budget so
1432 // the test stays robust under load when other tests are also
1433 // hammering the rayon pool and the GPU.
1434 for _ in 0..400 {
1435 std::thread::sleep(Duration::from_millis(5));
1436 let _ = runner.process(device, queue);
1437 if matches!(runner.status(id), UploadStatus::Ready) {
1438 return;
1439 }
1440 }
1441 panic!("GPU-gated job did not reach Ready");
1442 });
1443 }
1444
1445 // -----------------------------------------------------------------
1446 // Plugin-facing facade
1447 // -----------------------------------------------------------------
1448
1449 fn make_resources_for_jobs()
1450 -> Option<(wgpu::Device, wgpu::Queue, super::super::DeviceResources)> {
1451 let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
1452 let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
1453 power_preference: wgpu::PowerPreference::LowPower,
1454 compatible_surface: None,
1455 force_fallback_adapter: false,
1456 }))
1457 .ok()?;
1458 let (device, queue) =
1459 pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()?;
1460 let resources =
1461 super::super::DeviceResources::new(&device, wgpu::TextureFormat::Rgba8UnormSrgb, 1);
1462 Some((device, queue, resources))
1463 }
1464
1465 fn drive_resources<F: FnMut(&super::super::DeviceResources) -> bool>(
1466 resources: &mut super::super::DeviceResources,
1467 device: &wgpu::Device,
1468 queue: &wgpu::Queue,
1469 mut predicate: F,
1470 ) {
1471 for _ in 0..200 {
1472 resources.process_uploads(device, queue);
1473 if predicate(resources) {
1474 return;
1475 }
1476 std::thread::sleep(Duration::from_millis(5));
1477 }
1478 }
1479
1480 #[test]
1481 fn plugin_jobs_round_trip_typed_result() {
1482 let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
1483 eprintln!("skipping: no wgpu adapter available");
1484 return;
1485 };
1486 let id = {
1487 let jobs = super::Jobs::new(&resources);
1488 jobs.submit_cpu(|| 41_u32 + 1)
1489 };
1490
1491 // Before the worker completes, take is None.
1492 assert!(super::Jobs::new(&resources).take::<u32>(id).is_none());
1493
1494 drive_resources(&mut resources, &device, &queue, |r| {
1495 matches!(r.upload_status(id), UploadStatus::Ready)
1496 });
1497
1498 let jobs = super::Jobs::new(&resources);
1499 assert_eq!(jobs.take::<u32>(id), Some(42));
1500 // Second take returns None (already drained).
1501 assert_eq!(jobs.take::<u32>(id), None);
1502 }
1503
1504 #[test]
1505 fn plugin_jobs_panic_surfaces_as_failed() {
1506 let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
1507 eprintln!("skipping: no wgpu adapter available");
1508 return;
1509 };
1510 let id = {
1511 let jobs = super::Jobs::new(&resources);
1512 jobs.submit_cpu(|| -> u32 { panic!("plugin worker exploded") })
1513 };
1514
1515 drive_resources(&mut resources, &device, &queue, |r| {
1516 !matches!(r.upload_status(id), UploadStatus::Pending { .. })
1517 });
1518
1519 match resources.upload_status(id) {
1520 UploadStatus::Failed(ViewportError::JobWorkerLost { reason }) => {
1521 assert_eq!(reason, "worker panicked");
1522 }
1523 other => panic!("expected Failed, got {other:?}"),
1524 }
1525 // Failed jobs leave the slot intact but with no value; take returns None.
1526 assert!(super::Jobs::new(&resources).take::<u32>(id).is_none());
1527 }
1528
1529 #[test]
1530 fn plugin_jobs_wrong_type_returns_none() {
1531 let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
1532 eprintln!("skipping: no wgpu adapter available");
1533 return;
1534 };
1535 let id = {
1536 let jobs = super::Jobs::new(&resources);
1537 jobs.submit_cpu(|| 7_i64)
1538 };
1539
1540 drive_resources(&mut resources, &device, &queue, |r| {
1541 matches!(r.upload_status(id), UploadStatus::Ready)
1542 });
1543
1544 let jobs = super::Jobs::new(&resources);
1545 // Asking for a different type drops the box silently.
1546 assert!(jobs.take::<u32>(id).is_none());
1547 // After the type mismatch the slot is gone, so subsequent takes
1548 // continue to return None even with the correct type.
1549 assert!(jobs.take::<i64>(id).is_none());
1550 }
1551
1552 // -----------------------------------------------------------------
1553 // Optional `future` feature
1554 // -----------------------------------------------------------------
1555
1556 /// Smoke test the `JobHandle` future under a non-tokio executor.
1557 ///
1558 /// Drives the upload runner from a helper thread so the polled future
1559 /// can observe completion through its registered callback. Uses
1560 /// `pollster::block_on` because the crate stays runtime-free; the same
1561 /// future works unchanged under tokio.
1562 #[cfg(feature = "future")]
1563 #[test]
1564 fn job_handle_resolves_via_future() {
1565 let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
1566 eprintln!("skipping: no wgpu adapter available");
1567 return;
1568 };
1569
1570 // Submit a CPU job so we have an id to await.
1571 let id = {
1572 let jobs = super::Jobs::new(&resources);
1573 jobs.submit_cpu(|| 7_u32)
1574 };
1575 let handle = resources.upload_handle(id);
1576
1577 // Drive `process_uploads` from a worker thread while the main
1578 // thread blocks on the future. The two threads share the resources
1579 // through a `Mutex` so the worker can call `process_uploads(&mut)`.
1580 let shared = Arc::new(std::sync::Mutex::new((resources, device, queue)));
1581 let driver_handle = {
1582 let shared = shared.clone();
1583 std::thread::spawn(move || {
1584 for _ in 0..400 {
1585 {
1586 let mut g = shared.lock().unwrap();
1587 let (resources, device, queue) = &mut *g;
1588 resources.process_uploads(device, queue);
1589 if resources.all_uploads_complete() {
1590 return;
1591 }
1592 }
1593 std::thread::sleep(Duration::from_millis(5));
1594 }
1595 })
1596 };
1597
1598 let result = pollster::block_on(handle);
1599 assert!(matches!(result, Ok(())));
1600 driver_handle.join().ok();
1601 }
1602
1603 // -----------------------------------------------------------------
1604 // Frame budget
1605 // -----------------------------------------------------------------
1606
1607 /// Submit a batch of apply-bearing jobs, run one short-budget pass,
1608 /// and check that the leftover applies survive to the next pass.
1609 ///
1610 /// Each apply spins for a few hundred microseconds, so a 100 us
1611 /// budget is virtually guaranteed to drop work to the next frame
1612 /// without being so tight that the runner spins forever. Across
1613 /// successive unbounded passes everything drains and every job ends
1614 /// in Ready.
1615 #[test]
1616 fn budgeted_drain_spills_to_next_call() {
1617 let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
1618 eprintln!("skipping: no wgpu adapter available");
1619 return;
1620 };
1621
1622 let mut ids = Vec::with_capacity(32);
1623 {
1624 let jobs = super::Jobs::new(&resources);
1625 for i in 0..32_u32 {
1626 // The closure body runs on the worker; the value
1627 // arrives via the plugin facade's typed slot, which
1628 // materializes through an apply step on the main
1629 // thread. That apply step is what the budget caps.
1630 ids.push(jobs.submit_cpu(move || i * 2));
1631 }
1632 }
1633
1634 // Drain workers first so every job is sitting on pending_apply.
1635 for _ in 0..400 {
1636 {
1637 let runner = resources.jobs.lock().unwrap();
1638 if runner.pending_apply_len() == ids.len() {
1639 break;
1640 }
1641 }
1642 std::thread::sleep(Duration::from_millis(2));
1643 // Advance the runner without running applies. A zero-duration
1644 // budget pops nothing (the exhausted check fires before the
1645 // first pop), so pending_apply grows as workers report in.
1646 resources.process_uploads_with_budget(
1647 &device,
1648 &queue,
1649 super::FrameBudget::from_now(Duration::from_nanos(0)),
1650 );
1651 }
1652 let after_workers = resources.jobs.lock().unwrap().pending_apply_len();
1653 assert!(
1654 after_workers > 0,
1655 "expected at least one apply queued, got 0"
1656 );
1657
1658 // One unbounded pass clears the rest and lands every job in
1659 // Ready.
1660 resources.process_uploads(&device, &queue);
1661 for &id in &ids {
1662 assert!(matches!(resources.upload_status(id), UploadStatus::Ready));
1663 }
1664 assert!(resources.all_uploads_complete());
1665 }
1666
1667 /// Creates a headless wgpu device + queue for the duration of `f`.
1668 ///
1669 /// Skips the test (via early return) if no adapter is available. CI
1670 /// builds without a GPU should pass the CPU-only tests above and skip
1671 /// the GPU-gated one.
1672 fn with_test_gpu<F: FnOnce(&wgpu::Device, &wgpu::Queue)>(f: F) {
1673 let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
1674 backends: wgpu::Backends::PRIMARY | wgpu::Backends::SECONDARY,
1675 ..Default::default()
1676 });
1677 let adapter =
1678 match pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
1679 power_preference: wgpu::PowerPreference::LowPower,
1680 compatible_surface: None,
1681 force_fallback_adapter: false,
1682 })) {
1683 Ok(a) => a,
1684 Err(_) => {
1685 eprintln!("skipping GPU-gated test: no adapter available");
1686 return;
1687 }
1688 };
1689 let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
1690 label: Some("upload_jobs_test_device"),
1691 required_features: wgpu::Features::empty(),
1692 required_limits: wgpu::Limits::downlevel_defaults(),
1693 memory_hints: wgpu::MemoryHints::Performance,
1694 experimental_features: wgpu::ExperimentalFeatures::default(),
1695 trace: wgpu::Trace::Off,
1696 }))
1697 .expect("device creation");
1698 f(&device, &queue);
1699 }
1700}