use std::sync::atomic::{AtomicBool, Ordering};
use crate::parallel::{num_threads, Reduce};
pub fn join<O1: Send, O2: Send>(left: impl FnOnce() -> O1 + Send, right: impl FnOnce() -> O2 + Send) -> (O1, O2) {
crossbeam_utils::thread::scope(|s| {
let left = s.spawn(|_| left());
let right = s.spawn(|_| right());
(left.join().unwrap(), right.join().unwrap())
})
.unwrap()
}
pub fn threads<'env, F, R>(f: F) -> std::thread::Result<R>
where
F: FnOnce(&crossbeam_utils::thread::Scope<'env>) -> R,
{
crossbeam_utils::thread::scope(f)
}
pub fn in_parallel<I, S, O, R>(
input: impl Iterator<Item = I> + Send,
thread_limit: Option<usize>,
new_thread_state: impl Fn(usize) -> S + Send + Clone,
consume: impl Fn(I, &mut S) -> O + Send + Clone,
mut reducer: R,
) -> Result<<R as Reduce>::Output, <R as Reduce>::Error>
where
R: Reduce<Input = O>,
I: Send,
O: Send,
{
let num_threads = num_threads(thread_limit);
crossbeam_utils::thread::scope(move |s| {
let receive_result = {
let (send_input, receive_input) = crossbeam_channel::bounded::<I>(num_threads);
let (send_result, receive_result) = crossbeam_channel::bounded::<O>(num_threads);
for thread_id in 0..num_threads {
s.spawn({
let send_result = send_result.clone();
let receive_input = receive_input.clone();
let new_thread_state = new_thread_state.clone();
let consume = consume.clone();
move |_| {
let mut state = new_thread_state(thread_id);
for item in receive_input {
if send_result.send(consume(item, &mut state)).is_err() {
break;
}
}
}
});
}
s.spawn(move |_| {
for item in input {
if send_input.send(item).is_err() {
break;
}
}
});
receive_result
};
for item in receive_result {
drop(reducer.feed(item)?);
}
reducer.finalize()
})
.expect("no panic")
}
pub fn in_parallel_with_slice<I, S, E>(
input: &[I],
thread_limit: Option<usize>,
new_thread_state: impl FnMut(usize) -> S + Send + Clone,
consume: impl FnMut(&I, &mut S) -> Result<(), E> + Send + Clone,
mut periodic: impl FnMut() -> Option<std::time::Duration> + Send,
) -> Result<Vec<S>, E>
where
I: Send + Sync,
E: Send + Sync,
S: Send,
{
let num_threads = num_threads(thread_limit);
let num_items = input.len();
let mut results = Vec::with_capacity(num_threads);
let stop_everything = &AtomicBool::default();
crossbeam_utils::thread::scope({
move |s| {
s.spawn({
move |_| loop {
if stop_everything.load(Ordering::Relaxed) {
break;
}
match periodic() {
Some(duration) => std::thread::sleep(duration),
None => {
stop_everything.store(true, Ordering::Relaxed);
break;
}
}
}
});
let threads: Vec<_> = (0..num_threads)
.map(|thread_id| {
s.spawn({
let mut new_thread_state = new_thread_state.clone();
let mut consume = consume.clone();
move |_| {
let mut state = new_thread_state(thread_id);
for input_index in (thread_id..num_items).step_by(num_threads) {
if stop_everything.load(Ordering::Relaxed) {
break;
}
let item = &input[input_index];
if let Err(err) = consume(item, &mut state) {
stop_everything.store(true, Ordering::Relaxed);
return Err(err);
}
}
Ok(state)
}
})
})
.collect();
for thread in threads {
match thread.join() {
Ok(res) => {
results.push(res?);
}
Err(err) => {
stop_everything.store(true, Ordering::Relaxed);
std::panic::resume_unwind(err);
}
}
}
stop_everything.store(true, Ordering::Relaxed);
Ok(results)
}
})
.expect("no panic")
}