|
|
|
@@ -0,0 +1,230 @@ |
|
|
|
use chrono::{DateTime, Timelike, Utc}; |
|
|
|
use std::time::Duration; |
|
|
|
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, clap::ValueEnum)] |
|
|
|
#[value(rename_all = "kebab-case")] |
|
|
|
pub enum RpmMode { |
|
|
|
Sustained, |
|
|
|
Burst, |
|
|
|
TokenBucket, |
|
|
|
SlidingWindow, |
|
|
|
WindowBoundary, |
|
|
|
Diagnose, |
|
|
|
} |
|
|
|
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)] |
|
|
|
pub enum ProbePhase { |
|
|
|
Burst, |
|
|
|
RefillProbe, |
|
|
|
SlidingProbe, |
|
|
|
BeforeBoundary, |
|
|
|
AfterBoundary, |
|
|
|
} |
|
|
|
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)] |
|
|
|
pub struct ScheduledProbe { |
|
|
|
pub offset: Duration, |
|
|
|
pub phase: ProbePhase, |
|
|
|
pub second: Option<u64>, |
|
|
|
} |
|
|
|
|
|
|
|
#[derive(Debug, Clone, PartialEq, Eq)] |
|
|
|
pub struct WindowBoundaryPlan { |
|
|
|
pub boundary_offset: Duration, |
|
|
|
pub probes: Vec<ScheduledProbe>, |
|
|
|
} |
|
|
|
|
|
|
|
pub fn sustained_schedule(duration: Duration, rpm: u32) -> Vec<Duration> { |
|
|
|
if rpm == 0 || duration.is_zero() { |
|
|
|
return Vec::new(); |
|
|
|
} |
|
|
|
|
|
|
|
let interval = rpm_interval(rpm); |
|
|
|
let mut offset = Duration::ZERO; |
|
|
|
let mut starts = Vec::new(); |
|
|
|
|
|
|
|
while offset < duration { |
|
|
|
starts.push(offset); |
|
|
|
offset += interval; |
|
|
|
} |
|
|
|
|
|
|
|
starts |
|
|
|
} |
|
|
|
|
|
|
|
pub fn burst_schedule(burst: u32) -> Vec<Duration> { |
|
|
|
(0..burst).map(|_| Duration::ZERO).collect() |
|
|
|
} |
|
|
|
|
|
|
|
pub fn token_bucket_schedule(rpm: u32, burst: u32, probe_seconds: u64) -> Vec<ScheduledProbe> { |
|
|
|
let mut probes = burst_probes(burst, ProbePhase::Burst); |
|
|
|
if rpm == 0 || probe_seconds == 0 { |
|
|
|
return probes; |
|
|
|
} |
|
|
|
|
|
|
|
let interval = rpm_interval(rpm); |
|
|
|
let end = Duration::from_secs(probe_seconds); |
|
|
|
let mut offset = interval; |
|
|
|
|
|
|
|
while offset <= end { |
|
|
|
probes.push(ScheduledProbe { |
|
|
|
offset, |
|
|
|
phase: ProbePhase::RefillProbe, |
|
|
|
second: Some(offset.as_secs().max(1)), |
|
|
|
}); |
|
|
|
offset += interval; |
|
|
|
} |
|
|
|
|
|
|
|
probes |
|
|
|
} |
|
|
|
|
|
|
|
pub fn sliding_window_schedule(burst: u32, probe_seconds: u64) -> Vec<ScheduledProbe> { |
|
|
|
let mut probes = burst_probes(burst, ProbePhase::Burst); |
|
|
|
|
|
|
|
for second in 1..=probe_seconds { |
|
|
|
probes.push(ScheduledProbe { |
|
|
|
offset: Duration::from_secs(second), |
|
|
|
phase: ProbePhase::SlidingProbe, |
|
|
|
second: Some(second), |
|
|
|
}); |
|
|
|
} |
|
|
|
|
|
|
|
probes |
|
|
|
} |
|
|
|
|
|
|
|
pub fn window_boundary_plan(now: DateTime<Utc>, burst: u32, offset_ms: u64) -> WindowBoundaryPlan { |
|
|
|
let millis_into_minute = |
|
|
|
u64::from(now.second()) * 1_000 + u64::from(now.nanosecond()) / 1_000_000; |
|
|
|
let minute_ms = 60_000; |
|
|
|
let until_next_boundary_ms = if millis_into_minute == 0 { |
|
|
|
minute_ms |
|
|
|
} else { |
|
|
|
minute_ms - millis_into_minute |
|
|
|
}; |
|
|
|
let boundary_offset = Duration::from_millis(until_next_boundary_ms); |
|
|
|
let offset = Duration::from_millis(offset_ms); |
|
|
|
let before_offset = boundary_offset.saturating_sub(offset); |
|
|
|
let after_offset = boundary_offset + offset; |
|
|
|
let mut probes = Vec::with_capacity((burst as usize).saturating_mul(2)); |
|
|
|
|
|
|
|
for _ in 0..burst { |
|
|
|
probes.push(ScheduledProbe { |
|
|
|
offset: before_offset, |
|
|
|
phase: ProbePhase::BeforeBoundary, |
|
|
|
second: None, |
|
|
|
}); |
|
|
|
} |
|
|
|
for _ in 0..burst { |
|
|
|
probes.push(ScheduledProbe { |
|
|
|
offset: after_offset, |
|
|
|
phase: ProbePhase::AfterBoundary, |
|
|
|
second: None, |
|
|
|
}); |
|
|
|
} |
|
|
|
|
|
|
|
WindowBoundaryPlan { |
|
|
|
boundary_offset, |
|
|
|
probes, |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
fn burst_probes(burst: u32, phase: ProbePhase) -> Vec<ScheduledProbe> { |
|
|
|
(0..burst) |
|
|
|
.map(|_| ScheduledProbe { |
|
|
|
offset: Duration::ZERO, |
|
|
|
phase, |
|
|
|
second: Some(0), |
|
|
|
}) |
|
|
|
.collect() |
|
|
|
} |
|
|
|
|
|
|
|
fn rpm_interval(rpm: u32) -> Duration { |
|
|
|
Duration::from_secs_f64(60.0 / f64::from(rpm)) |
|
|
|
} |
|
|
|
|
|
|
|
#[cfg(test)] |
|
|
|
mod tests { |
|
|
|
use super::*; |
|
|
|
use chrono::TimeZone; |
|
|
|
|
|
|
|
#[test] |
|
|
|
fn sustained_120_rpm_starts_every_500ms() { |
|
|
|
let starts = sustained_schedule(Duration::from_millis(1_100), 120); |
|
|
|
|
|
|
|
assert_eq!( |
|
|
|
starts, |
|
|
|
vec![ |
|
|
|
Duration::from_millis(0), |
|
|
|
Duration::from_millis(500), |
|
|
|
Duration::from_millis(1_000) |
|
|
|
] |
|
|
|
); |
|
|
|
} |
|
|
|
|
|
|
|
#[test] |
|
|
|
fn zero_rpm_sustained_has_no_starts() { |
|
|
|
assert!(sustained_schedule(Duration::from_secs(10), 0).is_empty()); |
|
|
|
} |
|
|
|
|
|
|
|
#[test] |
|
|
|
fn burst_five_starts_at_zero() { |
|
|
|
assert_eq!(burst_schedule(5), vec![Duration::ZERO; 5]); |
|
|
|
} |
|
|
|
|
|
|
|
#[test] |
|
|
|
fn token_bucket_120_rpm_has_burst_and_refill_probes() { |
|
|
|
let probes = token_bucket_schedule(120, 120, 2); |
|
|
|
let burst_count = probes |
|
|
|
.iter() |
|
|
|
.filter(|probe| probe.phase == ProbePhase::Burst) |
|
|
|
.count(); |
|
|
|
let refill_offsets = probes |
|
|
|
.iter() |
|
|
|
.filter(|probe| probe.phase == ProbePhase::RefillProbe) |
|
|
|
.map(|probe| probe.offset) |
|
|
|
.collect::<Vec<_>>(); |
|
|
|
|
|
|
|
assert_eq!(burst_count, 120); |
|
|
|
assert_eq!( |
|
|
|
refill_offsets, |
|
|
|
vec![ |
|
|
|
Duration::from_millis(500), |
|
|
|
Duration::from_millis(1_000), |
|
|
|
Duration::from_millis(1_500), |
|
|
|
Duration::from_millis(2_000), |
|
|
|
] |
|
|
|
); |
|
|
|
} |
|
|
|
|
|
|
|
#[test] |
|
|
|
fn sliding_window_probe_lasts_requested_seconds() { |
|
|
|
let probes = sliding_window_schedule(3, 4); |
|
|
|
let sliding_offsets = probes |
|
|
|
.iter() |
|
|
|
.filter(|probe| probe.phase == ProbePhase::SlidingProbe) |
|
|
|
.map(|probe| probe.offset) |
|
|
|
.collect::<Vec<_>>(); |
|
|
|
|
|
|
|
assert_eq!( |
|
|
|
sliding_offsets, |
|
|
|
vec![ |
|
|
|
Duration::from_secs(1), |
|
|
|
Duration::from_secs(2), |
|
|
|
Duration::from_secs(3), |
|
|
|
Duration::from_secs(4), |
|
|
|
] |
|
|
|
); |
|
|
|
} |
|
|
|
|
|
|
|
#[test] |
|
|
|
fn window_boundary_plan_places_batches_around_next_minute() { |
|
|
|
let now = Utc.with_ymd_and_hms(2026, 5, 6, 12, 34, 50).unwrap(); |
|
|
|
let plan = window_boundary_plan(now, 2, 500); |
|
|
|
|
|
|
|
assert_eq!(plan.boundary_offset, Duration::from_secs(10)); |
|
|
|
assert_eq!(plan.probes.len(), 4); |
|
|
|
assert_eq!(plan.probes[0].phase, ProbePhase::BeforeBoundary); |
|
|
|
assert_eq!(plan.probes[0].offset, Duration::from_millis(9_500)); |
|
|
|
assert_eq!(plan.probes[2].phase, ProbePhase::AfterBoundary); |
|
|
|
assert_eq!(plan.probes[2].offset, Duration::from_millis(10_500)); |
|
|
|
} |
|
|
|
} |