use std::any::Any;
use std::collections::{HashMap, VecDeque};
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};
use std::sync::mpsc;
use std::time::{Duration, Instant};
use crate::error::ViewportError;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct JobId(u64);
#[derive(Debug, Clone)]
pub enum UploadStatus {
Pending {
progress: f32,
},
Ready,
Failed(ViewportError),
Unknown,
}
#[derive(Clone)]
pub struct ProgressHandle {
inner: Arc<AtomicU32>,
}
impl ProgressHandle {
fn new() -> Self {
Self {
inner: Arc::new(AtomicU32::new(0)),
}
}
pub fn set(&self, fraction: f32) {
let clamped = fraction.clamp(0.0, 1.0);
self.inner.store(clamped.to_bits(), Ordering::Relaxed);
}
fn read(&self) -> f32 {
f32::from_bits(self.inner.load(Ordering::Relaxed))
}
}
pub type ApplyFn = Box<dyn FnOnce(&mut super::ViewportGpuResources) + Send>;
pub struct ResultSlot<T> {
inner: Arc<std::sync::Mutex<Option<T>>>,
}
impl<T> Clone for ResultSlot<T> {
fn clone(&self) -> Self {
Self {
inner: Arc::clone(&self.inner),
}
}
}
impl<T> ResultSlot<T> {
pub fn new() -> Self {
Self {
inner: Arc::new(std::sync::Mutex::new(None)),
}
}
pub fn set(&self, value: T) {
let mut guard = self.inner.lock().expect("result slot poisoned");
*guard = Some(value);
}
pub fn take(&self) -> Option<T> {
let mut guard = self.inner.lock().expect("result slot poisoned");
guard.take()
}
}
impl<T> Default for ResultSlot<T> {
fn default() -> Self {
Self::new()
}
}
pub struct JobProduct {
pub gpu: Option<wgpu::SubmissionIndex>,
pub apply: Option<ApplyFn>,
}
impl Default for JobProduct {
fn default() -> Self {
Self {
gpu: None,
apply: None,
}
}
}
impl JobProduct {
pub fn empty() -> Self {
Self::default()
}
pub fn with_gpu(gpu: wgpu::SubmissionIndex) -> Self {
Self {
gpu: Some(gpu),
apply: None,
}
}
pub fn with_apply(apply: ApplyFn) -> Self {
Self {
gpu: None,
apply: Some(apply),
}
}
pub fn with_gpu_and_apply(gpu: wgpu::SubmissionIndex, apply: ApplyFn) -> Self {
Self {
gpu: Some(gpu),
apply: Some(apply),
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct FrameBudget {
deadline: Option<Instant>,
}
impl FrameBudget {
pub fn unbounded() -> Self {
Self { deadline: None }
}
pub fn from_now(duration: Duration) -> Self {
Self {
deadline: Some(Instant::now() + duration),
}
}
fn exhausted(&self) -> bool {
match self.deadline {
Some(t) => Instant::now() >= t,
None => false,
}
}
}
#[allow(dead_code)]
enum WorkerOutcome {
Done(JobProduct, Duration),
Failed(ViewportError, Duration),
}
type CompletionCallback = Box<dyn FnOnce(&UploadStatus) + Send>;
pub struct Completion {
pub id: JobId,
pub status: UploadStatus,
pub apply: Option<ApplyFn>,
pub callback: Option<CompletionCallback>,
}
struct JobSlot {
progress: ProgressHandle,
rx: mpsc::Receiver<WorkerOutcome>,
awaiting: Option<(Option<wgpu::SubmissionIndex>, Option<ApplyFn>)>,
callback: Option<CompletionCallback>,
}
#[allow(dead_code)]
pub struct JobRunner {
next_id: u64,
slots: HashMap<u64, JobSlot>,
pending_apply: VecDeque<PendingApply>,
finished: HashMap<u64, UploadStatus>,
durations: HashMap<u64, Duration>,
}
pub struct PendingApply {
pub id: JobId,
pub status: UploadStatus,
pub apply: ApplyFn,
pub callback: Option<CompletionCallback>,
}
impl Default for JobRunner {
fn default() -> Self {
Self::new()
}
}
#[allow(dead_code)]
impl JobRunner {
pub fn new() -> Self {
Self {
next_id: 1,
slots: HashMap::new(),
pending_apply: VecDeque::new(),
finished: HashMap::new(),
durations: HashMap::new(),
}
}
pub fn duration(&self, id: JobId) -> Option<Duration> {
self.durations.get(&id.0).copied()
}
pub fn add_apply_duration(&mut self, id: JobId, apply: Duration) {
let entry = self.durations.entry(id.0).or_insert(Duration::ZERO);
*entry = entry.saturating_add(apply);
}
pub fn drop_duration(&mut self, id: JobId) {
self.durations.remove(&id.0);
}
fn issue_id(&mut self) -> JobId {
let id = self.next_id;
self.next_id = self
.next_id
.checked_add(1)
.expect("upload job id space exhausted");
JobId(id)
}
pub(crate) fn submit_cpu<F>(&mut self, work: F) -> JobId
where
F: FnOnce(&ProgressHandle) -> Result<JobProduct, ViewportError> + Send + 'static,
{
let id = self.issue_id();
let progress = ProgressHandle::new();
let worker_progress = progress.clone();
let (tx, rx) = mpsc::channel();
rayon::spawn(move || {
let t0 = Instant::now();
let outcome = match catch_unwind(AssertUnwindSafe(|| work(&worker_progress))) {
Ok(Ok(product)) => WorkerOutcome::Done(product, t0.elapsed()),
Ok(Err(e)) => WorkerOutcome::Failed(e, t0.elapsed()),
Err(_) => WorkerOutcome::Failed(
ViewportError::JobWorkerLost {
reason: "worker panicked",
},
t0.elapsed(),
),
};
let _ = tx.send(outcome);
});
self.slots.insert(
id.0,
JobSlot {
progress,
rx,
awaiting: None,
callback: None,
},
);
id
}
pub(crate) fn submit_with_gpu<F>(
&mut self,
device: &wgpu::Device,
queue: &wgpu::Queue,
work: F,
) -> JobId
where
F: FnOnce(
&wgpu::Device,
&wgpu::Queue,
&ProgressHandle,
) -> Result<JobProduct, ViewportError>
+ Send
+ 'static,
{
let id = self.issue_id();
let progress = ProgressHandle::new();
let worker_progress = progress.clone();
let (tx, rx) = mpsc::channel();
let device = device.clone();
let queue = queue.clone();
rayon::spawn(move || {
let t0 = Instant::now();
let outcome =
match catch_unwind(AssertUnwindSafe(|| work(&device, &queue, &worker_progress))) {
Ok(Ok(product)) => WorkerOutcome::Done(product, t0.elapsed()),
Ok(Err(e)) => WorkerOutcome::Failed(e, t0.elapsed()),
Err(_) => WorkerOutcome::Failed(
ViewportError::JobWorkerLost {
reason: "worker panicked",
},
t0.elapsed(),
),
};
let _ = tx.send(outcome);
});
self.slots.insert(
id.0,
JobSlot {
progress,
rx,
awaiting: None,
callback: None,
},
);
id
}
pub fn on_complete<F>(&mut self, id: JobId, cb: F)
where
F: FnOnce(&UploadStatus) + Send + 'static,
{
if let Some(slot) = self.slots.get_mut(&id.0) {
slot.callback = Some(Box::new(cb));
return;
}
if let Some(status) = self.finished.get(&id.0) {
cb(status);
}
}
pub fn status(&self, id: JobId) -> UploadStatus {
if let Some(slot) = self.slots.get(&id.0) {
return UploadStatus::Pending {
progress: slot.progress.read(),
};
}
if self.pending_apply.iter().any(|pa| pa.id.0 == id.0) {
return UploadStatus::Pending { progress: 1.0 };
}
if let Some(status) = self.finished.get(&id.0) {
return status.clone();
}
UploadStatus::Unknown
}
pub fn pending(&self) -> usize {
self.slots.len() + self.pending_apply.len()
}
pub fn all_complete(&self) -> bool {
self.slots.is_empty() && self.pending_apply.is_empty()
}
pub fn pop_pending_apply(&mut self) -> Option<PendingApply> {
self.pending_apply.pop_front()
}
pub fn requeue_pending_apply(&mut self, pa: PendingApply) {
self.pending_apply.push_front(pa);
}
pub fn mark_applied(&mut self, id: JobId, status: UploadStatus) {
self.finished.insert(id.0, status);
}
pub fn pending_apply_len(&self) -> usize {
self.pending_apply.len()
}
pub fn process(&mut self, device: &wgpu::Device, _queue: &wgpu::Queue) -> Vec<Completion> {
self.finished.clear();
let _ = device.poll(wgpu::PollType::Poll);
let mut completions = Vec::new();
let ids: Vec<u64> = self.slots.keys().copied().collect();
for id in ids {
if self.slots.get(&id).is_some_and(|s| s.awaiting.is_none()) {
let outcome = self
.slots
.get(&id)
.map(|s| s.rx.try_recv())
.expect("slot existed");
match outcome {
Ok(WorkerOutcome::Done(product, worker_dur)) => {
self.durations.insert(id, worker_dur);
let JobProduct { gpu, apply } = product;
match gpu {
None => {
self.finish(id, UploadStatus::Ready, apply, &mut completions);
continue;
}
Some(sub) => {
if let Some(slot) = self.slots.get_mut(&id) {
slot.awaiting = Some((Some(sub), apply));
}
}
}
}
Ok(WorkerOutcome::Failed(e, worker_dur)) => {
self.durations.insert(id, worker_dur);
self.finish(id, UploadStatus::Failed(e), None, &mut completions);
continue;
}
Err(mpsc::TryRecvError::Empty) => continue,
Err(mpsc::TryRecvError::Disconnected) => {
self.finish(
id,
UploadStatus::Failed(ViewportError::JobWorkerLost {
reason: "worker channel closed without result",
}),
None,
&mut completions,
);
continue;
}
}
}
let pending_sub = self
.slots
.get(&id)
.and_then(|s| s.awaiting.as_ref())
.and_then(|(g, _)| g.clone());
if let Some(sub) = pending_sub {
let result = device.poll(wgpu::PollType::Wait {
submission_index: Some(sub),
timeout: Some(Duration::from_millis(0)),
});
match result {
Ok(wgpu::PollStatus::QueueEmpty) | Ok(wgpu::PollStatus::WaitSucceeded) => {
let apply = self
.slots
.get_mut(&id)
.and_then(|s| s.awaiting.take())
.and_then(|(_, a)| a);
self.finish(id, UploadStatus::Ready, apply, &mut completions);
}
Ok(wgpu::PollStatus::Poll) => {
}
Err(_) => {
}
}
}
}
completions
}
fn finish(
&mut self,
id: u64,
status: UploadStatus,
apply: Option<ApplyFn>,
completions: &mut Vec<Completion>,
) {
let Some(mut slot) = self.slots.remove(&id) else {
return;
};
let callback = slot.callback.take();
match (status, apply) {
(UploadStatus::Ready, Some(apply_fn)) => {
self.pending_apply.push_back(PendingApply {
id: JobId(id),
status: UploadStatus::Ready,
apply: apply_fn,
callback,
});
return;
}
(status, _) => {
completions.push(Completion {
id: JobId(id),
status: status.clone(),
apply: None,
callback,
});
self.finished.insert(id, status);
return;
}
}
}
}
impl super::ViewportGpuResources {
pub fn process_uploads(&mut self, device: &wgpu::Device, queue: &wgpu::Queue) {
self.process_uploads_with_budget(device, queue, FrameBudget::unbounded());
}
pub fn process_uploads_with_budget(
&mut self,
device: &wgpu::Device,
queue: &wgpu::Queue,
budget: FrameBudget,
) {
let completions = {
let mut runner = self.jobs.lock().expect("upload job runner poisoned");
runner.process(device, queue)
};
for Completion {
id: _,
status,
apply: _,
callback,
} in completions
{
if let Some(cb) = callback {
cb(&status);
}
}
loop {
if budget.exhausted() {
break;
}
let next = {
let mut runner = self.jobs.lock().expect("upload job runner poisoned");
runner.pop_pending_apply()
};
let Some(PendingApply {
id,
status,
apply,
callback,
}) = next
else {
break;
};
let t = Instant::now();
apply(self);
let apply_d = t.elapsed();
{
let mut runner = self.jobs.lock().expect("upload job runner poisoned");
runner.add_apply_duration(id, apply_d);
runner.mark_applied(id, status.clone());
}
if let Some(cb) = callback {
cb(&status);
}
}
}
pub fn job_duration(&self, id: JobId) -> Option<Duration> {
let runner = self.jobs.lock().expect("upload job runner poisoned");
runner.duration(id)
}
pub fn drop_job_duration(&mut self, id: JobId) {
let mut runner = self.jobs.lock().expect("upload job runner poisoned");
runner.drop_duration(id);
}
pub fn upload_status(&self, id: JobId) -> UploadStatus {
let runner = self.jobs.lock().expect("upload job runner poisoned");
runner.status(id)
}
pub fn uploads_pending(&self) -> usize {
let runner = self.jobs.lock().expect("upload job runner poisoned");
runner.pending()
}
pub fn all_uploads_complete(&self) -> bool {
let runner = self.jobs.lock().expect("upload job runner poisoned");
runner.all_complete()
}
pub fn on_upload_complete<F>(&mut self, id: JobId, cb: F)
where
F: FnOnce(&UploadStatus) + Send + 'static,
{
let mut runner = self.jobs.lock().expect("upload job runner poisoned");
runner.on_complete(id, cb);
}
pub fn drain_until(
&mut self,
device: &wgpu::Device,
queue: &wgpu::Queue,
id: JobId,
) -> crate::error::ViewportResult<()> {
loop {
self.process_uploads(device, queue);
match self.upload_status(id) {
UploadStatus::Ready => return Ok(()),
UploadStatus::Failed(e) => return Err(e),
UploadStatus::Pending { .. } => {
std::thread::sleep(Duration::from_micros(200));
}
UploadStatus::Unknown => {
return Err(crate::error::ViewportError::JobResultMissing {
reason: "drain target was already reaped",
});
}
}
}
}
}
#[cfg(feature = "future")]
pub struct JobHandle {
id: JobId,
state: Arc<std::sync::Mutex<JobHandleState>>,
}
#[cfg(feature = "future")]
struct JobHandleState {
done: Option<crate::error::ViewportResult<()>>,
waker: Option<std::task::Waker>,
}
#[cfg(feature = "future")]
impl JobHandle {
pub fn id(&self) -> JobId {
self.id
}
}
#[cfg(feature = "future")]
impl std::future::Future for JobHandle {
type Output = crate::error::ViewportResult<()>;
fn poll(
self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
let mut guard = self.state.lock().expect("job handle poisoned");
if let Some(result) = guard.done.take() {
return std::task::Poll::Ready(result);
}
guard.waker = Some(cx.waker().clone());
std::task::Poll::Pending
}
}
#[cfg(feature = "future")]
impl super::ViewportGpuResources {
pub fn upload_handle(&mut self, id: JobId) -> JobHandle {
let state = Arc::new(std::sync::Mutex::new(JobHandleState {
done: None,
waker: None,
}));
let state_for_cb = state.clone();
self.on_upload_complete(id, move |status| {
let result = match status {
UploadStatus::Ready => Ok(()),
UploadStatus::Failed(e) => Err(e.clone()),
UploadStatus::Pending { .. } => {
Err(crate::error::ViewportError::JobNotReady)
}
UploadStatus::Unknown => Err(crate::error::ViewportError::JobResultMissing {
reason: "job vanished before completion callback fired",
}),
};
let mut guard = state_for_cb.lock().expect("job handle poisoned");
guard.done = Some(result);
if let Some(waker) = guard.waker.take() {
waker.wake();
}
});
JobHandle { id, state }
}
}
type PluginResultSlot = ResultSlot<Box<dyn Any + Send>>;
pub struct Jobs<'a> {
resources: &'a super::ViewportGpuResources,
}
impl<'a> Jobs<'a> {
pub(crate) fn new(resources: &'a super::ViewportGpuResources) -> Self {
Self { resources }
}
pub fn submit_cpu<T, F>(&self, work: F) -> JobId
where
T: Send + 'static,
F: FnOnce() -> T + Send + 'static,
{
let slot: PluginResultSlot = ResultSlot::new();
let slot_for_apply = slot.clone();
let id = {
let mut runner = self
.resources
.jobs
.lock()
.expect("upload job runner poisoned");
runner.submit_cpu(move |_progress| {
let value: T = work();
let boxed: Box<dyn Any + Send> = Box::new(value);
Ok(JobProduct::with_apply(Box::new(
move |_resources: &mut super::ViewportGpuResources| {
slot_for_apply.set(boxed);
},
)))
})
};
self.resources
.plugin_job_results
.lock()
.expect("plugin job result map poisoned")
.insert(id, slot);
id
}
pub fn status(&self, id: JobId) -> UploadStatus {
self.resources.upload_status(id)
}
pub fn take<T: Any + Send + 'static>(&self, id: JobId) -> Option<T> {
let mut map = self
.resources
.plugin_job_results
.lock()
.expect("plugin job result map poisoned");
let slot = map.get(&id)?.clone();
let boxed = slot.take()?;
match boxed.downcast::<T>() {
Ok(value) => {
map.remove(&id);
Some(*value)
}
Err(_boxed_back) => {
map.remove(&id);
None
}
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
use super::*;
fn drain_until<F>(
runner: &mut JobRunner,
device: &wgpu::Device,
queue: &wgpu::Queue,
max_iterations: usize,
mut predicate: F,
) where
F: FnMut(&JobRunner) -> bool,
{
for _ in 0..max_iterations {
let _ = runner.process(device, queue);
if predicate(runner) {
return;
}
std::thread::sleep(Duration::from_millis(5));
}
}
#[test]
fn cpu_job_reports_ready_after_drain() {
let mut runner = JobRunner::new();
let id = runner.submit_cpu(|_p| Ok(JobProduct::empty()));
assert_eq!(runner.pending(), 1);
with_test_gpu(|device, queue| {
drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
});
assert!(matches!(runner.status(id), UploadStatus::Ready));
assert_eq!(runner.pending(), 0);
assert!(runner.all_complete());
}
#[test]
fn cpu_job_progress_is_observable() {
let mut runner = JobRunner::new();
let gate = Arc::new(AtomicBool::new(false));
let gate_for_worker = gate.clone();
let id = runner.submit_cpu(move |p| {
p.set(0.25);
while !gate_for_worker.load(Ordering::Relaxed) {
std::thread::sleep(Duration::from_millis(1));
}
p.set(1.0);
Ok(JobProduct::empty())
});
let mut observed = None;
for _ in 0..200 {
if let UploadStatus::Pending { progress } = runner.status(id) {
if progress > 0.0 {
observed = Some(progress);
break;
}
}
std::thread::sleep(Duration::from_millis(5));
}
let progress = observed.expect("worker never published progress");
assert!(
(0.2..=0.3).contains(&progress),
"expected ~0.25, got {progress}"
);
gate.store(true, Ordering::Relaxed);
with_test_gpu(|device, queue| {
drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
});
assert!(matches!(runner.status(id), UploadStatus::Ready));
}
#[test]
fn worker_error_surfaces_as_failed() {
let mut runner = JobRunner::new();
let id = runner.submit_cpu(|_| {
Err(ViewportError::InvalidGaussianSplatData {
reason: "test error",
})
});
with_test_gpu(|device, queue| {
drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
});
match runner.status(id) {
UploadStatus::Failed(ViewportError::InvalidGaussianSplatData { reason }) => {
assert_eq!(reason, "test error");
}
other => panic!("expected Failed, got {other:?}"),
}
}
#[test]
fn worker_panic_surfaces_as_failed() {
let mut runner = JobRunner::new();
let id = runner.submit_cpu(|_| panic!("worker exploded"));
with_test_gpu(|device, queue| {
drain_until(&mut runner, device, queue, 200, |r| r.all_complete());
});
match runner.status(id) {
UploadStatus::Failed(ViewportError::JobWorkerLost { reason }) => {
assert_eq!(reason, "worker panicked");
}
other => panic!("expected Failed(JobWorkerLost), got {other:?}"),
}
}
#[test]
fn callback_fires_on_completion() {
let mut runner = JobRunner::new();
let seen = Arc::new(Mutex::new(None));
let seen_clone = seen.clone();
let id = runner.submit_cpu(|_| Ok(JobProduct::empty()));
runner.on_complete(id, move |status| {
*seen_clone.lock().unwrap() = Some(matches!(status, UploadStatus::Ready));
});
with_test_gpu(|device, queue| {
for _ in 0..200 {
for c in runner.process(device, queue) {
if let Some(cb) = c.callback {
cb(&c.status);
}
}
if matches!(runner.status(id), UploadStatus::Ready) {
break;
}
std::thread::sleep(Duration::from_millis(5));
}
});
let observed = seen.lock().unwrap().clone();
assert_eq!(observed, Some(true));
}
#[test]
fn unknown_id_returns_unknown() {
let runner = JobRunner::new();
let made_up = JobId(99_999);
assert!(matches!(runner.status(made_up), UploadStatus::Unknown));
}
#[test]
fn many_concurrent_jobs_all_complete() {
let mut runner = JobRunner::new();
let mut ids = Vec::with_capacity(256);
for _ in 0..256 {
ids.push(runner.submit_cpu(|p| {
p.set(0.5);
std::thread::sleep(Duration::from_millis(2));
Ok(JobProduct::empty())
}));
}
assert_eq!(runner.pending(), 256);
let mut seen_ready = std::collections::HashSet::new();
with_test_gpu(|device, queue| {
for _ in 0..400 {
let _ = runner.process(device, queue);
for id in &ids {
if seen_ready.contains(id) {
continue;
}
if let UploadStatus::Ready = runner.status(*id) {
seen_ready.insert(*id);
}
}
if seen_ready.len() == ids.len() {
break;
}
std::thread::sleep(Duration::from_millis(5));
}
});
assert_eq!(
seen_ready.len(),
ids.len(),
"stragglers: {}",
runner.pending()
);
assert!(runner.all_complete());
}
#[test]
fn gpu_job_waits_for_submission() {
let mut runner = JobRunner::new();
with_test_gpu(|device, queue| {
let id = runner.submit_with_gpu(device, queue, |device, queue, _p| {
let buf = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("upload_jobs_test_buf"),
size: 16,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
queue.write_buffer(
&buf,
0,
&[1u8, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
);
let mut enc =
device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None });
let dst = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("upload_jobs_test_dst"),
size: 16,
usage: wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
enc.copy_buffer_to_buffer(&buf, 0, &dst, 0, 16);
let sub = queue.submit(std::iter::once(enc.finish()));
Ok(JobProduct::with_gpu(sub))
});
for _ in 0..400 {
std::thread::sleep(Duration::from_millis(5));
let _ = runner.process(device, queue);
if matches!(runner.status(id), UploadStatus::Ready) {
return;
}
}
panic!("GPU-gated job did not reach Ready");
});
}
fn make_resources_for_jobs() -> Option<(
wgpu::Device,
wgpu::Queue,
super::super::ViewportGpuResources,
)> {
let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::LowPower,
compatible_surface: None,
force_fallback_adapter: false,
}))
.ok()?;
let (device, queue) =
pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()?;
let resources = super::super::ViewportGpuResources::new(
&device,
wgpu::TextureFormat::Rgba8UnormSrgb,
1,
);
Some((device, queue, resources))
}
fn drive_resources<F: FnMut(&super::super::ViewportGpuResources) -> bool>(
resources: &mut super::super::ViewportGpuResources,
device: &wgpu::Device,
queue: &wgpu::Queue,
mut predicate: F,
) {
for _ in 0..200 {
resources.process_uploads(device, queue);
if predicate(resources) {
return;
}
std::thread::sleep(Duration::from_millis(5));
}
}
#[test]
fn plugin_jobs_round_trip_typed_result() {
let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
eprintln!("skipping: no wgpu adapter available");
return;
};
let id = {
let jobs = super::Jobs::new(&resources);
jobs.submit_cpu(|| 41_u32 + 1)
};
assert!(super::Jobs::new(&resources).take::<u32>(id).is_none());
drive_resources(&mut resources, &device, &queue, |r| {
matches!(r.upload_status(id), UploadStatus::Ready)
});
let jobs = super::Jobs::new(&resources);
assert_eq!(jobs.take::<u32>(id), Some(42));
assert_eq!(jobs.take::<u32>(id), None);
}
#[test]
fn plugin_jobs_panic_surfaces_as_failed() {
let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
eprintln!("skipping: no wgpu adapter available");
return;
};
let id = {
let jobs = super::Jobs::new(&resources);
jobs.submit_cpu(|| -> u32 { panic!("plugin worker exploded") })
};
drive_resources(&mut resources, &device, &queue, |r| {
!matches!(r.upload_status(id), UploadStatus::Pending { .. })
});
match resources.upload_status(id) {
UploadStatus::Failed(ViewportError::JobWorkerLost { reason }) => {
assert_eq!(reason, "worker panicked");
}
other => panic!("expected Failed, got {other:?}"),
}
assert!(super::Jobs::new(&resources).take::<u32>(id).is_none());
}
#[test]
fn plugin_jobs_wrong_type_returns_none() {
let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
eprintln!("skipping: no wgpu adapter available");
return;
};
let id = {
let jobs = super::Jobs::new(&resources);
jobs.submit_cpu(|| 7_i64)
};
drive_resources(&mut resources, &device, &queue, |r| {
matches!(r.upload_status(id), UploadStatus::Ready)
});
let jobs = super::Jobs::new(&resources);
assert!(jobs.take::<u32>(id).is_none());
assert!(jobs.take::<i64>(id).is_none());
}
#[cfg(feature = "future")]
#[test]
fn job_handle_resolves_via_future() {
let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
eprintln!("skipping: no wgpu adapter available");
return;
};
let id = {
let jobs = super::Jobs::new(&resources);
jobs.submit_cpu(|| 7_u32)
};
let handle = resources.upload_handle(id);
let shared = Arc::new(std::sync::Mutex::new((resources, device, queue)));
let driver_handle = {
let shared = shared.clone();
std::thread::spawn(move || {
for _ in 0..400 {
{
let mut g = shared.lock().unwrap();
let (resources, device, queue) = &mut *g;
resources.process_uploads(device, queue);
if resources.all_uploads_complete() {
return;
}
}
std::thread::sleep(Duration::from_millis(5));
}
})
};
let result = pollster::block_on(handle);
assert!(matches!(result, Ok(())));
driver_handle.join().ok();
}
#[test]
fn budgeted_drain_spills_to_next_call() {
let Some((device, queue, mut resources)) = make_resources_for_jobs() else {
eprintln!("skipping: no wgpu adapter available");
return;
};
let mut ids = Vec::with_capacity(32);
{
let jobs = super::Jobs::new(&resources);
for i in 0..32_u32 {
ids.push(jobs.submit_cpu(move || i * 2));
}
}
for _ in 0..400 {
{
let runner = resources.jobs.lock().unwrap();
if runner.pending_apply_len() == ids.len() {
break;
}
}
std::thread::sleep(Duration::from_millis(2));
resources.process_uploads_with_budget(
&device,
&queue,
super::FrameBudget::from_now(Duration::from_nanos(0)),
);
}
let after_workers = resources.jobs.lock().unwrap().pending_apply_len();
assert!(
after_workers > 0,
"expected at least one apply queued, got 0"
);
resources.process_uploads(&device, &queue);
for &id in &ids {
assert!(matches!(resources.upload_status(id), UploadStatus::Ready));
}
assert!(resources.all_uploads_complete());
}
fn with_test_gpu<F: FnOnce(&wgpu::Device, &wgpu::Queue)>(f: F) {
let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor {
backends: wgpu::Backends::PRIMARY | wgpu::Backends::SECONDARY,
..Default::default()
});
let adapter =
match pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
power_preference: wgpu::PowerPreference::LowPower,
compatible_surface: None,
force_fallback_adapter: false,
})) {
Ok(a) => a,
Err(_) => {
eprintln!("skipping GPU-gated test: no adapter available");
return;
}
};
let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
label: Some("upload_jobs_test_device"),
required_features: wgpu::Features::empty(),
required_limits: wgpu::Limits::downlevel_defaults(),
memory_hints: wgpu::MemoryHints::Performance,
experimental_features: wgpu::ExperimentalFeatures::default(),
trace: wgpu::Trace::Off,
}))
.expect("device creation");
f(&device, &queue);
}
}