1use std::time::Duration;
9
10#[derive(Debug, Clone)]
14pub enum LatencyProfile {
15 Constant(Duration),
17 LinearRamp {
19 start: Duration,
21 step: Duration,
23 },
24 StepSchedule(Vec<(usize, Duration)>),
30}
31
32pub struct LatencyInjector {
50 profile: LatencyProfile,
51}
52
53impl LatencyInjector {
54 pub fn new(profile: LatencyProfile) -> Self {
56 Self { profile }
57 }
58
59 pub fn delay_for(&self, attempt: usize) -> Duration {
61 match &self.profile {
62 LatencyProfile::Constant(d) => *d,
63 LatencyProfile::LinearRamp { start, step } => {
64 let n = attempt.saturating_sub(1) as u32;
65 *start + step.saturating_mul(n)
66 }
67 LatencyProfile::StepSchedule(boundaries) => {
68 if boundaries.is_empty() {
69 return Duration::ZERO;
70 }
71 for (threshold, delay) in boundaries.iter() {
72 if attempt <= *threshold {
73 return *delay;
74 }
75 }
76 boundaries.last().unwrap().1
77 }
78 }
79 }
80
81 pub fn apply_blocking(&self, attempt: usize) {
85 std::thread::sleep(self.delay_for(attempt));
86 }
87}
88
89#[cfg(test)]
90mod tests {
91 use super::*;
92
93 #[test]
94 fn constant_profile_returns_same_duration() {
95 let inj = LatencyInjector::new(LatencyProfile::Constant(Duration::from_micros(50)));
96 for attempt in 1..=10 {
97 assert_eq!(inj.delay_for(attempt), Duration::from_micros(50));
98 }
99 }
100
101 #[test]
102 fn linear_ramp_increases() {
103 let inj = LatencyInjector::new(LatencyProfile::LinearRamp {
104 start: Duration::from_micros(10),
105 step: Duration::from_micros(5),
106 });
107 assert_eq!(inj.delay_for(1), Duration::from_micros(10));
108 assert_eq!(inj.delay_for(2), Duration::from_micros(15));
109 assert_eq!(inj.delay_for(5), Duration::from_micros(30));
110 }
111
112 #[test]
113 fn step_schedule_picks_correct_band() {
114 let inj = LatencyInjector::new(LatencyProfile::StepSchedule(vec![
115 (10, Duration::from_micros(1)),
116 (20, Duration::from_micros(5)),
117 (50, Duration::from_micros(20)),
118 ]));
119 assert_eq!(inj.delay_for(1), Duration::from_micros(1));
120 assert_eq!(inj.delay_for(10), Duration::from_micros(1));
121 assert_eq!(inj.delay_for(11), Duration::from_micros(5));
122 assert_eq!(inj.delay_for(20), Duration::from_micros(5));
123 assert_eq!(inj.delay_for(21), Duration::from_micros(20));
124 assert_eq!(inj.delay_for(100), Duration::from_micros(20));
125 }
126
127 #[test]
128 fn empty_step_schedule_yields_zero() {
129 let inj = LatencyInjector::new(LatencyProfile::StepSchedule(vec![]));
130 assert_eq!(inj.delay_for(1), Duration::ZERO);
131 }
132
133 #[test]
134 fn apply_blocking_sleeps_at_least_the_delay() {
135 let inj = LatencyInjector::new(LatencyProfile::Constant(Duration::from_millis(10)));
136 let start = std::time::Instant::now();
137 inj.apply_blocking(1);
138 assert!(start.elapsed() >= Duration::from_millis(10));
139 }
140}