diff --git a/src/main.rs b/src/main.rs index 19ecb40..1735c67 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,7 @@ mod config; mod metrics; mod protocols; mod report; +mod rpm_modes; mod runner; use anyhow::Result; diff --git a/src/rpm_modes.rs b/src/rpm_modes.rs new file mode 100644 index 0000000..0ede11b --- /dev/null +++ b/src/rpm_modes.rs @@ -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, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WindowBoundaryPlan { + pub boundary_offset: Duration, + pub probes: Vec, +} + +pub fn sustained_schedule(duration: Duration, rpm: u32) -> Vec { + 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 { + (0..burst).map(|_| Duration::ZERO).collect() +} + +pub fn token_bucket_schedule(rpm: u32, burst: u32, probe_seconds: u64) -> Vec { + 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 { + 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, 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 { + (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::>(); + + 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::>(); + + 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)); + } +}