From 73492af4b09448684ffde3d55454434ec3ed490b Mon Sep 17 00:00:00 2001
From: dullbananas <dull.bananas0@gmail.com>
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<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)