From 73492af4b09448684ffde3d55454434ec3ed490b Mon Sep 17 00:00:00 2001 From: dullbananas Date: Mon, 10 Jul 2023 13:52:37 -0700 Subject: [PATCH] Shrink capacity in `RateLimitStorage::remove_older_than` (#3536) * Shrink capacity in `RateLimitStorage::remove_older_than` * Update rate_limiter.rs * rerun ci * rerun ci * rerun ci * Update rate_limiter.rs --- crates/utils/src/rate_limit/rate_limiter.rs | 22 +++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/crates/utils/src/rate_limit/rate_limiter.rs b/crates/utils/src/rate_limit/rate_limiter.rs index cea05bc0..ed3dc569 100644 --- a/crates/utils/src/rate_limit/rate_limiter.rs +++ b/crates/utils/src/rate_limit/rate_limiter.rs @@ -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}, }; @@ -206,13 +207,13 @@ 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() @@ -220,6 +221,15 @@ impl RateLimitStorage { } } +fn retain_and_shrink(map: &mut HashMap, 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) -- 2.44.1