use std::sync::Arc;
use std::thread;
use std::time::Duration;
use qubit_function::Runnable;
use qubit_progress::{
Progress,
reporter::ProgressReporter,
};
use crate::BatchExecutionError;
use crate::BatchOutcome;
use crate::execute::{
BatchExecutionState,
BatchExecutor,
SequentialBatchExecutor,
};
use crate::utils::run_scoped_parallel;
use super::ParallelBatchExecutorBuildError;
use super::ParallelBatchExecutorBuilder;
use super::indexed_task::run_parallel_task;
#[derive(Clone)]
pub struct ParallelBatchExecutor {
pub(crate) thread_count: usize,
pub(crate) sequential_threshold: usize,
pub(crate) report_interval: Duration,
pub(crate) reporter: Arc<dyn ProgressReporter>,
}
impl ParallelBatchExecutor {
pub const DEFAULT_REPORT_INTERVAL: Duration = Duration::from_secs(5);
pub const DEFAULT_SEQUENTIAL_THRESHOLD: usize = 100;
#[inline]
pub fn default_thread_count() -> usize {
thread::available_parallelism()
.map(usize::from)
.unwrap_or(1)
}
#[inline]
pub fn builder() -> ParallelBatchExecutorBuilder {
ParallelBatchExecutorBuilder::default()
}
#[inline]
pub fn new(thread_count: usize) -> Result<Self, ParallelBatchExecutorBuildError> {
Self::builder().thread_count(thread_count).build()
}
#[inline]
pub const fn thread_count(&self) -> usize {
self.thread_count
}
#[inline]
pub const fn sequential_threshold(&self) -> usize {
self.sequential_threshold
}
#[inline]
pub const fn report_interval(&self) -> Duration {
self.report_interval
}
#[inline]
pub fn reporter(&self) -> &Arc<dyn ProgressReporter> {
&self.reporter
}
fn sequential_executor(&self) -> SequentialBatchExecutor {
SequentialBatchExecutor::builder()
.report_interval(self.report_interval)
.reporter_arc(Arc::clone(&self.reporter))
.build()
}
}
impl Default for ParallelBatchExecutor {
fn default() -> Self {
Self::builder()
.build()
.expect("default parallel batch executor should build")
}
}
impl BatchExecutor for ParallelBatchExecutor {
fn execute_with_count<T, E, I>(
&self,
tasks: I,
count: usize,
) -> Result<BatchOutcome<E>, BatchExecutionError<E>>
where
I: IntoIterator<Item = T>,
T: Runnable<E> + Send,
E: Send,
{
if count <= self.sequential_threshold || self.thread_count <= 1 {
return self.sequential_executor().execute_with_count(tasks, count);
}
let state = Arc::new(BatchExecutionState::new(count));
let progress = Progress::new(self.reporter.as_ref(), self.report_interval);
progress.report_started(state.progress_counters());
let mut actual_count = 0usize;
let worker_count = self.thread_count.min(count);
thread::scope(|scope| {
let reporter_state = Arc::clone(&state);
let running_progress =
progress.spawn_running_reporter(scope, move || reporter_state.progress_counters());
let running_point_handle = running_progress.point_handle();
let observer_state = Arc::clone(&state);
let worker_state = Arc::clone(&state);
actual_count = run_scoped_parallel(
tasks,
count,
worker_count,
move || observer_state.record_task_observed(),
move |index, task| {
run_parallel_task(&worker_state, index, task);
running_point_handle.report();
},
);
running_progress.stop_and_join();
});
let state = Arc::into_inner(state)
.expect("parallel batch execution state should have a single owner");
if actual_count < count {
let failed = progress.report_failed(state.progress_counters());
let result = state.into_outcome(failed.elapsed());
Err(BatchExecutionError::CountShortfall {
expected: count,
actual: actual_count,
outcome: result,
})
} else if actual_count > count {
let failed = progress.report_failed(state.progress_counters());
let result = state.into_outcome(failed.elapsed());
Err(BatchExecutionError::CountExceeded {
expected: count,
observed_at_least: actual_count,
outcome: result,
})
} else {
let finished = progress.report_finished(state.progress_counters());
let result = state.into_outcome(finished.elapsed());
Ok(result)
}
}
}