瀏覽代碼

feat: add rpm mode schedules

main
orangels 1 周之前
父節點
當前提交
7c0d1e8af0
共有 2 個文件被更改,包括 231 次插入0 次删除
  1. +1
    -0
      src/main.rs
  2. +230
    -0
      src/rpm_modes.rs

+ 1
- 0
src/main.rs 查看文件

@@ -4,6 +4,7 @@ mod config;
mod metrics; mod metrics;
mod protocols; mod protocols;
mod report; mod report;
mod rpm_modes;
mod runner; mod runner;


use anyhow::Result; use anyhow::Result;


+ 230
- 0
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<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));
}
}

Loading…
取消
儲存