use std::cell::UnsafeCell;
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use serde::{Deserialize, Serialize};
pub fn new_diagnostic_id() -> String {
static NEXT: AtomicU64 = AtomicU64::new(1);
let id = NEXT
.try_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
.expect("diagnostic identity exhausted");
id.to_string()
}
#[derive(Default)]
struct TaskTiming {
waits: Vec<(u32, u32, u64, u64)>,
start: u64,
end: u64,
code: Option<i64>,
}
struct ExecutionTrace {
start: Instant,
tasks: Vec<Mutex<TaskTiming>>,
}
impl ExecutionTrace {
fn nanos(&self) -> u64 {
self.start.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct LaneTaskSpec {
pub lane: u32,
pub waits: Vec<(u32, u32)>,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum LaneScheduleError {
#[error("task {task} names lane {lane}, but the schedule has {lanes} lanes")]
Lane { task: usize, lane: u32, lanes: u32 },
#[error("task {task} has an invalid wait on lane {lane} for {completed} tasks")]
Wait {
task: usize,
lane: u32,
completed: u32,
},
}
#[derive(Debug, Clone)]
pub struct LaneSchedule {
lanes: u32,
tasks: Vec<LaneTaskSpec>,
lane_tasks: Vec<Vec<usize>>,
diagnostic_id: String,
}
impl PartialEq for LaneSchedule {
fn eq(&self, other: &Self) -> bool {
self.lanes == other.lanes
&& self.tasks == other.tasks
&& self.lane_tasks == other.lane_tasks
}
}
impl Eq for LaneSchedule {}
impl LaneSchedule {
pub fn new(lanes: u32, tasks: Vec<LaneTaskSpec>) -> Result<Self, LaneScheduleError> {
let lanes = lanes.max(1);
let mut lane_tasks = vec![Vec::new(); lanes as usize];
for (task, spec) in tasks.iter().enumerate() {
if spec.lane >= lanes {
return Err(LaneScheduleError::Lane {
task,
lane: spec.lane,
lanes,
});
}
for &(lane, completed) in &spec.waits {
let listed_before = lane_tasks
.get(lane as usize)
.is_some_and(|listed: &Vec<usize>| completed as usize <= listed.len());
if lane == spec.lane || completed == 0 || !listed_before {
return Err(LaneScheduleError::Wait {
task,
lane,
completed,
});
}
}
lane_tasks[spec.lane as usize].push(task);
}
Ok(Self {
lanes,
tasks,
lane_tasks,
diagnostic_id: new_diagnostic_id(),
})
}
pub fn with_diagnostic_id(mut self, name: String) -> Self {
self.diagnostic_id = name;
self
}
pub fn lanes(&self) -> u32 {
self.lanes
}
pub fn tasks(&self) -> &[LaneTaskSpec] {
&self.tasks
}
pub fn run_sequential(&self, runner: &dyn LaneTaskRunner) -> Result<(), LaneTaskFailure> {
for task in 0..self.tasks.len() {
let code = runner.run_task(task);
if code != 0 {
return Err(LaneTaskFailure { task, code });
}
}
Ok(())
}
}
pub trait LaneTaskRunner: Sync {
fn run_task(&self, task: usize) -> i64;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct LaneTaskFailure {
pub task: usize,
pub code: i64,
}
pub const LANE_TASK_PANICKED: i64 = i64::MIN;
#[repr(align(256))]
struct Padded<T>(T);
const _: () = assert!(std::mem::align_of::<Padded<u8>>() == celox_state_layout::LANE_SEGMENT_ALIGN);
#[derive(Clone, Copy)]
struct Job {
schedule: *const LaneSchedule,
runner: *const (dyn LaneTaskRunner + 'static),
trace: *const ExecutionTrace,
}
const LANE_DONE: u64 = u32::MAX as u64;
const MAX_EPOCH: u64 = u32::MAX as u64;
struct Shared {
lanes: usize,
job_epoch: Padded<AtomicU64>,
job: UnsafeCell<Option<Job>>,
progress: Box<[Padded<AtomicU64>]>,
abort: Padded<AtomicBool>,
failure: Mutex<Option<LaneTaskFailure>>,
sleeping: Box<[Padded<AtomicBool>]>,
shutdown: AtomicBool,
idle_spin: Duration,
}
unsafe impl Sync for Shared {}
unsafe impl Send for Shared {}
impl Shared {
fn record_failure(&self, failure: LaneTaskFailure) {
let mut slot = self
.failure
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if slot.is_none_or(|current| {
let priority = |failure: LaneTaskFailure| {
(crate::comb_fatal_site(failure.code).is_some(), failure.task)
};
priority(failure) < priority(current)
}) {
*slot = Some(failure);
}
self.abort.0.store(true, Ordering::Release);
}
fn execute_lane(&self, lane: usize, epoch: u64, job: Job) {
let (schedule, runner) = unsafe { (&*job.schedule, &*job.runner) };
let trace = unsafe { job.trace.as_ref() };
let base = epoch << 32;
let Some(tasks) = schedule.lane_tasks.get(lane) else {
return;
};
for (sequence, &task) in tasks.iter().enumerate() {
let mut timing = trace.map(|_| TaskTiming::default());
let mut aborted = self.abort.0.load(Ordering::Acquire);
for &(other, completed) in &schedule.tasks[task].waits {
if aborted {
break;
}
let target = base | u64::from(completed);
let progress = &self.progress[other as usize].0;
let mut spins = 0u32;
let wait_start = trace.map(ExecutionTrace::nanos);
while progress.load(Ordering::Acquire) < target {
if self.abort.0.load(Ordering::Acquire) {
aborted = true;
break;
}
wait_pause(&mut spins);
}
if let (Some(trace), Some(timing), Some(start)) = (trace, &mut timing, wait_start) {
timing.waits.push((other, completed, start, trace.nanos()));
}
}
if !aborted {
if let (Some(trace), Some(timing)) = (trace, &mut timing) {
timing.start = trace.nanos();
}
let code = catch_unwind(AssertUnwindSafe(|| runner.run_task(task)))
.unwrap_or(LANE_TASK_PANICKED);
if let (Some(trace), Some(timing)) = (trace, &mut timing) {
timing.end = trace.nanos();
timing.code = Some(code);
}
if code != 0 {
self.record_failure(LaneTaskFailure { task, code });
}
}
if let (Some(trace), Some(timing)) = (trace, timing) {
*trace.tasks[task].lock().unwrap_or_else(|e| e.into_inner()) = timing;
}
self.progress[lane]
.0
.store(base | (sequence as u64 + 1), Ordering::Release);
}
}
}
fn wait_pause(spins: &mut u32) {
*spins = spins.wrapping_add(1);
if *spins < 2048 {
std::hint::spin_loop();
} else {
std::thread::yield_now();
}
}
pub struct LanePool {
shared: Arc<Shared>,
workers: Vec<JoinHandle<()>>,
epoch: u64,
diagnostic_id: String,
}
impl std::fmt::Debug for LanePool {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LanePool")
.field("lanes", &self.shared.lanes)
.finish_non_exhaustive()
}
}
impl LanePool {
pub const DEFAULT_IDLE_SPIN: Duration = Duration::from_micros(200);
pub fn new(lanes: u32) -> std::io::Result<Self> {
Self::with_idle_spin(lanes, Self::DEFAULT_IDLE_SPIN)
}
pub fn with_idle_spin(lanes: u32, idle_spin: Duration) -> std::io::Result<Self> {
let lanes = lanes.max(1) as usize;
let shared = Arc::new(Shared {
lanes,
job_epoch: Padded(AtomicU64::new(0)),
job: UnsafeCell::new(None),
progress: (0..lanes).map(|_| Padded(AtomicU64::new(0))).collect(),
abort: Padded(AtomicBool::new(false)),
failure: Mutex::new(None),
sleeping: (0..lanes).map(|_| Padded(AtomicBool::new(false))).collect(),
shutdown: AtomicBool::new(false),
idle_spin,
});
let mut workers = Vec::with_capacity(lanes - 1);
for lane in 1..lanes {
let worker_shared = Arc::clone(&shared);
let spawned = std::thread::Builder::new()
.name(format!("celox-lane-{lane}"))
.spawn(move || worker(&worker_shared, lane));
match spawned {
Ok(handle) => workers.push(handle),
Err(error) => {
let pool = Self {
shared,
workers,
epoch: 0,
diagnostic_id: new_diagnostic_id(),
};
drop(pool);
return Err(error);
}
}
}
Ok(Self {
shared,
workers,
epoch: 0,
diagnostic_id: new_diagnostic_id(),
})
}
pub fn lanes(&self) -> u32 {
self.shared.lanes as u32
}
pub fn run(
&mut self,
schedule: &LaneSchedule,
runner: &dyn LaneTaskRunner,
) -> Result<(), LaneTaskFailure> {
if schedule.lanes as usize > self.shared.lanes {
tracing::trace!(target: "celox::parallel", event = "pool_fallback", reason = "insufficient_workers", requested = schedule.lanes, available = self.shared.lanes);
return schedule.run_sequential(runner);
}
let trace =
tracing::enabled!(target: "celox::parallel", tracing::Level::TRACE).then(|| {
ExecutionTrace {
start: Instant::now(),
tasks: (0..schedule.tasks.len())
.map(|_| Mutex::new(TaskTiming::default()))
.collect(),
}
});
let shared = &*self.shared;
if self.epoch >= MAX_EPOCH {
for progress in shared.progress.iter() {
progress.0.store(0, Ordering::Relaxed);
}
self.epoch = 0;
}
self.epoch += 1;
let epoch = self.epoch;
let runner: &'static dyn LaneTaskRunner = unsafe {
std::mem::transmute::<&dyn LaneTaskRunner, &'static dyn LaneTaskRunner>(runner)
};
let job = Job {
schedule: schedule as *const LaneSchedule,
runner: runner as *const dyn LaneTaskRunner,
trace: trace
.as_ref()
.map_or(std::ptr::null(), |trace| trace as *const ExecutionTrace),
};
unsafe {
*shared.job.get() = Some(job);
}
shared.abort.0.store(false, Ordering::Relaxed);
shared.job_epoch.0.store(epoch, Ordering::SeqCst);
for (index, worker) in self.workers.iter().enumerate() {
let sleeping = &shared.sleeping[index + 1].0;
if sleeping.load(Ordering::SeqCst) && sleeping.swap(false, Ordering::SeqCst) {
worker.thread().unpark();
}
}
shared.execute_lane(0, epoch, job);
let done = (epoch << 32) | LANE_DONE;
for lane in 1..shared.lanes {
let mut spins = 0u32;
while shared.progress[lane].0.load(Ordering::Acquire) != done {
wait_pause(&mut spins);
}
}
if let Some(trace) = &trace {
let elapsed_ns = trace.nanos();
let schedule_id = &schedule.diagnostic_id;
let pool = &self.diagnostic_id;
tracing::trace!(target: "celox::parallel", event = "dispatch", pool, epoch,
schedule = schedule_id, lanes = schedule.lanes, elapsed_ns);
for (task, timing) in trace.tasks.iter().enumerate() {
let timing = timing.lock().unwrap_or_else(|e| e.into_inner());
tracing::trace!(target: "celox::parallel", event = "task_run", pool, epoch,
schedule = schedule_id, task, lane = schedule.tasks[task].lane,
start_ns = timing.start, end_ns = timing.end, skipped = timing.code.is_none(), code = timing.code.unwrap_or(0));
for &(other, completed, start_ns, end_ns) in &timing.waits {
tracing::trace!(target: "celox::parallel", event = "task_wait", pool, epoch,
task, lane = schedule.tasks[task].lane, other, completed, start_ns, end_ns);
}
}
}
unsafe {
*shared.job.get() = None;
}
let failure = shared
.failure
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.take();
match failure {
Some(failure) => Err(failure),
None => Ok(()),
}
}
}
fn worker(shared: &Shared, lane: usize) {
let mut seen = 0u64;
loop {
let mut idle_since = None::<Instant>;
let mut spins = 0u32;
let epoch = loop {
let epoch = shared.job_epoch.0.load(Ordering::Acquire);
if epoch != seen || shared.shutdown.load(Ordering::Acquire) {
break epoch;
}
spins = spins.wrapping_add(1);
std::hint::spin_loop();
if !spins.is_multiple_of(256) {
continue;
}
let since = *idle_since.get_or_insert_with(Instant::now);
if since.elapsed() < shared.idle_spin {
continue;
}
shared.sleeping[lane].0.store(true, Ordering::SeqCst);
let epoch = shared.job_epoch.0.load(Ordering::SeqCst);
if epoch != seen || shared.shutdown.load(Ordering::SeqCst) {
shared.sleeping[lane].0.store(false, Ordering::SeqCst);
break epoch;
}
std::thread::park();
shared.sleeping[lane].0.store(false, Ordering::SeqCst);
idle_since = None;
};
if shared.shutdown.load(Ordering::Acquire) {
return;
}
seen = epoch;
let job = unsafe { *shared.job.get() };
if let Some(job) = job {
shared.execute_lane(lane, epoch, job);
}
shared.progress[lane]
.0
.store((epoch << 32) | LANE_DONE, Ordering::Release);
}
}
impl Drop for LanePool {
fn drop(&mut self) {
self.shared.shutdown.store(true, Ordering::SeqCst);
self.shared.job_epoch.0.fetch_add(1, Ordering::SeqCst);
for worker in &self.workers {
worker.thread().unpark();
}
for worker in self.workers.drain(..) {
let _ = worker.join();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::AtomicI64;
struct Recorder {
log: Mutex<Vec<usize>>,
done: Vec<AtomicBool>,
predecessors: Vec<Vec<usize>>,
fail: Option<usize>,
calls: AtomicI64,
}
impl LaneTaskRunner for Recorder {
fn run_task(&self, task: usize) -> i64 {
self.calls.fetch_add(1, Ordering::Relaxed);
for &predecessor in &self.predecessors[task] {
assert!(
self.done[predecessor].load(Ordering::Acquire),
"task {task} ran before {predecessor}"
);
}
self.log.lock().unwrap().push(task);
self.done[task].store(true, Ordering::Release);
if self.fail == Some(task) { 7 } else { 0 }
}
}
fn chain_schedule() -> (LaneSchedule, Vec<Vec<usize>>) {
let tasks = vec![
LaneTaskSpec {
lane: 0,
waits: vec![],
},
LaneTaskSpec {
lane: 1,
waits: vec![(0, 1)],
},
LaneTaskSpec {
lane: 0,
waits: vec![],
},
LaneTaskSpec {
lane: 2,
waits: vec![(1, 1), (0, 2)],
},
];
let predecessors = vec![vec![], vec![0], vec![0], vec![1, 2]];
(LaneSchedule::new(3, tasks).unwrap(), predecessors)
}
fn recorder(predecessors: Vec<Vec<usize>>, fail: Option<usize>) -> Recorder {
Recorder {
log: Mutex::new(Vec::new()),
done: (0..predecessors.len())
.map(|_| AtomicBool::new(false))
.collect(),
predecessors,
fail,
calls: AtomicI64::new(0),
}
}
#[test]
fn schedules_reject_forward_or_self_waits() {
let forward = vec![LaneTaskSpec {
lane: 0,
waits: vec![(1, 1)],
}];
assert!(LaneSchedule::new(2, forward).is_err());
let self_wait = vec![
LaneTaskSpec {
lane: 0,
waits: vec![],
},
LaneTaskSpec {
lane: 0,
waits: vec![(0, 1)],
},
];
assert!(LaneSchedule::new(1, self_wait).is_err());
}
#[test]
fn repeated_executions_honour_waits() {
let (schedule, predecessors) = chain_schedule();
let mut pool = LanePool::with_idle_spin(3, Duration::from_micros(0)).unwrap();
for _ in 0..2000 {
let runner = recorder(predecessors.clone(), None);
pool.run(&schedule, &runner).unwrap();
assert_eq!(runner.log.lock().unwrap().len(), 4);
}
}
#[test]
fn failures_skip_pending_work_and_report_the_first_task() {
let (schedule, predecessors) = chain_schedule();
let mut pool = LanePool::new(3).unwrap();
let runner = recorder(predecessors.clone(), Some(0));
assert_eq!(
pool.run(&schedule, &runner),
Err(LaneTaskFailure { task: 0, code: 7 })
);
assert!(!runner.done[1].load(Ordering::Acquire));
assert!(!runner.done[3].load(Ordering::Acquire));
let runner = recorder(predecessors, None);
pool.run(&schedule, &runner).unwrap();
assert_eq!(runner.calls.load(Ordering::Relaxed), 4);
}
#[test]
fn concurrent_loop_and_internal_failures_take_precedence_over_fatal_captures() {
struct Runner {
started: std::sync::Barrier,
codes: [i64; 2],
}
impl LaneTaskRunner for Runner {
fn run_task(&self, task: usize) -> i64 {
self.started.wait();
self.codes[task]
}
}
let schedule = LaneSchedule::new(
2,
vec![
LaneTaskSpec {
lane: 0,
waits: vec![],
},
LaneTaskSpec {
lane: 1,
waits: vec![],
},
],
)
.unwrap();
let mut pool = LanePool::new(2).unwrap();
for code in [1, 2000, LANE_TASK_PANICKED] {
for task in [0, 1] {
let mut codes = [crate::comb_fatal_code(1); 2];
codes[task] = code;
let runner = Runner {
started: std::sync::Barrier::new(2),
codes,
};
assert_eq!(
pool.run(&schedule, &runner),
Err(LaneTaskFailure { task, code })
);
}
}
let runner = Runner {
started: std::sync::Barrier::new(2),
codes: [crate::comb_fatal_code(2), crate::comb_fatal_code(1)],
};
assert_eq!(
pool.run(&schedule, &runner),
Err(LaneTaskFailure {
task: 0,
code: crate::comb_fatal_code(2),
})
);
}
#[test]
fn epochs_restart_without_releasing_waits_early() {
let (schedule, predecessors) = chain_schedule();
let mut pool = LanePool::new(3).unwrap();
pool.epoch = MAX_EPOCH - 2;
for _ in 0..4 {
let runner = recorder(predecessors.clone(), None);
pool.run(&schedule, &runner).unwrap();
assert_eq!(runner.calls.load(Ordering::Relaxed), 4);
}
assert!(pool.epoch < 4);
}
#[test]
fn sequential_fallback_preserves_list_order() {
let (schedule, predecessors) = chain_schedule();
let runner = recorder(predecessors, None);
schedule.run_sequential(&runner).unwrap();
assert_eq!(*runner.log.lock().unwrap(), vec![0, 1, 2, 3]);
}
#[test]
fn diagnostic_ids_distinguish_equal_topologies_and_recreated_pools() {
let (first, _) = chain_schedule();
let (second, _) = chain_schedule();
assert_ne!(first.diagnostic_id, second.diagnostic_id);
assert_eq!(first, second);
assert_eq!(first, second.clone().with_diagnostic_id("different".into()));
let different = LaneSchedule::new(
1,
vec![LaneTaskSpec {
lane: 0,
waits: vec![],
}],
)
.unwrap();
assert_ne!(first, different);
assert_eq!(first.clone().diagnostic_id, first.diagnostic_id);
let first_pool = LanePool::new(1).unwrap();
let id = first_pool.diagnostic_id.clone();
drop(first_pool);
let second_pool = LanePool::new(1).unwrap();
assert_ne!(id, second_pool.diagnostic_id);
}
#[test]
fn traced_failure_and_recovery_record_skips_and_real_waits() {
use std::collections::BTreeMap;
#[derive(Clone)]
struct Capture(Arc<Mutex<Vec<BTreeMap<String, String>>>>);
struct Fields(BTreeMap<String, String>);
impl tracing::field::Visit for Fields {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
self.0.insert(field.name().into(), format!("{value:?}"));
}
}
impl tracing::Subscriber for Capture {
fn max_level_hint(&self) -> Option<tracing::metadata::LevelFilter> {
Some(tracing::metadata::LevelFilter::TRACE)
}
fn enabled(&self, metadata: &tracing::Metadata<'_>) -> bool {
metadata.target() == "celox::parallel"
}
fn register_callsite(
&self,
_: &'static tracing::Metadata<'static>,
) -> tracing::subscriber::Interest {
tracing::subscriber::Interest::sometimes()
}
fn new_span(&self, _: &tracing::span::Attributes<'_>) -> tracing::span::Id {
tracing::span::Id::from_u64(1)
}
fn record(&self, _: &tracing::span::Id, _: &tracing::span::Record<'_>) {}
fn record_follows_from(&self, _: &tracing::span::Id, _: &tracing::span::Id) {}
fn enter(&self, _: &tracing::span::Id) {}
fn exit(&self, _: &tracing::span::Id) {}
fn event(&self, event: &tracing::Event<'_>) {
let mut fields = Fields(BTreeMap::new());
event.record(&mut fields);
self.0.lock().unwrap().push(fields.0);
}
}
let events = Arc::new(Mutex::new(Vec::new()));
tracing::subscriber::with_default(Capture(events.clone()), || {
let (schedule, predecessors) = chain_schedule();
let mut pool = LanePool::new(3).unwrap();
assert!(
pool.run(&schedule, &recorder(predecessors.clone(), Some(0)))
.is_err()
);
pool.run(&schedule, &recorder(predecessors, None)).unwrap();
});
let events = events.lock().unwrap();
let rows = |name: &str| {
let name = format!("{name:?}");
events.iter().filter(move |e| e["event"] == name)
};
assert_eq!(rows("dispatch").count(), 2, "{events:?}");
let runs = rows("task_run").collect::<Vec<_>>();
assert_eq!(runs.len(), 8);
assert!(runs.iter().any(|e| e["skipped"] == "true"));
for event in runs.iter().filter(|e| e["epoch"] == "2") {
assert_eq!(event["skipped"], "false");
assert_eq!(event["code"], "0");
assert!(
event["end_ns"].parse::<u64>().unwrap()
>= event["start_ns"].parse::<u64>().unwrap()
);
}
assert_eq!(rows("task_wait").filter(|e| e["epoch"] == "2").count(), 3);
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KernelChoice {
Parallel,
Sequential,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum KernelSelection {
#[default]
Measured,
Parallel,
Sequential,
}
#[derive(Debug)]
pub struct KernelSelector {
diagnostic_id: String,
state: SelectorState,
recheck_interval: u64,
}
#[derive(Debug, Clone)]
enum SelectorState {
Calibrating {
round: u32,
position: u32,
samples: [Vec<u64>; 2],
},
Decided {
choice: KernelChoice,
remaining: u64,
},
}
impl Clone for KernelSelector {
fn clone(&self) -> Self {
Self {
diagnostic_id: new_diagnostic_id(),
state: self.state.clone(),
recheck_interval: self.recheck_interval,
}
}
}
impl KernelSelector {
const BLOCK: u32 = 16;
const DISCARDED: u32 = 4;
const ROUNDS: u32 = 3;
const PARALLEL_MARGIN_PERCENT: u64 = 95;
const FIRST_RECHECK: u64 = 1 << 16;
const MAX_RECHECK: u64 = 1 << 26;
pub fn new(selection: KernelSelection) -> Self {
let decided = |choice| SelectorState::Decided {
choice,
remaining: u64::MAX,
};
match selection {
KernelSelection::Measured => Self {
diagnostic_id: new_diagnostic_id(),
state: Self::calibrating(),
recheck_interval: Self::FIRST_RECHECK,
},
KernelSelection::Parallel => Self {
diagnostic_id: new_diagnostic_id(),
state: decided(KernelChoice::Parallel),
recheck_interval: 0,
},
KernelSelection::Sequential => Self {
diagnostic_id: new_diagnostic_id(),
state: decided(KernelChoice::Sequential),
recheck_interval: 0,
},
}
}
pub fn diagnostic_id(&self) -> &str {
&self.diagnostic_id
}
fn calibrating() -> SelectorState {
SelectorState::Calibrating {
round: 0,
position: 0,
samples: [Vec::new(), Vec::new()],
}
}
pub fn next(&self) -> (KernelChoice, bool) {
match &self.state {
SelectorState::Decided { choice, .. } => (*choice, false),
SelectorState::Calibrating { round, .. } => {
let choice = if round % 2 == 0 {
KernelChoice::Parallel
} else {
KernelChoice::Sequential
};
(choice, true)
}
}
}
pub fn decided(&self) -> Option<KernelChoice> {
match self.state {
SelectorState::Decided { choice, .. } => Some(choice),
SelectorState::Calibrating { .. } => None,
}
}
pub fn advance(&mut self, calls: u64) {
if self.recheck_interval == 0 {
return;
}
if let SelectorState::Decided { remaining, .. } = &mut self.state {
*remaining = remaining.saturating_sub(calls);
if *remaining == 0 {
self.state = Self::calibrating();
}
}
}
pub fn record(&mut self, choice: KernelChoice, nanos: u64) {
let SelectorState::Calibrating {
round,
position,
samples,
} = &mut self.state
else {
return;
};
if *position >= Self::DISCARDED {
samples[usize::from(choice == KernelChoice::Sequential)].push(nanos);
}
*position += 1;
if *position < Self::BLOCK {
return;
}
*position = 0;
*round += 1;
if *round < 2 * Self::ROUNDS {
return;
}
let median = |values: &mut Vec<u64>| {
values.sort_unstable();
values.get(values.len() / 2).copied().unwrap_or(u64::MAX)
};
let [parallel, sequential] = samples;
let parallel = median(parallel);
let sequential = median(sequential);
let choice = if parallel.saturating_mul(100)
< sequential.saturating_mul(Self::PARALLEL_MARGIN_PERCENT)
{
KernelChoice::Parallel
} else {
KernelChoice::Sequential
};
tracing::debug!(target: "celox::parallel::selection", event = "selector_decision",
selector = self.diagnostic_id.as_str(), parallel_ns = parallel, sequential_ns = sequential, choice = ?choice);
tracing::debug!(
"[parallel] kernel selection: parallel median {parallel} ns, sequential median {sequential} ns -> {choice:?}"
);
self.state = SelectorState::Decided {
choice,
remaining: self.recheck_interval,
};
self.recheck_interval = self
.recheck_interval
.saturating_mul(2)
.min(Self::MAX_RECHECK);
}
}
#[cfg(test)]
mod selector_tests {
use super::*;
fn calibrate(parallel: u64, sequential: u64) -> KernelChoice {
let mut selector = KernelSelector::new(KernelSelection::Measured);
while selector.decided().is_none() {
let (choice, timed) = selector.next();
assert!(timed);
selector.record(
choice,
match choice {
KernelChoice::Parallel => parallel,
KernelChoice::Sequential => sequential,
},
);
}
selector.decided().unwrap()
}
#[test]
fn selector_identity_survives_moves_without_reuse() {
let selector = KernelSelector::new(KernelSelection::Measured);
let id = selector.diagnostic_id().to_owned();
let mut moved = Box::new(selector);
assert_eq!(moved.diagnostic_id(), id);
let cloned = moved.clone();
assert_ne!(cloned.diagnostic_id(), id);
assert_eq!(cloned.next(), moved.next());
while moved.decided().is_none() {
let (choice, _) = moved.next();
moved.record(choice, 100);
}
assert_eq!(moved.diagnostic_id(), id);
let decided_clone = moved.clone();
assert_ne!(decided_clone.diagnostic_id(), id);
assert_eq!(decided_clone.next(), moved.next());
drop(moved);
let replacement = KernelSelector::new(KernelSelection::Measured);
assert_ne!(replacement.diagnostic_id(), id);
assert_ne!(replacement.diagnostic_id(), cloned.diagnostic_id());
}
#[test]
fn measured_selection_keeps_the_faster_kernel() {
assert_eq!(calibrate(500, 1000), KernelChoice::Parallel);
assert_eq!(calibrate(1000, 500), KernelChoice::Sequential);
assert_eq!(calibrate(980, 1000), KernelChoice::Sequential);
}
#[test]
fn measured_decisions_are_checked_again_with_growing_intervals() {
let mut selector = KernelSelector::new(KernelSelection::Measured);
let mut calibrations = 0;
let mut checks = Vec::new();
let mut calls = 0u64;
while checks.len() < 3 {
let (choice, timed) = selector.next();
if timed {
let was_measuring = selector.decided().is_none();
selector.record(
choice,
if choice == KernelChoice::Parallel {
1
} else {
2
},
);
if was_measuring && selector.decided().is_some() {
calibrations += 1;
checks.push(calls);
}
} else {
selector.advance(1 << 12);
}
calls += 1;
}
assert_eq!(calibrations, 3);
let first = checks[1] - checks[0];
let second = checks[2] - checks[1];
assert!(second > first);
assert_eq!(selector.decided(), Some(KernelChoice::Parallel));
}
#[test]
fn fixed_selections_do_not_measure() {
for (selection, expected) in [
(KernelSelection::Parallel, KernelChoice::Parallel),
(KernelSelection::Sequential, KernelChoice::Sequential),
] {
let mut selector = KernelSelector::new(selection);
assert_eq!(selector.next(), (expected, false));
selector.advance(u64::MAX);
assert_eq!(selector.next(), (expected, false));
}
}
}