use dashmap::DashMap;
use std::{
panic::Location,
sync::Arc,
time::{Duration, Instant},
};
const PREFERRED_SAMPLE_INTERVAL: u32 = 10;
const RESAMPLE_INTERVAL: u32 = 100;
const MAX_RESAMPLE_SHIFT: u32 = 5;
const SERIAL_SAMPLE_BUDGET_NS: u64 = 10_000_000;
const EWMA_PREVIOUS_WEIGHT: u64 = 4;
const EWMA_NEXT_WEIGHT: u64 = 1;
const EWMA_WEIGHT: u64 = EWMA_PREVIOUS_WEIGHT + EWMA_NEXT_WEIGHT;
type Entries = DashMap<Key, Entry>;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum Execution {
Serial,
Parallel,
}
#[derive(Clone, Debug, Default)]
pub(super) struct Policy {
entries: Arc<Entries>,
}
impl Policy {
pub(super) fn try_run<R, E>(
&self,
caller: &'static Location<'static>,
len: usize,
work: usize,
parallelism: usize,
run: impl FnOnce(Execution) -> Result<R, E>,
) -> Result<R, E> {
if parallelism <= 1 {
return run(Execution::Serial);
}
let key = Key::new(caller, len, work, parallelism);
let (execution, measure) = self.entries.entry(key).or_default().choose(parallelism);
let start = measure.then(Instant::now);
let result = run(execution);
if let (Some(start), Ok(_)) = (start, &result) {
let mut entry = self.entries.entry(key).or_default();
entry.record(execution, start.elapsed());
}
result
}
#[cfg(test)]
pub(super) fn len(&self) -> usize {
self.entries.len()
}
#[cfg(test)]
pub(super) fn get_entry(
&self,
caller: &'static Location<'static>,
len: usize,
work: usize,
parallelism: usize,
) -> Option<(Option<u64>, Option<u64>)> {
let key = Key::new(caller, len, work, parallelism);
self.entries.get(&key).map(|e| (e.serial_ns, e.parallel_ns))
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
struct Key {
file: &'static str,
line: u32,
column: u32,
len_bucket: u8,
work_bucket: u8,
parallelism: usize,
}
impl Key {
const fn new(
caller: &'static Location<'static>,
len: usize,
work: usize,
parallelism: usize,
) -> Self {
Self {
file: caller.file(),
line: caller.line(),
column: caller.column(),
len_bucket: len_bucket(len),
work_bucket: len_bucket(work),
parallelism,
}
}
}
#[derive(Clone, Copy, Debug, Default)]
struct Entry {
serial_ns: Option<u64>,
parallel_ns: Option<u64>,
since_probe: u32,
}
impl Entry {
fn projected_serial(parallel_ns: u64, parallelism: usize) -> u64 {
parallel_ns.saturating_mul(u64::try_from(parallelism).unwrap_or(u64::MAX))
}
fn preferred(serial_ns: u64, parallel_ns: u64, parallelism: usize) -> Execution {
if Self::projected_serial(parallel_ns, parallelism) >= SERIAL_SAMPLE_BUDGET_NS
|| serial_ns >= SERIAL_SAMPLE_BUDGET_NS
|| parallel_ns < serial_ns
{
Execution::Parallel
} else {
Execution::Serial
}
}
fn choose(&mut self, parallelism: usize) -> (Execution, bool) {
let Some(parallel_ns) = self.parallel_ns else {
self.since_probe = u32::MAX;
return (Execution::Parallel, true);
};
let can_sample_serial =
Self::projected_serial(parallel_ns, parallelism) < SERIAL_SAMPLE_BUDGET_NS;
let (preferred, interval) =
self.serial_ns
.map_or((Execution::Parallel, RESAMPLE_INTERVAL), |serial_ns| {
let preferred = Self::preferred(serial_ns, parallel_ns, parallelism);
let (winner_ns, loser_ns) = match preferred {
Execution::Serial => (serial_ns, parallel_ns),
Execution::Parallel => (parallel_ns, serial_ns),
};
let slowdown = loser_ns / winner_ns.max(1);
let shift = slowdown
.saturating_sub(1)
.min(u64::from(MAX_RESAMPLE_SHIFT)) as u32;
(preferred, RESAMPLE_INTERVAL << shift)
});
self.since_probe = self.since_probe.saturating_add(1);
if self.since_probe >= interval {
self.since_probe = 0;
let probe = match preferred {
Execution::Serial => Execution::Parallel,
Execution::Parallel if can_sample_serial => Execution::Serial,
Execution::Parallel => Execution::Parallel,
};
return (probe, true);
}
(
preferred,
self.since_probe.is_multiple_of(PREFERRED_SAMPLE_INTERVAL),
)
}
fn record(&mut self, execution: Execution, elapsed: Duration) {
let elapsed_ns = u64::try_from(elapsed.as_nanos()).unwrap_or(u64::MAX);
let estimate = match execution {
Execution::Serial => &mut self.serial_ns,
Execution::Parallel => &mut self.parallel_ns,
};
*estimate = Some(estimate.map_or(elapsed_ns, |current| update_ewma(current, elapsed_ns)));
}
}
fn update_ewma(current: u64, next: u64) -> u64 {
let weighted = u128::from(current) * u128::from(EWMA_PREVIOUS_WEIGHT)
+ u128::from(next) * u128::from(EWMA_NEXT_WEIGHT);
(weighted / u128::from(EWMA_WEIGHT))
.try_into()
.unwrap_or(u64::MAX)
}
const fn len_bucket(len: usize) -> u8 {
if len == 0 {
0
} else {
(usize::BITS - len.leading_zeros()) as u8
}
}
#[cfg(test)]
mod tests {
use super::{
Entry, Execution, Policy, MAX_RESAMPLE_SHIFT, PREFERRED_SAMPLE_INTERVAL, RESAMPLE_INTERVAL,
};
use std::{panic::Location, time::Duration};
const PARALLELISM: usize = 4;
fn choose(entry: &mut Entry) -> (Execution, bool) {
entry.choose(PARALLELISM)
}
#[test]
fn starts_parallel_then_seeds_serial_immediately() {
let mut entry = Entry::default();
assert_eq!(choose(&mut entry), (Execution::Parallel, true));
entry.record(Execution::Parallel, Duration::from_micros(100));
assert_eq!(choose(&mut entry), (Execution::Serial, true));
entry.record(Execution::Serial, Duration::from_micros(95));
for i in 1..RESAMPLE_INTERVAL {
assert_eq!(
choose(&mut entry),
(Execution::Serial, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert_eq!(choose(&mut entry), (Execution::Parallel, true));
}
#[test]
fn defers_serial_seed_when_projection_exceeds_budget() {
let mut entry = Entry::default();
assert_eq!(choose(&mut entry), (Execution::Parallel, true));
entry.record(Execution::Parallel, Duration::from_millis(10));
assert_eq!(choose(&mut entry), (Execution::Parallel, true));
assert!(entry.serial_ns.is_none());
for i in 1..RESAMPLE_INTERVAL {
assert_eq!(
choose(&mut entry),
(Execution::Parallel, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
}
#[test]
fn never_seeds_serial_when_projection_exceeds_budget() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(10));
for i in 1..=(2 * RESAMPLE_INTERVAL) {
assert_eq!(
choose(&mut entry),
(Execution::Parallel, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert!(entry.serial_ns.is_none());
}
#[test]
fn never_runs_serial_on_big_work() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(25));
for _ in 0..10_000 {
let (execution, measure) = entry.choose(12);
assert_eq!(execution, Execution::Parallel);
if measure {
entry.record(Execution::Parallel, Duration::from_millis(25));
}
}
assert!(entry.serial_ns.is_none());
}
#[test]
fn big_serial_estimate_biases_parallel() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(30));
entry.record(Execution::Serial, Duration::from_millis(12));
for _ in 0..(2 * RESAMPLE_INTERVAL) {
let (execution, _) = choose(&mut entry);
assert_eq!(execution, Execution::Parallel);
}
}
#[test]
fn projection_gates_preferred_serial() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(25));
entry.record(Execution::Serial, Duration::from_millis(8));
for _ in 0..(2 * RESAMPLE_INTERVAL) {
let (execution, _) = entry.choose(12);
assert_eq!(execution, Execution::Parallel);
}
}
#[test]
fn prefers_serial_when_faster() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(100));
entry.record(Execution::Serial, Duration::from_micros(95));
assert_eq!(choose(&mut entry), (Execution::Serial, false));
}
#[test]
fn prefers_parallel_when_it_wins_wall_time() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(100));
entry.record(Execution::Serial, Duration::from_micros(200));
assert_eq!(choose(&mut entry), (Execution::Parallel, false));
}
#[test]
fn prefers_serial_on_tie() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(100));
entry.record(Execution::Serial, Duration::from_micros(100));
assert_eq!(choose(&mut entry), (Execution::Serial, false));
}
#[test]
fn pre_seed_parallel_samples_blend() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(10));
entry.record(Execution::Parallel, Duration::from_millis(20));
assert_eq!(entry.parallel_ns, Some(12_000_000));
}
#[test]
fn blends_preferred_samples_with_integer_math() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_nanos(1000));
entry.record(Execution::Serial, Duration::from_nanos(100));
entry.record(Execution::Serial, Duration::from_nanos(200));
assert_eq!(entry.serial_ns, Some(120));
}
#[test]
fn blends_preferred_parallel_samples() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_nanos(100));
entry.record(Execution::Serial, Duration::from_nanos(1000));
entry.record(Execution::Parallel, Duration::from_nanos(200));
assert_eq!(entry.parallel_ns, Some(120));
}
#[test]
fn probes_blend_into_stale_estimates() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(25));
entry.record(Execution::Serial, Duration::from_millis(100));
entry.record(Execution::Serial, Duration::from_millis(5));
assert_eq!(entry.serial_ns, Some(81_000_000));
assert_eq!(
Entry::preferred(
entry.serial_ns.unwrap(),
entry.parallel_ns.unwrap(),
PARALLELISM
),
Execution::Parallel
);
}
#[test]
fn seeds_serial_once_projection_shrinks_into_budget() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(10));
for _ in 1..RESAMPLE_INTERVAL {
let (execution, measure) = choose(&mut entry);
assert_eq!(execution, Execution::Parallel);
if measure {
entry.record(Execution::Parallel, Duration::from_micros(100));
}
}
assert_eq!(choose(&mut entry), (Execution::Serial, true));
}
#[test]
fn resamples_other_execution() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(100));
entry.record(Execution::Serial, Duration::from_micros(80));
for i in 1..RESAMPLE_INTERVAL {
assert_eq!(
choose(&mut entry),
(Execution::Serial, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert_eq!(choose(&mut entry), (Execution::Parallel, true));
}
#[test]
fn resamples_serial_when_parallel_wins() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(100));
entry.record(Execution::Serial, Duration::from_micros(150));
for i in 1..RESAMPLE_INTERVAL {
assert_eq!(
choose(&mut entry),
(Execution::Parallel, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert_eq!(choose(&mut entry), (Execution::Serial, true));
}
#[test]
fn resample_interval_doubles_per_slowdown_multiple() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(250));
entry.record(Execution::Serial, Duration::from_micros(100));
for i in 1..(2 * RESAMPLE_INTERVAL) {
assert_eq!(
choose(&mut entry),
(Execution::Serial, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert_eq!(choose(&mut entry), (Execution::Parallel, true));
}
#[test]
fn resample_interval_is_capped() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(100));
entry.record(Execution::Serial, Duration::from_micros(900));
let interval = RESAMPLE_INTERVAL << MAX_RESAMPLE_SHIFT;
for i in 1..interval {
assert_eq!(
choose(&mut entry),
(Execution::Parallel, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert_eq!(choose(&mut entry), (Execution::Serial, true));
}
#[test]
fn recovers_from_poisoned_estimate() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_millis(2));
entry.record(Execution::Serial, Duration::from_millis(1));
assert_eq!(
Entry::preferred(
entry.serial_ns.unwrap(),
entry.parallel_ns.unwrap(),
PARALLELISM
),
Execution::Serial
);
let mut probes = 0;
let mut flipped_at = None;
for i in 1..=1_000 {
if Entry::preferred(
entry.serial_ns.unwrap(),
entry.parallel_ns.unwrap(),
PARALLELISM,
) == Execution::Parallel
{
flipped_at = Some(i - 1);
break;
}
let (execution, measure) = choose(&mut entry);
match execution {
Execution::Parallel => {
assert!(measure);
probes += 1;
entry.record(Execution::Parallel, Duration::from_micros(500));
}
Execution::Serial => {
if measure {
entry.record(Execution::Serial, Duration::from_millis(1));
}
}
}
}
assert_eq!(probes, 5);
assert_eq!(flipped_at, Some(600));
}
#[test]
fn seed_offered_once_per_interval() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(250));
for round in 0..2 {
for i in 1..RESAMPLE_INTERVAL {
assert_eq!(
choose(&mut entry),
(Execution::Parallel, i % PREFERRED_SAMPLE_INTERVAL == 0),
"round {round}"
);
}
assert_eq!(choose(&mut entry), (Execution::Serial, true));
}
entry.record(Execution::Serial, Duration::from_millis(1));
assert_eq!(choose(&mut entry).0, Execution::Parallel);
}
#[test]
fn probes_big_serial_when_projection_is_affordable() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(500));
entry.record(Execution::Serial, Duration::from_millis(15));
assert_eq!(
Entry::preferred(
entry.serial_ns.unwrap(),
entry.parallel_ns.unwrap(),
PARALLELISM
),
Execution::Parallel
);
let interval = RESAMPLE_INTERVAL << MAX_RESAMPLE_SHIFT;
for i in 1..interval {
assert_eq!(
choose(&mut entry),
(Execution::Parallel, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert_eq!(choose(&mut entry), (Execution::Serial, true));
entry.record(Execution::Serial, Duration::from_micros(300));
assert_eq!(entry.serial_ns, Some(12_060_000));
}
#[test]
fn poisoned_probe_cannot_flip_preference() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(800));
entry.record(Execution::Serial, Duration::from_millis(3));
assert_eq!(
Entry::preferred(
entry.serial_ns.unwrap(),
entry.parallel_ns.unwrap(),
PARALLELISM
),
Execution::Parallel
);
entry.record(Execution::Serial, Duration::from_micros(20));
assert_eq!(entry.serial_ns, Some(2_404_000));
assert_eq!(
Entry::preferred(
entry.serial_ns.unwrap(),
entry.parallel_ns.unwrap(),
PARALLELISM
),
Execution::Parallel
);
}
#[test]
fn refreshes_preferred_parallel_sample() {
let mut entry = Entry::default();
entry.record(Execution::Parallel, Duration::from_micros(100));
entry.record(Execution::Serial, Duration::from_micros(410));
for i in 1..PREFERRED_SAMPLE_INTERVAL {
assert_eq!(
choose(&mut entry),
(Execution::Parallel, i % PREFERRED_SAMPLE_INTERVAL == 0)
);
}
assert_eq!(choose(&mut entry), (Execution::Parallel, true));
}
#[test]
fn try_run_records_success_not_errors() {
let policy = Policy::default();
let location = Location::caller();
let len = 10;
let work = 10;
let result: Result<(), ()> = policy.try_run(location, len, work, PARALLELISM, |_| Err(()));
assert!(result.is_err());
let (serial_ns, parallel_ns) = policy.get_entry(location, len, work, PARALLELISM).unwrap();
assert!(serial_ns.is_none() && parallel_ns.is_none());
let result: Result<(), ()> = policy.try_run(location, len, work, PARALLELISM, |_| Ok(()));
assert!(result.is_ok());
let (serial_ns, parallel_estimate) =
policy.get_entry(location, len, work, PARALLELISM).unwrap();
assert!(parallel_estimate.is_some());
assert!(serial_ns.is_none());
for _ in 0..20 {
let _: Result<(), ()> =
policy.try_run(
location,
len,
work,
PARALLELISM,
|execution| match execution {
Execution::Parallel => Err(()),
Execution::Serial => Ok(()),
},
);
}
let (_, parallel_ns) = policy.get_entry(location, len, work, PARALLELISM).unwrap();
assert_eq!(parallel_ns, parallel_estimate);
}
}