use std::{
sync::Arc,
thread,
};
use rayon::ThreadPoolBuilder as RayonThreadPoolBuilder;
use crate::{
rayon_executor_service::RayonExecutorService,
rayon_executor_service_build_error::RayonExecutorServiceBuildError,
rayon_executor_service_state::RayonExecutorServiceState,
};
const DEFAULT_THREAD_NAME_PREFIX: &str = "qubit-rayon-executor";
#[derive(Debug, Clone)]
pub struct RayonExecutorServiceBuilder {
num_threads: usize,
thread_name_prefix: String,
stack_size: Option<usize>,
}
impl RayonExecutorServiceBuilder {
#[inline]
pub fn num_threads(mut self, num_threads: usize) -> Self {
self.num_threads = num_threads;
self
}
#[inline]
pub fn thread_name_prefix(mut self, prefix: &str) -> Self {
self.thread_name_prefix = prefix.to_owned();
self
}
#[inline]
pub fn stack_size(mut self, stack_size: usize) -> Self {
self.stack_size = Some(stack_size);
self
}
pub fn build(self) -> Result<RayonExecutorService, RayonExecutorServiceBuildError> {
if self.num_threads == 0 {
return Err(RayonExecutorServiceBuildError::ZeroThreadCount);
}
if self.stack_size == Some(0) {
return Err(RayonExecutorServiceBuildError::ZeroStackSize);
}
let prefix = self.thread_name_prefix;
let mut builder = RayonThreadPoolBuilder::new()
.num_threads(self.num_threads)
.thread_name(move |index| format!("{prefix}-{index}"));
if let Some(stack_size) = self.stack_size {
builder = builder.stack_size(stack_size);
}
let pool = Arc::new(builder.build()?);
Ok(RayonExecutorService {
pool,
state: Arc::new(RayonExecutorServiceState::new()),
})
}
}
impl Default for RayonExecutorServiceBuilder {
fn default() -> Self {
Self {
num_threads: default_rayon_thread_count(),
thread_name_prefix: DEFAULT_THREAD_NAME_PREFIX.to_owned(),
stack_size: None,
}
}
}
fn default_rayon_thread_count() -> usize {
thread::available_parallelism()
.map(usize::from)
.unwrap_or(1)
}