From 99d01e186ada2bb3f5cf3fdc447e0cbc5dfae089 Mon Sep 17 00:00:00 2001 From: Sander Saarend Date: Mon, 25 Mar 2024 13:56:03 +0200 Subject: [PATCH] Fix rate limiter (#4560) --- crates/utils/src/rate_limit/rate_limiter.rs | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/crates/utils/src/rate_limit/rate_limiter.rs b/crates/utils/src/rate_limit/rate_limiter.rs index d452fcbb3..66d830c7e 100644 --- a/crates/utils/src/rate_limit/rate_limiter.rs +++ b/crates/utils/src/rate_limit/rate_limiter.rs @@ -158,7 +158,7 @@ impl MapLevel for Map { // Evaluated if `some_children_remaining` is false let total_has_refill_in_future = || { - group.total.into_iter().all(|(action_type, bucket)| { + group.total.into_iter().any(|(action_type, bucket)| { #[allow(clippy::indexing_slicing)] let config = configs[action_type]; bucket.update(now, config).tokens != config.capacity @@ -416,5 +416,23 @@ mod tests { rate_limiter.remove_full_buckets(now); assert!(rate_limiter.ipv4_buckets.is_empty()); assert!(rate_limiter.ipv6_buckets.is_empty()); + + // `remove full buckets` should not remove empty buckets + let ip = "1.1.1.1".parse().unwrap(); + // empty the bucket with 2 requests + assert!(rate_limiter.check(ActionType::Post, ip, now)); + assert!(rate_limiter.check(ActionType::Post, ip, now)); + + rate_limiter.remove_full_buckets(now); + assert!(!rate_limiter.ipv4_buckets.is_empty()); + + // `remove full buckets` should not remove partial buckets + now.secs += 2; + let ip = "1.1.1.1".parse().unwrap(); + // Only make one request, so bucket still has 1 token + assert!(rate_limiter.check(ActionType::Post, ip, now)); + + rate_limiter.remove_full_buckets(now); + assert!(!rate_limiter.ipv4_buckets.is_empty()); } }