use crate::awake::Awake;
#[cfg(feature = "parallel")]
use crate::awake::LendGuard;
use crate::cost::{pays_at, width_for_by, DispatchKind};
use crate::error::{Error, Refused};
use crate::policy::{WaitOverride, WaitPolicy};
use crate::pool::Pool;
use crate::reserve::LeaderGuard;
use crate::sched;
use crate::split;
use crate::team::{split_raw, spread_raw, Plan, TeamExec};
use crate::track::Track;
use core::fmt;
use core::num::NonZeroUsize;
use core::ops::Range;
use core::sync::atomic::{AtomicUsize, Ordering};
use core::time::Duration;
#[must_use = "an Exec is a budget of execution slots; pass it to a kernel"]
pub struct Exec<'p> {
pool: Option<&'p Pool>,
budget: NonZeroUsize,
child: bool,
track: Track,
team: Option<TeamAuth<'p>>,
_leader: Option<LeaderGuard<'p>>,
}
#[derive(Clone, Copy)]
pub(crate) struct TeamAuth<'p> {
pub(crate) awake: &'p Awake,
pub(crate) ids: &'p [usize],
pub(crate) policy: WaitPolicy,
pub(crate) over: bool,
}
impl Exec<'static> {
pub const fn serial() -> Exec<'static> {
Exec {
pool: None,
budget: NonZeroUsize::MIN,
child: false,
track: Track::NONE,
team: None,
_leader: None,
}
}
}
impl<'p> Exec<'p> {
pub fn new(pool: &'p Pool, budget: usize) -> Result<Exec<'p>, Error> {
let budget = NonZeroUsize::new(budget).ok_or(Error::ZeroBudget)?;
Ok(Exec {
pool: Some(pool),
budget: budget.min(pool.width()),
child: false,
track: Track::root(),
team: None,
_leader: None,
})
}
pub(crate) fn from_team(
pool: &'p Pool,
auth: TeamAuth<'p>,
guard: LeaderGuard<'p>,
) -> Exec<'p> {
Exec {
pool: Some(pool),
budget: NonZeroUsize::new(auth.ids.len() + 1).unwrap_or(NonZeroUsize::MIN),
child: false,
track: Track::root(),
team: Some(auth),
_leader: Some(guard),
}
}
pub fn budget(&self) -> NonZeroUsize {
self.budget
}
pub fn dispatch_cost(&self, width: NonZeroUsize) -> Duration {
let width = width.min(self.budget);
if width.get() < 2 {
return Duration::MAX;
}
match (self.team, self.pool) {
(Some(t), _) => t.awake.costs.cost(DispatchKind::Team, width),
(None, Some(pool)) => {
let kind = if sched::in_pool(pool) {
DispatchKind::InPool
} else {
DispatchKind::Outside
};
pool.awake.costs.cost(kind, width)
}
(None, None) => Duration::MAX,
}
}
pub fn pays(&self, estimate: Duration, sections: u32, width: NonZeroUsize) -> bool {
pays_at(estimate, sections, self.dispatch_cost(width))
}
pub fn width_for(&self, estimate: Duration, sections: u32) -> NonZeroUsize {
width_for_by(estimate, sections, self.budget, |k| self.dispatch_cost(k))
}
pub fn is_serial(&self) -> bool {
self.budget.get() == 1
}
pub fn carries_team(&self) -> bool {
self.team.is_some()
}
pub fn limit(self, max: NonZeroUsize) -> Exec<'p> {
let budget = self.budget.min(max);
Exec {
budget,
team: self.team.map(|t| TeamAuth {
ids: &t.ids[..budget.get() - 1],
..t
}),
..self
}
}
pub fn reborrow(&mut self) -> Exec<'_> {
Exec {
pool: self.pool,
budget: self.budget,
child: self.child,
track: self.track.alias(),
team: self.team,
_leader: None,
}
}
pub fn fork<F>(self, jobs: usize, f: F)
where
F: Fn(usize, Exec<'_>) + Sync,
{
if jobs == 0 {
return;
}
let budget = self.budget.get();
let tasks = jobs.min(budget);
let _released = self.track.release();
let next = AtomicUsize::new(0);
let this = &self;
let body = |k: usize, team_ids: Option<&[usize]>| {
let _live = this.track.task();
let share = split::share(budget, tasks, k);
let child = || Exec {
pool: this.pool,
budget: NonZeroUsize::new(share).unwrap_or(NonZeroUsize::MIN),
child: true,
track: this.track.child(),
team: this
.team
.zip(team_ids)
.map(|(t, ids)| TeamAuth { ids, ..t }),
_leader: None,
};
if jobs <= tasks {
f(k, child());
} else {
loop {
let job = next.fetch_add(1, Ordering::Relaxed);
if job >= jobs {
break;
}
f(job, child());
}
}
};
match (self.team, self.pool) {
(Some(t), _) if tasks > 1 => {
let plan = Plan::Even {
size: budget,
children: tasks,
};
split_raw(Some(t.awake), t.ids, &plan, t.over, |k, ids| {
body(k, Some(ids));
});
}
(None, Some(pool)) if tasks > 1 => {
sched::run_tasks(pool, tasks, &self.track, |k| body(k, None));
}
_ => {
for k in 0..tasks {
body(k, self.team.map(|t| t.ids));
}
}
}
}
pub fn for_ranges<F>(self, len: usize, grain: NonZeroUsize, f: F)
where
F: Fn(Range<usize>) + Sync,
{
if len == 0 {
return;
}
let (chunk, tasks) = split::range_plan(len, grain.get(), self.budget.get());
let _released = self.track.release();
let track = &self.track;
let body = |k: usize| {
let _live = track.task();
f(split::chunk_range(len, chunk, k));
};
match (self.team, self.pool) {
(Some(t), _) if tasks > 1 => {
spread_raw(Some(t.awake), t.ids, tasks, t.over, body);
}
(None, Some(pool)) if tasks > 1 => sched::run_tasks(pool, tasks, track, body),
_ => {
for k in 0..tasks {
body(k);
}
}
}
}
pub fn broadcast<F>(self, f: F) -> Result<(), Refused>
where
F: Fn(usize) + Sync,
{
if let Some(refused) = self.refusal() {
return Err(refused);
}
if let Some(t) = self.team {
let _released = self.track.release();
let track = &self.track;
spread_raw(Some(t.awake), t.ids, self.budget.get(), t.over, |i| {
let _live = track.task();
f(i);
});
return Ok(());
}
let pool = match self.pool {
Some(pool) if self.budget.get() > 1 => pool,
_ => {
let _live = self.track.task();
f(0);
return Ok(());
}
};
let _released = self.track.release();
#[cfg(feature = "parallel")]
sched::broadcast(pool, &self.track, self.budget.get(), f);
#[cfg(not(feature = "parallel"))]
{
let _ = pool;
f(0);
}
Ok(())
}
pub fn can_broadcast(&self) -> bool {
self.refusal().is_none()
}
fn refusal(&self) -> Option<Refused> {
if self.team.is_some() {
return None;
}
let pool = match self.pool {
Some(pool) if self.budget.get() > 1 => pool,
_ => return None,
};
if self.child {
return Some(Refused::ForkChild);
}
if sched::in_pool(pool) {
return Some(Refused::CallerInPool);
}
None
}
pub fn external<R: Send>(self, f: impl FnOnce(NonZeroUsize) -> R + Send) -> R {
#[cfg(feature = "parallel")]
if let Some(pool) = self.pool {
if self.budget.get() > 1 {
let width = self.budget;
let _released = self.track.release();
let _lent = self.team.map(|t| LendGuard::new(t.awake, t.ids));
sched::unpin_worker(pool);
return sched::enter(pool, &self.track, move || f(width));
}
}
f(NonZeroUsize::MIN)
}
pub fn team<R>(
&mut self,
f: impl for<'team> FnOnce(TeamExec<'team>) -> R,
) -> Result<R, Refused> {
self.team_with(WaitOverride::new(), f)
}
pub fn team_with<R>(
&mut self,
over: WaitOverride,
f: impl for<'team> FnOnce(TeamExec<'team>) -> R,
) -> Result<R, Refused> {
if let Some(t) = self.team {
let mut team = TeamExec {
awake: Some(t.awake),
ids: t.ids,
policy: t.policy.with(over),
track: self.track.alias(),
over: t.over,
};
if self._leader.is_some() {
let oversubscribed = crate::pin::oversubscribed(team.size().get());
team.over = oversubscribed;
if let Some(awake) = team.awake {
awake.set_spin_budget(team.ids, awake.spin_budget(oversubscribed));
}
if let Some(auth) = self.team.as_mut() {
auth.over = oversubscribed;
}
if team.size().get() > 1 {
pin_scope(&mut team);
}
}
return Ok(f(team));
}
let pool = match self.pool {
Some(pool) if self.budget.get() > 1 => pool,
_ => {
let policy = self.pool.map_or(WaitPolicy::DEFAULT, Pool::wait_policy);
return Ok(f(TeamExec {
awake: None,
ids: &[],
policy: policy.with(over),
track: self.track.alias(),
over: false,
}));
}
};
let policy = pool.wait_policy().with(over);
let oversubscribed = crate::pin::oversubscribed(self.budget.get());
let reservation = pool.awake.acquire(
pool,
self.budget.get() - 1,
policy,
pool.awake.spin_budget(oversubscribed),
)?;
let mut team = TeamExec {
awake: Some(&pool.awake),
ids: &reservation.ids,
policy,
track: self.track.alias(),
over: oversubscribed,
};
let pinned = team.size().get() > 1 && pin_scope(&mut team);
if !pinned {
team.acknowledge();
}
Ok(f(team))
}
#[cfg(feature = "ledger")]
pub fn ledger_root(&self) -> Option<u64> {
self.track.root_id()
}
}
fn pin_scope(team: &mut TeamExec<'_>) -> bool {
if !team.awake.is_some_and(|awake| awake.pin) {
return false;
}
let Some(set) = crate::pin::read_here() else {
return false;
};
team.pin_section(&set);
true
}
impl fmt::Debug for Exec<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Exec")
.field("budget", &self.budget)
.field("pooled", &self.pool.is_some())
.field("child", &self.child)
.field("team", &self.team.is_some())
.finish()
}
}