use moirai_core::{error::ExecutorResult, Priority};
use crate::schedule::class::WorkClass;
use crate::schedule::runtime::{ScheduleMetrics, ThreadScheduler};
pub trait WorkSubmit: Send + Sync + 'static {
fn schedule<C, F>(
&self,
priority: Priority,
locality_hint: Option<usize>,
task: F,
) -> ExecutorResult<()>
where
C: WorkClass,
F: FnOnce(usize) + Send + 'static;
}
pub trait SchedulerControl {
fn pending_tasks(&self) -> usize;
fn active_workers(&self) -> usize;
fn worker_count(&self) -> usize;
fn has_work(&self) -> bool;
fn join(&self) -> ExecutorResult<()>;
fn shutdown(&self);
fn metrics(&self) -> ScheduleMetrics;
}
pub trait DataParallel {
fn for_each_indexed<C, F>(
&self,
priority: Priority,
locality_hint: Option<usize>,
count: usize,
task: F,
) -> ExecutorResult<()>
where
C: WorkClass,
F: Fn(usize) + Send + Sync;
fn map_reduce_indexed<C, T, Map, Reduce>(
&self,
priority: Priority,
locality_hint: Option<usize>,
count: usize,
identity: T,
map: Map,
reduce: Reduce,
) -> ExecutorResult<T>
where
C: WorkClass,
T: Send + Clone,
Map: Fn(usize) -> T + Send + Sync,
Reduce: Fn(T, T) -> T + Send + Sync;
}
pub trait WorkScheduler: WorkSubmit + SchedulerControl + DataParallel + Clone {}
impl<S> WorkScheduler for S where S: WorkSubmit + SchedulerControl + DataParallel + Clone {}
impl<const QUEUE_CAPACITY: usize, const SPIN_LIMIT: usize> WorkSubmit
for ThreadScheduler<QUEUE_CAPACITY, SPIN_LIMIT>
{
fn schedule<C, F>(
&self,
priority: Priority,
locality_hint: Option<usize>,
task: F,
) -> ExecutorResult<()>
where
C: WorkClass,
F: FnOnce(usize) + Send + 'static,
{
ThreadScheduler::schedule::<C, F>(self, priority, locality_hint, task)
}
}
impl<const QUEUE_CAPACITY: usize, const SPIN_LIMIT: usize> SchedulerControl
for ThreadScheduler<QUEUE_CAPACITY, SPIN_LIMIT>
{
fn pending_tasks(&self) -> usize {
ThreadScheduler::pending_tasks(self)
}
fn active_workers(&self) -> usize {
ThreadScheduler::active_workers(self)
}
fn worker_count(&self) -> usize {
ThreadScheduler::worker_count(self)
}
fn has_work(&self) -> bool {
ThreadScheduler::has_work(self)
}
fn join(&self) -> ExecutorResult<()> {
ThreadScheduler::join(self)
}
fn shutdown(&self) {
ThreadScheduler::shutdown(self)
}
fn metrics(&self) -> ScheduleMetrics {
ThreadScheduler::metrics(self)
}
}
impl<const QUEUE_CAPACITY: usize, const SPIN_LIMIT: usize> DataParallel
for ThreadScheduler<QUEUE_CAPACITY, SPIN_LIMIT>
{
fn for_each_indexed<C, F>(
&self,
priority: Priority,
locality_hint: Option<usize>,
count: usize,
task: F,
) -> ExecutorResult<()>
where
C: WorkClass,
F: Fn(usize) + Send + Sync,
{
ThreadScheduler::for_each_indexed::<C, F>(self, priority, locality_hint, count, task)
}
fn map_reduce_indexed<C, T, Map, Reduce>(
&self,
priority: Priority,
locality_hint: Option<usize>,
count: usize,
identity: T,
map: Map,
reduce: Reduce,
) -> ExecutorResult<T>
where
C: WorkClass,
T: Send + Clone,
Map: Fn(usize) -> T + Send + Sync,
Reduce: Fn(T, T) -> T + Send + Sync,
{
ThreadScheduler::map_reduce_indexed::<C, T, Map, Reduce>(
self,
priority,
locality_hint,
count,
identity,
map,
reduce,
)
}
}