Source code
Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
#[cfg(not(feature = "shuttle"))]
use std::sync::Mutex;
#[cfg(not(feature = "shuttle"))]
use std::sync::atomic::{AtomicUsize, Ordering};
#[cfg(not(feature = "shuttle"))]
use std::thread;
#[cfg(feature = "shuttle")]
use shuttle::sync::Mutex;
#[cfg(feature = "shuttle")]
use shuttle::sync::atomic::{AtomicUsize, Ordering};
#[cfg(feature = "shuttle")]
use shuttle::thread;
use crate::api::JxlParallelRunner;
use crate::error::Error;
pub struct TestParallelRunner {
pub max_threads: usize,
}
impl JxlParallelRunner for TestParallelRunner {
fn run(&mut self, num: usize, fun: &crate::api::JxlParallelRunnerFun<'_>) -> Result<(), Error> {
if num <= 1 || self.max_threads <= 1 {
for i in 0..num {
fun(i)?;
}
return Ok(());
}
let num_threads = self.max_threads.min(num);
let next_task = AtomicUsize::new(0);
let error = Mutex::new(None);
thread::scope(|s| {
let mut handles = Vec::with_capacity(num_threads);
for _ in 0..num_threads {
handles.push(s.spawn(|| {
loop {
if error.lock().unwrap().is_some() {
break;
}
let task = next_task.fetch_add(1, Ordering::Relaxed);
if task >= num {
break;
}
if let Err(e) = fun(task) {
let mut err = error.lock().unwrap();
if err.is_none() {
*err = Some(e);
}
break;
}
}
}));
}
for handle in handles {
if let Err(e) = handle.join() {
std::panic::resume_unwind(e);
}
}
});
if let Some(err) = error.into_inner().unwrap() {
Err(err)
} else {
Ok(())
}
}
fn num_threads(&self) -> usize {
self.max_threads
}
}