use std::fmt::Debug;
use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
use std::sync::Arc;
use rayon::prelude::*;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CpuInnerParallelism {
None,
Rayon,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CpuExecutorReentrancy {
Rejected,
SameExecutor,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CpuExecutorAffinity {
TenferroPinnedVerified,
CallerDeclaredUnverified,
None,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CpuExecutorShutdown {
TenferroOwned,
CallerOwned,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct CpuDomainExecutorCapabilities {
pub worker_count: NonZeroUsize,
pub outer_parallelism: bool,
pub inner_parallelism: CpuInnerParallelism,
pub reentrancy: CpuExecutorReentrancy,
pub affinity: CpuExecutorAffinity,
pub shutdown: CpuExecutorShutdown,
}
#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)]
pub enum CpuDomainExecutorError {
#[error("CPU domain executor admission failed: {message}")]
Admission {
message: String,
},
#[error("CPU domain executor scheduling failed: {message}")]
Scheduling {
message: String,
},
#[error("CPU domain executor cancelled work: {message}")]
Cancellation {
message: String,
},
#[error("CPU domain executor worker panicked: {message}")]
PanicBridge {
message: String,
},
}
pub trait ScopedCpuJob: Send {
fn run(&mut self) -> Result<(), CpuDomainExecutorError>;
}
pub trait ScopedCpuJobs: Sync {
fn len(&self) -> usize;
fn is_empty(&self) -> bool {
self.len() == 0
}
fn run(&self, index: usize) -> Result<(), CpuDomainExecutorError>;
}
pub trait CpuDomainExecutor: Debug + Send + Sync + 'static {
fn capabilities(&self) -> CpuDomainExecutorCapabilities;
fn submit(&self, jobs: &dyn ScopedCpuJobs) -> Result<(), CpuDomainExecutorError>;
fn install(&self, job: &mut dyn ScopedCpuJob) -> Result<(), CpuDomainExecutorError>;
}
pub struct RayonCpuDomainExecutor {
pool: Arc<rayon::ThreadPool>,
}
impl std::fmt::Debug for RayonCpuDomainExecutor {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("RayonCpuDomainExecutor")
.field("worker_count", &self.pool.current_num_threads())
.finish_non_exhaustive()
}
}
impl RayonCpuDomainExecutor {
pub fn new(pool: Arc<rayon::ThreadPool>) -> Self {
Self { pool }
}
}
impl CpuDomainExecutor for RayonCpuDomainExecutor {
fn capabilities(&self) -> CpuDomainExecutorCapabilities {
let worker_count =
NonZeroUsize::new(self.pool.current_num_threads()).unwrap_or(NonZeroUsize::MIN);
CpuDomainExecutorCapabilities {
worker_count,
outer_parallelism: worker_count.get() > 1,
inner_parallelism: CpuInnerParallelism::Rayon,
reentrancy: CpuExecutorReentrancy::SameExecutor,
affinity: CpuExecutorAffinity::None,
shutdown: CpuExecutorShutdown::CallerOwned,
}
}
fn submit(&self, jobs: &dyn ScopedCpuJobs) -> Result<(), CpuDomainExecutorError> {
self.pool.install(|| {
(0..jobs.len())
.into_par_iter()
.try_for_each(|index| jobs.run(index))
})
}
fn install(&self, job: &mut dyn ScopedCpuJob) -> Result<(), CpuDomainExecutorError> {
self.pool.install(|| job.run())
}
}
pub(crate) struct ScopedJob<F, R> {
operation: Option<F>,
result: Option<R>,
}
pub(crate) fn scoped_job<F, R>(operation: F) -> ScopedJob<F, R>
where
F: FnOnce() -> R + Send,
R: Send,
{
ScopedJob {
operation: Some(operation),
result: None,
}
}
impl<F, R> ScopedJob<F, R> {
fn into_result(self) -> Result<R, CpuDomainExecutorError> {
self.result
.ok_or_else(|| CpuDomainExecutorError::Scheduling {
message: "executor returned success without running the scoped CPU job".to_string(),
})
}
}
impl<F, R> ScopedCpuJob for ScopedJob<F, R>
where
F: FnOnce() -> R + Send,
R: Send,
{
fn run(&mut self) -> Result<(), CpuDomainExecutorError> {
let operation =
self.operation
.take()
.ok_or_else(|| CpuDomainExecutorError::Scheduling {
message: "executor attempted to run a scoped CPU job more than once"
.to_string(),
})?;
self.result = Some(operation());
Ok(())
}
}
pub(crate) fn install_scoped<F, R>(
executor: &dyn CpuDomainExecutor,
operation: F,
) -> Result<R, CpuDomainExecutorError>
where
F: FnOnce() -> R + Send,
R: Send,
{
let mut job = scoped_job(operation);
executor.install(&mut job)?;
job.into_result()
}
pub(crate) struct IndexedJobs<F> {
len: usize,
run: F,
invalid_index_attempt: InvalidIndexAudit,
}
const INVALID_INDEX_EMPTY: u8 = 0;
const INVALID_INDEX_WRITING: u8 = 1;
const INVALID_INDEX_READY: u8 = 2;
struct InvalidIndexAudit {
state: AtomicU8,
index: AtomicUsize,
}
impl InvalidIndexAudit {
const fn new() -> Self {
Self {
state: AtomicU8::new(INVALID_INDEX_EMPTY),
index: AtomicUsize::new(0),
}
}
fn record(&self, index: usize) {
if self
.state
.compare_exchange(
INVALID_INDEX_EMPTY,
INVALID_INDEX_WRITING,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok()
{
self.index.store(index, Ordering::Relaxed);
self.state.store(INVALID_INDEX_READY, Ordering::Release);
}
}
fn load(&self) -> Option<usize> {
(self.state.load(Ordering::Acquire) == INVALID_INDEX_READY)
.then(|| self.index.load(Ordering::Relaxed))
}
}
pub(crate) fn indexed_jobs<F>(len: usize, run: F) -> IndexedJobs<F>
where
F: Fn(usize) -> Result<(), CpuDomainExecutorError> + Sync,
{
IndexedJobs {
len,
run,
invalid_index_attempt: InvalidIndexAudit::new(),
}
}
impl<F> IndexedJobs<F> {
pub(crate) fn invalid_index_attempt(&self) -> Option<usize> {
self.invalid_index_attempt.load()
}
}
impl<F> ScopedCpuJobs for IndexedJobs<F>
where
F: Fn(usize) -> Result<(), CpuDomainExecutorError> + Sync,
{
fn len(&self) -> usize {
self.len
}
fn run(&self, index: usize) -> Result<(), CpuDomainExecutorError> {
if index >= self.len {
self.invalid_index_attempt.record(index);
return Err(CpuDomainExecutorError::Scheduling {
message: format!(
"executor requested scoped CPU job index {index}, but the submission has {} jobs",
self.len
),
});
}
(self.run)(index)
}
}
#[cfg(test)]
mod tests;