grishka/lemmy
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

Shrink capacity in RateLimitStorage::remove_older_than (#3536)

Browse source 
* Shrink capacity in `RateLimitStorage::remove_older_than`

* Update rate_limiter.rs

* rerun ci

* rerun ci

* rerun ci

* Update rate_limiter.rs
This commit is contained in:
dullbananas 2023-07-10 13:52:37 -07:00 committed by GitHub
parent ef9dc5d0b6
commit 73492af4b0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 16 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
crates/utils/src/rate_limit/rate_limiter.rs
View file

@ -2,6 +2,7 @@ use enum_map::{enum_map, EnumMap};
use once_cell::sync::Lazy; use once_cell::sync::Lazy;
use std::{ use std::{
collections::HashMap, collections::HashMap,
hash::Hash,
net::{IpAddr, Ipv4Addr, Ipv6Addr}, net::{IpAddr, Ipv4Addr, Ipv6Addr},
time::{Duration, Instant}, time::{Duration, Instant},
}; };
@ -206,13 +207,13 @@ impl RateLimitStorage {
.all(|bucket| bucket.last_checked.to_instant() > instant) .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| { retain_and_shrink(&mut self.ipv6_buckets, |_, group_48| {
group_48.children.retain(|_, group_56| { retain_and_shrink(&mut group_48.children, |_, group_56| {
group_56 retain_and_shrink(&mut group_56.children, |_, group_64| {
.children is_recently_used(group_64)
.retain(|_, group_64| is_recently_used(group_64)); });
!group_56.children.is_empty() !group_56.children.is_empty()
}); });
!group_48.children.is_empty() !group_48.children.is_empty()
@ -220,6 +221,15 @@ impl RateLimitStorage {
} }
} }
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) { fn split_ipv6(ip: Ipv6Addr) -> ([u8; 6], u8, u8) {
let [a0, a1, a2, a3, a4, a5, b, c, ..] = ip.octets(); let [a0, a1, a2, a3, a4, a5, b, c, ..] = ip.octets();
([a0, a1, a2, a3, a4, a5], b, c) ([a0, a1, a2, a3, a4, a5], b, c)