Skip to content

Commit

Permalink
Shrink capacity in RateLimitStorage::remove_older_than (#3536)
Browse files Browse the repository at this point in the history
* Shrink capacity in `RateLimitStorage::remove_older_than`

* Update rate_limiter.rs

* rerun ci

* rerun ci

* rerun ci

* Update rate_limiter.rs
  • Loading branch information
dullbananas authored and Nutomic committed Jul 21, 2023
1 parent c666d1b commit cabb14c
Showing 1 changed file with 16 additions and 6 deletions.
22 changes: 16 additions & 6 deletions crates/utils/src/rate_limit/rate_limiter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ use enum_map::{enum_map, EnumMap};
use once_cell::sync::Lazy;
use std::{
collections::HashMap,
hash::Hash,
net::{IpAddr, Ipv4Addr, Ipv6Addr},
time::{Duration, Instant},
};
Expand Down Expand Up @@ -206,20 +207,29 @@ impl RateLimitStorage {
.all(|bucket| bucket.last_checked.to_instant() > instant)
};

self.ipv4_buckets.retain(|_, group| is_recently_used(group));
retain_and_shrink(&mut self.ipv4_buckets, |_, group| is_recently_used(group));

self.ipv6_buckets.retain(|_, group_48| {
group_48.children.retain(|_, group_56| {
group_56
.children
.retain(|_, group_64| is_recently_used(group_64));
retain_and_shrink(&mut self.ipv6_buckets, |_, group_48| {
retain_and_shrink(&mut group_48.children, |_, group_56| {
retain_and_shrink(&mut group_56.children, |_, group_64| {
is_recently_used(group_64)
});
!group_56.children.is_empty()
});
!group_48.children.is_empty()
})
}
}

fn retain_and_shrink<K, V, F>(map: &mut HashMap<K, V>, f: F)
where
K: Eq + Hash,
F: FnMut(&K, &mut V) -> bool,
{
map.retain(f);
map.shrink_to_fit();
}

fn split_ipv6(ip: Ipv6Addr) -> ([u8; 6], u8, u8) {
let [a0, a1, a2, a3, a4, a5, b, c, ..] = ip.octets();
([a0, a1, a2, a3, a4, a5], b, c)
Expand Down

0 comments on commit cabb14c

Please sign in to comment.