use std::time::{Duration, Instant};
use crate::{
HardLimitFactor, LocalRateLimiterOptions, RateGroupSizeMs, RateLimit, RateLimitDecision,
SuppressedLocalRateLimiter, SuppressionFactorCacheMs, WindowSizeSeconds,
};
fn limiter(
window_size_seconds: u64,
rate_group_size_ms: u64,
hard_limit_factor: f64,
) -> SuppressedLocalRateLimiter {
SuppressedLocalRateLimiter::new(LocalRateLimiterOptions {
window_size_seconds: WindowSizeSeconds::try_from(window_size_seconds).unwrap(),
rate_group_size_ms: RateGroupSizeMs::try_from(rate_group_size_ms).unwrap(),
hard_limit_factor: HardLimitFactor::try_from(hard_limit_factor).unwrap(),
suppression_factor_cache_ms: SuppressionFactorCacheMs::default(),
})
}
fn limiter_with_cache_ms(
window_size_seconds: u64,
rate_group_size_ms: u64,
hard_limit_factor: f64,
suppression_factor_cache_ms: u64,
) -> SuppressedLocalRateLimiter {
SuppressedLocalRateLimiter::new(LocalRateLimiterOptions {
window_size_seconds: WindowSizeSeconds::try_from(window_size_seconds).unwrap(),
rate_group_size_ms: RateGroupSizeMs::try_from(rate_group_size_ms).unwrap(),
hard_limit_factor: HardLimitFactor::try_from(hard_limit_factor).unwrap(),
suppression_factor_cache_ms: SuppressionFactorCacheMs::try_from(
suppression_factor_cache_ms,
)
.unwrap(),
})
}
fn fill_past_soft_limit(
limiter: &SuppressedLocalRateLimiter,
key: &str,
rate_limit: &RateLimit,
window_size_seconds: u64,
desired_suppression_factor: f64,
) {
let soft_limit = window_size_seconds * **rate_limit as u64;
let _ = limiter.inc(key, rate_limit, soft_limit);
limiter.test_set_suppression_factor(key, Instant::now(), desired_suppression_factor);
}
#[test]
fn does_not_suppress_when_window_limit_not_exceeded() {
let limiter = limiter(10, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
let _ = limiter.inc(key, &rate_limit, 10);
std::thread::sleep(Duration::from_millis(1001));
let _ = limiter.inc(key, &rate_limit, 20);
let decision = limiter.inc(key, &rate_limit, 1);
assert!(
matches!(
decision,
RateLimitDecision::Suppressed {
suppression_factor,
..
} if suppression_factor < 1e-12
) || matches!(decision, RateLimitDecision::Allowed),
"decision: {:?}",
decision
);
}
#[test]
fn verify_suppression_factor_calculation_spread() {
let limiter = limiter(10, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(1f64).unwrap();
for _ in 0..20 {
let _ = limiter.inc(key, &rate_limit, 1);
std::thread::sleep(Duration::from_millis(3000 / 20));
}
std::thread::sleep(Duration::from_millis(1200));
let expected_suppression_factor = 1f64 - (1f64 / 2.0f64);
limiter.test_set_suppression_factor(key, Instant::now() - Duration::from_millis(101), 0.0);
let suppression_factor = limiter.get_suppression_factor(key);
assert!(
(suppression_factor - expected_suppression_factor).abs() < 1e-12,
"suppression_factor: {suppression_factor:?} expected: {expected_suppression_factor:?}"
);
}
#[test]
fn verify_suppression_factor_calculation_last_second() {
let limiter = limiter(10, 100, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(1f64).unwrap();
let _ = limiter.inc(key, &rate_limit, 10);
std::thread::sleep(Duration::from_millis(1001));
let _ = limiter.inc(key, &rate_limit, 20);
let expected_suppression_factor = 1f64 - (1f64 / 20f64);
limiter.test_set_suppression_factor(key, Instant::now() - Duration::from_millis(101), 0.0);
let suppression_factor = limiter.get_suppression_factor(key);
assert!(
(suppression_factor - expected_suppression_factor).abs() < 1e-12,
"suppression_factor: {suppression_factor:?} expected: {expected_suppression_factor:?}"
);
}
#[test]
fn verify_hard_limit_rejects_local() {
let limiter = limiter(10, 100, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(1f64).unwrap();
fill_past_soft_limit(&limiter, key, &rate_limit, 10, 0.0);
let decision = limiter.inc(key, &rate_limit, 91);
assert!(
matches!(decision, RateLimitDecision::Rejected { .. })
|| matches!(decision, RateLimitDecision::Suppressed { suppression_factor, .. } if suppression_factor == 1.0f64),
"decision: {:?}",
decision
);
}
#[test]
fn suppressed_inc_denied_returns_suppressed_and_does_not_increment_accepted() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
fill_past_soft_limit(&limiter, key, &rate_limit, 1, 0.25);
let mut rng = |p: f64| {
assert!((p - 0.75).abs() < 1e-12, "p: {p:?}");
false
};
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert!(
matches!(decision, RateLimitDecision::Rejected { .. })
|| matches!(
decision,
RateLimitDecision::Suppressed {
suppression_factor,
is_allowed: false
} if (suppression_factor - 0.25).abs() < 1e-12
),
"decision: {decision:?}"
);
}
#[test]
fn suppressed_inc_allowed_returns_suppressed_is_allowed_true_and_increments_accepted() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
fill_past_soft_limit(&limiter, key, &rate_limit, 1, 0.25);
let mut rng = |p: f64| {
assert!((p - 0.75).abs() < 1e-12, "p: {p:?}");
true
};
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert!(
matches!(decision, RateLimitDecision::Rejected { .. })
|| matches!(
decision,
RateLimitDecision::Suppressed {
suppression_factor,
is_allowed: true
} if (suppression_factor - 0.25).abs() < 1e-12
),
"decision: {decision:?}"
);
}
#[test]
fn hard_limit_is_enforced_after_suppression_factor_cache_expires() {
let limiter = limiter_with_cache_ms(10, 1000, 2f64, 1);
let key = "k";
let rate_limit = RateLimit::try_from(1f64).unwrap();
let d1 = limiter.inc(key, &rate_limit, 1);
assert!(matches!(d1, RateLimitDecision::Allowed), "d1: {d1:?}");
std::thread::sleep(Duration::from_millis(5));
let d2 = limiter.inc(key, &rate_limit, 1);
assert!(matches!(d2, RateLimitDecision::Allowed), "d2: {d2:?}");
std::thread::sleep(Duration::from_millis(5));
let d3 = limiter.inc(key, &rate_limit, 20);
assert!(
matches!(d3, RateLimitDecision::Rejected { .. })
|| matches!(
d3,
RateLimitDecision::Suppressed {
suppression_factor,
is_allowed: false
} if (suppression_factor - 1.0).abs() < 1e-12
),
"d3: {d3:?}"
);
assert!((limiter.get_suppression_factor(key) - 1.0).abs() < 1e-12);
}
#[test]
#[should_panic(
expected = "SuppressedLocalRateLimiter::get_suppression_factor: suppression factor > 1"
)]
fn suppression_factor_gt_one_panics() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
fill_past_soft_limit(&limiter, key, &rate_limit, 1, 2f64);
let mut rng = |_p: f64| true;
let _ = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
}
#[test]
#[should_panic(
expected = "SuppressedLocalRateLimiter::get_suppression_factor: negative suppression factor"
)]
fn suppression_factor_negative_panics() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
fill_past_soft_limit(&limiter, key, &rate_limit, 1, -0.01);
let mut rng = |_p: f64| true;
let _ = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
}
#[test]
fn suppression_factor_cache_is_recomputed_after_cache_window() {
let limiter = limiter(1, 50, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
let old_ts = Instant::now() - Duration::from_millis(101);
limiter.test_set_suppression_factor(key, old_ts, 0.9);
let _ = limiter.inc(key, &rate_limit, 1);
limiter.test_set_suppression_factor(key, old_ts, 0.9);
let val = limiter.get_suppression_factor(key);
assert!((val - 0f64).abs() < 1e-12, "val: {val}");
let (_ts, cached_val) = limiter
.test_get_suppression_factor(key)
.expect("expected suppression factor to be persisted");
assert!((cached_val - 0f64).abs() < 1e-12);
}
#[test]
fn cleanup_removes_stale_suppression_factors_and_keeps_fresh() {
let limiter = limiter(1, 1000, 10f64);
limiter.test_set_suppression_factor("stale", Instant::now() - Duration::from_millis(250), 0.5);
limiter.test_set_suppression_factor("fresh", Instant::now(), 0.25);
limiter.cleanup(100);
assert!(
limiter.test_get_suppression_factor("stale").is_none(),
"expected stale suppression factor removed"
);
assert!(
limiter.test_get_suppression_factor("fresh").is_some(),
"expected fresh suppression factor retained"
);
}
#[test]
fn cleanup_does_not_break_next_inc_when_cache_entry_removed() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
limiter.test_set_suppression_factor(key, Instant::now() - Duration::from_millis(250), 0.9);
limiter.cleanup(100);
assert!(limiter.test_get_suppression_factor(key).is_none());
let val = limiter.get_suppression_factor(key);
assert!(matches!(val, v if (v - 0f64).abs() < 1e-12));
let (_ts, cached_val) = limiter
.test_get_suppression_factor(key)
.expect("expected suppression factor persisted after recompute");
assert!((cached_val - 0f64).abs() < 1e-12);
}
#[test]
fn cleanup_is_millisecond_granularity_for_suppression_factor() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
limiter.test_set_suppression_factor(key, Instant::now() - Duration::from_millis(150), 0.5);
limiter.cleanup(100);
assert!(limiter.test_get_suppression_factor(key).is_none());
}
#[test]
fn suppression_factor_zero_returns_allowed_and_rng_not_called() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
limiter.test_set_suppression_factor(key, Instant::now(), 0.0);
assert_eq!(limiter.get_suppression_factor(key), 0.0);
let mut rng = |_p: f64| panic!("rng must not be called when suppression_factor == 0");
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert!(matches!(decision, RateLimitDecision::Allowed));
}
#[test]
fn suppression_factor_one_must_not_return_allowed() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
fill_past_soft_limit(&limiter, key, &rate_limit, 1, 1.0);
let mut rng = |_p: f64| panic!("rng must not be called when suppression_factor == 1");
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert!(
matches!(decision, RateLimitDecision::Rejected { .. })
|| matches!(decision, RateLimitDecision::Suppressed { .. }),
"decision: {decision:?}"
);
}
#[test]
fn suppression_rng_probability_is_one_minus_suppression_factor() {
let limiter = limiter(1, 1000, 10f64);
let key = "k";
let rate_limit = RateLimit::try_from(5f64).unwrap();
fill_past_soft_limit(&limiter, key, &rate_limit, 1, 0.25);
let mut rng = |p: f64| {
assert!((p - 0.75).abs() < 1e-12, "p: {p:?}");
false
};
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert!(
matches!(decision, RateLimitDecision::Rejected { .. })
|| matches!(
decision,
RateLimitDecision::Suppressed {
suppression_factor,
is_allowed: false
} if (suppression_factor - 0.25).abs() < 1e-12
),
"decision: {decision:?}"
);
}
#[test]
fn suppression_property_values_do_not_panic_and_follow_basic_contract() {
let limiter = limiter(1, 1000, 10f64);
let rate_limit = RateLimit::try_from(5f64).unwrap();
let sfs = [0.0, 0.1, 0.25, 0.5, 0.9, 1.0];
for sf in sfs {
let key = format!("k_{sf}");
let key = key.as_str();
if sf == 0.0 {
limiter.test_set_suppression_factor(key, Instant::now(), sf);
let mut rng = |_p: f64| panic!("rng must not be called when suppression_factor == 0");
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert!(
matches!(decision, RateLimitDecision::Allowed),
"sf={sf} decision={decision:?}"
);
continue;
}
if sf == 1.0 {
fill_past_soft_limit(&limiter, key, &rate_limit, 1, sf);
let mut rng = |_p: f64| panic!("rng must not be called when suppression_factor == 1");
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert!(
!matches!(decision, RateLimitDecision::Allowed),
"sf={sf} decision={decision:?}"
);
continue;
}
fill_past_soft_limit(&limiter, key, &rate_limit, 1, sf);
for expected_is_allowed in [false, true] {
limiter.test_set_suppression_factor(key, Instant::now(), sf);
let mut called = 0u64;
let mut rng = |p: f64| {
called += 1;
assert!((p - (1.0 - sf)).abs() < 1e-12, "sf={sf} p={p:?}");
expected_is_allowed
};
let decision = limiter.inc_with_rng(key, &rate_limit, 1, &mut rng);
assert_eq!(called, 1, "sf={sf} expected 1 rng call, got {called}");
match decision {
RateLimitDecision::Rejected { .. } => {}
RateLimitDecision::Suppressed {
suppression_factor,
is_allowed,
} => {
assert!(
(suppression_factor - sf).abs() < 1e-12,
"sf={sf} got={suppression_factor:?}"
);
assert_eq!(is_allowed, expected_is_allowed, "sf={sf}");
}
RateLimitDecision::Allowed => {
panic!("sf={sf} unexpectedly returned Allowed");
}
}
}
}
}
#[test]
fn suppressed_is_deterministically_allowed_until_base_capacity_boundary() {
let window_size_seconds = 10_u64;
let hard_limit_factor = 2f64;
let limiter = limiter(window_size_seconds, 1000, hard_limit_factor);
let key = "k";
let rate_limit = RateLimit::try_from(1f64).unwrap();
let d1 = limiter.inc(key, &rate_limit, 1);
assert!(matches!(d1, RateLimitDecision::Allowed), "d1: {d1:?}");
let d2 = limiter.inc(key, &rate_limit, 1);
assert!(matches!(d2, RateLimitDecision::Allowed), "d2: {d2:?}");
for i in 2..8u64 {
let d = limiter.inc(key, &rate_limit, 1);
assert!(matches!(d, RateLimitDecision::Allowed), "i={i} d={d:?}");
}
let d_boundary = limiter.inc(key, &rate_limit, 1);
assert!(
matches!(d_boundary, RateLimitDecision::Allowed),
"d_boundary: {d_boundary:?}"
);
let d_over = limiter.inc(key, &rate_limit, 1);
assert!(
!matches!(d_over, RateLimitDecision::Allowed),
"d_over should not be Allowed once past soft limit: {d_over:?}"
);
}
#[test]
fn suppressed_is_fully_denied_after_hard_limit_observed() {
let window_size_seconds = 10_u64;
let hard_limit_factor = 2f64;
let limiter = limiter_with_cache_ms(window_size_seconds, 1000, hard_limit_factor, 1);
let key = "k_hard";
let rate_limit = RateLimit::try_from(1f64).unwrap();
for i in 0..9u64 {
std::thread::sleep(Duration::from_millis(2)); let d = limiter.inc(key, &rate_limit, 1);
assert!(
matches!(d, RateLimitDecision::Allowed)
|| matches!(
d,
RateLimitDecision::Suppressed {
is_allowed: true,
..
}
),
"i={i} d={d:?}"
);
}
std::thread::sleep(Duration::from_millis(2));
let d_push = limiter.inc(key, &rate_limit, 6);
assert!(
matches!(
d_push,
RateLimitDecision::Suppressed {
suppression_factor,
is_allowed: false,
} if (suppression_factor - 1.0).abs() < 1e-12
),
"d_push: {d_push:?}"
);
std::thread::sleep(Duration::from_millis(2)); let d_next = limiter.inc(key, &rate_limit, 7);
assert!(
matches!(
d_next,
RateLimitDecision::Suppressed {
suppression_factor,
is_allowed: false,
} if (suppression_factor - 1.0).abs() < 1e-12
),
"d_next: {d_next:?}"
);
}