use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RateTracker {
#[serde(skip, default = "default_timestamps")]
timestamps: [u64; 64],
write_idx: u8,
valid_count: u8,
}
fn default_timestamps() -> [u64; 64] {
[0u64; 64]
}
impl Default for RateTracker {
fn default() -> Self {
Self::new()
}
}
impl RateTracker {
pub fn new() -> Self {
Self {
timestamps: [0; 64],
write_idx: 0,
valid_count: 0,
}
}
#[inline]
pub fn record(&mut self, now_ms: u64) -> f64 {
self.timestamps[self.write_idx as usize] = now_ms;
self.write_idx = (self.write_idx + 1) % 64;
if self.valid_count < 64 {
self.valid_count += 1;
}
self.current_rate(now_ms)
}
#[inline]
pub fn current_rate(&self, now_ms: u64) -> f64 {
if self.valid_count == 0 {
return 0.0;
}
let window_ms: u64 = 60_000; let cutoff = now_ms.saturating_sub(window_ms);
let count = self
.timestamps
.iter()
.take(self.valid_count as usize)
.filter(|&&ts| ts > cutoff)
.count();
count as f64 }
#[inline]
pub fn rate_in_window(&self, now_ms: u64, window_ms: u64) -> f64 {
if self.valid_count == 0 || window_ms == 0 {
return 0.0;
}
let cutoff = now_ms.saturating_sub(window_ms);
let count = self
.timestamps
.iter()
.take(self.valid_count as usize)
.filter(|&&ts| ts > cutoff)
.count();
(count as f64 * 60_000.0) / window_ms as f64
}
#[inline]
pub fn request_count(&self) -> u8 {
self.valid_count
}
#[inline]
pub fn last_request_time(&self) -> Option<u64> {
if self.valid_count == 0 {
return None;
}
let last_idx = if self.write_idx == 0 {
(self.valid_count - 1) as usize
} else {
(self.write_idx - 1) as usize
};
Some(self.timestamps[last_idx])
}
pub fn clear(&mut self) {
self.timestamps = [0; 64];
self.write_idx = 0;
self.valid_count = 0;
}
#[inline]
pub fn is_burst(&self, now_ms: u64, baseline_rate: f64, multiplier: f64) -> bool {
if baseline_rate <= 0.0 {
return false;
}
let current = self.current_rate(now_ms);
current > baseline_rate * multiplier
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_rate_tracker_new() {
let rt = RateTracker::new();
assert_eq!(rt.request_count(), 0);
assert_eq!(rt.current_rate(1000000), 0.0);
}
#[test]
fn test_rate_tracker_record() {
let mut rt = RateTracker::new();
let base_time = 1000000u64;
for i in 0..10 {
rt.record(base_time + i * 100);
}
let rate = rt.current_rate(base_time + 1000);
assert!((rate - 10.0).abs() < 0.1);
}
#[test]
fn test_rate_tracker_sliding_window() {
let mut rt = RateTracker::new();
let base_time = 1000000u64;
for i in 0..5 {
rt.record(base_time + i * 100);
}
let rate = rt.current_rate(base_time + 500);
assert_eq!(rate, 5.0);
let rate_later = rt.current_rate(base_time + 61_000);
assert_eq!(rate_later, 0.0);
}
#[test]
fn test_rate_tracker_buffer_wraparound() {
let mut rt = RateTracker::new();
let base_time = 1000000u64;
for i in 0..100 {
rt.record(base_time + i * 100);
}
assert_eq!(rt.request_count(), 64);
let rate = rt.current_rate(base_time + 10000);
assert!(rate > 0.0);
}
#[test]
fn test_rate_tracker_last_request_time() {
let mut rt = RateTracker::new();
assert!(rt.last_request_time().is_none());
rt.record(1000);
assert_eq!(rt.last_request_time(), Some(1000));
rt.record(2000);
assert_eq!(rt.last_request_time(), Some(2000));
rt.record(3000);
assert_eq!(rt.last_request_time(), Some(3000));
}
#[test]
fn test_rate_tracker_clear() {
let mut rt = RateTracker::new();
for i in 0..10 {
rt.record(1000 + i * 100);
}
assert_eq!(rt.request_count(), 10);
rt.clear();
assert_eq!(rt.request_count(), 0);
assert!(rt.last_request_time().is_none());
}
#[test]
fn test_rate_tracker_is_burst() {
let mut rt = RateTracker::new();
let base_time = 1000000u64;
for i in 0..20 {
rt.record(base_time + i * 100);
}
assert!(rt.is_burst(base_time + 2000, 5.0, 3.0));
assert!(!rt.is_burst(base_time + 2000, 5.0, 5.0));
}
#[test]
fn test_rate_tracker_is_burst_zero_baseline() {
let mut rt = RateTracker::new();
rt.record(1000);
assert!(!rt.is_burst(1000, 0.0, 2.0));
}
#[test]
fn test_rate_in_window() {
let mut rt = RateTracker::new();
let base_time = 1000000u64;
for i in 0..10 {
rt.record(base_time + i * 1000);
}
let rate = rt.rate_in_window(base_time + 10000, 10_000);
assert!((rate - 54.0).abs() < 1.0);
let rate_5s = rt.rate_in_window(base_time + 10000, 5_000);
assert!(rate_5s > 0.0);
}
#[test]
fn test_rate_in_window_zero() {
let rt = RateTracker::new();
assert_eq!(rt.rate_in_window(1000, 0), 0.0);
assert_eq!(rt.rate_in_window(1000, 10000), 0.0);
}
#[test]
fn test_rate_tracker_timestamps_outside_window() {
let mut rt = RateTracker::new();
for i in 0..5 {
rt.record(i * 100);
}
let rate = rt.current_rate(120_000); assert_eq!(rate, 0.0);
}
#[test]
fn test_rate_tracker_exact_boundary() {
let mut rt = RateTracker::new();
rt.record(0);
let rate = rt.current_rate(60_000);
assert_eq!(rate, 0.0);
let rate_earlier = rt.current_rate(59_999);
assert_eq!(rate_earlier, 0.0);
rt.record(1); let rate_with_in_window = rt.current_rate(60_000);
assert_eq!(rate_with_in_window, 1.0);
}
}