add and use lazy map cleaner

signal-gateway/src/lazy_map_cleaner.rs

@@ -0,0 +1,64 @@
+use crate::concurrent_map::LazyMap;
+use std::{sync::Mutex};
+
+/// Trait for types that expose a timestamp in seconds.
+/// For containers, this should be the oldest timestamp in the container.
+pub trait TsSecs {
+    fn ts_secs(&self) -> i64;
+}
+
+#[derive(Debug, Default)]
+struct CleanerState {
+    last_cleanup_size: usize,
+    last_checkup: i64,
+}
+
+/// A policy object that can cleanup entries from a LazyMap, if the value type has timestamps.
+/// The lazy map is cleaned whenever it doubles in size, or whenever enough time has passed
+/// since the last cleanup. This can prevent lazy maps from growing without bound.
+pub struct LazyMapCleaner {
+    max_age: i64,
+    min_checkup_period: i64,
+    state: Mutex<CleanerState>,
+}
+
+impl LazyMapCleaner {
+    /// Create a new lazy map cleaner, with given max age of items (in seconds)
+    pub fn new(max_age: i64) -> Self {
+        Self {
+            max_age,
+            min_checkup_period:  3600,
+            state: Default::default(),
+        }
+    }
+    
+    /// Set the min checkup period (in seconds)
+    /// Defaults to 1 hour if unset
+    #[allow(dead_code)]
+    pub fn with_min_checkup_period(self, min_checkup_period: i64) -> Self {
+        Self {
+            max_age: self.max_age,
+            min_checkup_period,
+            state: self.state,
+        }
+    }
+
+    /// Maybe clean a lazy map, if we have met conditions to do so
+    pub fn maybe_clean<K, V>(&self, now: i64, map: &LazyMap<K, V>)
+        where K: Eq + std::hash::Hash + Clone,
+              V: TsSecs
+    {
+        // Only lock if we can do so without contention, if someone else is checking to clean,
+        // we don't have to do so ourselves.
+        if let Ok(mut lk) = self.state.try_lock() {
+            // If the last cleanup size is 0, we'll still treat it as 1 and only look to clean if we reach 2.
+            if map.len() >= lk.last_cleanup_size.max(1) * 2
+            || lk.last_checkup + self.min_checkup_period < now {
+                let cutoff = now - self.max_age;
+                let new_cleanup_size = map.retain(|_key, val| val.ts_secs() >= cutoff);
+                lk.last_cleanup_size = new_cleanup_size;
+                lk.last_checkup = now;
+            }
+        }
+    }
+}

signal-gateway/src/lib.rs

@@ -12,6 +12,7 @@ pub mod message_handler;
 
 pub(crate) mod circular_buffer;
 pub(crate) mod concurrent_map;
+pub(crate) mod lazy_map_cleaner;
 pub(crate) mod limiter_sequence;
 pub(crate) mod log_format;
 pub(crate) mod log_message;

signal-gateway/src/rate_limiter.rs

@@ -1,6 +1,6 @@
 //! Rate limiting for log alerts.
 
-use crate::{concurrent_map::LazyMap, log_message::LogMessage};
+use crate::{concurrent_map::LazyMap, lazy_map_cleaner::{LazyMapCleaner, TsSecs}, log_message::LogMessage};
 use serde::Deserialize;
 use std::{
     str::FromStr,
@@ -163,18 +163,15 @@ impl SimpleRateLimiter {
 pub struct SourceLocationRateLimiter {
     /// Maps (file, line) -> rate limiter for that location
     limiters: LazyMap<(Box<str>, Box<str>), MultiRateLimiter>,
-    /// Threshold for creating new limiters (used for cleanup window calculation)
-    threshold: RateThreshold,
-    /// Maximum entries before triggering cleanup
-    last_cleanup_size: Mutex<usize>,
+    /// Manages cleanup of the lazy map
+    lazy_map_cleaner: LazyMapCleaner,
 }
 
 impl SourceLocationRateLimiter {
     pub fn new(threshold: RateThreshold) -> Self {
         Self {
-            limiters: LazyMap::with_capacity(16, move |_key| MultiRateLimiter::from(threshold)),
-            threshold,
-            last_cleanup_size: Mutex::new(8),
+            limiters: LazyMap::with_capacity(8, move |_key| MultiRateLimiter::from(threshold)),
+            lazy_map_cleaner: LazyMapCleaner::new(threshold.duration.as_secs() as i64),
         }
     }
 
@@ -189,27 +186,11 @@ impl SourceLocationRateLimiter {
 
         let result = self.limiters.get(&key, |limiter| limiter.evaluate(ts_sec));
 
-        // Opportunistically check for cleanup, but not if someone else is cleaning up
-        if let Ok(mut lk) = self.last_cleanup_size.try_lock() {
-            // If we're twice as large as the ending count from the last time we
-            // cleaned up, then let's try to clean up again.
-            if self.limiters.len() >= 2 * *lk {
-                *lk = self.cleanup(ts_sec);
-            }
-        }
+        // Maybe prune the lazy map
+        self.lazy_map_cleaner.maybe_clean(ts_sec, &self.limiters);
 
         result
     }
-
-    /// Remove entries where all timestamps are older than the window
-    fn cleanup(&self, now: i64) -> usize {
-        let window = self.threshold.duration.as_secs() as i64;
-        let cutoff = now - window;
-        self.limiters.retain(|_, limiter| {
-            // Keep if any timestamp is recent enough
-            limiter.get_latest() > cutoff
-        })
-    }
 }
 
 /// Implements rate-limiting criteria such as 'at least n in the last w seconds'
@@ -258,6 +239,12 @@ impl MultiRateLimiter {
     }
 }
 
+impl TsSecs for MultiRateLimiter {
+    fn ts_secs(&self) -> i64 {
+        self.get_latest()
+    }
+}
+
 struct MultiRateLimiterInner {
     /// Records the last n events as a ring buffer.
     /// Initialized to 0, representing "very distant past".