use std::time::Duration;
#[derive(Debug, Clone)]
pub enum LatencyProfile {
Constant(Duration),
LinearRamp {
start: Duration,
step: Duration,
},
StepSchedule(Vec<(usize, Duration)>),
}
pub struct LatencyInjector {
profile: LatencyProfile,
}
impl LatencyInjector {
pub fn new(profile: LatencyProfile) -> Self {
Self { profile }
}
pub fn delay_for(&self, attempt: usize) -> Duration {
match &self.profile {
LatencyProfile::Constant(d) => *d,
LatencyProfile::LinearRamp { start, step } => {
let n = attempt.saturating_sub(1) as u32;
*start + step.saturating_mul(n)
}
LatencyProfile::StepSchedule(boundaries) => {
if boundaries.is_empty() {
return Duration::ZERO;
}
for (threshold, delay) in boundaries.iter() {
if attempt <= *threshold {
return *delay;
}
}
boundaries.last().unwrap().1
}
}
}
pub fn apply_blocking(&self, attempt: usize) {
std::thread::sleep(self.delay_for(attempt));
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn constant_profile_returns_same_duration() {
let inj = LatencyInjector::new(LatencyProfile::Constant(Duration::from_micros(50)));
for attempt in 1..=10 {
assert_eq!(inj.delay_for(attempt), Duration::from_micros(50));
}
}
#[test]
fn linear_ramp_increases() {
let inj = LatencyInjector::new(LatencyProfile::LinearRamp {
start: Duration::from_micros(10),
step: Duration::from_micros(5),
});
assert_eq!(inj.delay_for(1), Duration::from_micros(10));
assert_eq!(inj.delay_for(2), Duration::from_micros(15));
assert_eq!(inj.delay_for(5), Duration::from_micros(30));
}
#[test]
fn step_schedule_picks_correct_band() {
let inj = LatencyInjector::new(LatencyProfile::StepSchedule(vec![
(10, Duration::from_micros(1)),
(20, Duration::from_micros(5)),
(50, Duration::from_micros(20)),
]));
assert_eq!(inj.delay_for(1), Duration::from_micros(1));
assert_eq!(inj.delay_for(10), Duration::from_micros(1));
assert_eq!(inj.delay_for(11), Duration::from_micros(5));
assert_eq!(inj.delay_for(20), Duration::from_micros(5));
assert_eq!(inj.delay_for(21), Duration::from_micros(20));
assert_eq!(inj.delay_for(100), Duration::from_micros(20));
}
#[test]
fn empty_step_schedule_yields_zero() {
let inj = LatencyInjector::new(LatencyProfile::StepSchedule(vec![]));
assert_eq!(inj.delay_for(1), Duration::ZERO);
}
#[test]
fn apply_blocking_sleeps_at_least_the_delay() {
let inj = LatencyInjector::new(LatencyProfile::Constant(Duration::from_millis(10)));
let start = std::time::Instant::now();
inj.apply_blocking(1);
assert!(start.elapsed() >= Duration::from_millis(10));
}
}