improvements to source location rate limiter cleanup scheme, concurrent map retain api and len impl

signal-gateway/src/concurrent_map.rs

@@ -18,7 +18,7 @@
 //! This is used instead of dash_map and once_map to avoid unnecessary complexity
 //! and dependencies, and give exactly the API needed in our application.
 
-use std::{borrow::Borrow, collections::HashMap, hash::Hash, sync::RwLock};
+use std::{borrow::Borrow, collections::HashMap, hash::Hash, sync::{atomic::{Ordering, AtomicUsize}, RwLock}};
 
 /// A concurrent hash map that uses read-preferring locking.
 ///
@@ -28,6 +28,7 @@ use std::{borrow::Borrow, collections::HashMap, hash::Hash, sync::RwLock};
 #[derive(Debug)]
 pub struct ConcurrentMap<K, V> {
     inner: RwLock<HashMap<K, V>>,
+    len_cache: AtomicUsize,
 }
 
 impl<K, V> ConcurrentMap<K, V>
@@ -38,6 +39,14 @@ where
     pub fn new() -> Self {
         Self::default()
     }
+    
+    /// Initialize a map with a given capacity
+    pub fn with_capacity(cap: usize) -> Self {
+        Self {
+            inner: RwLock::new(HashMap::with_capacity(cap)),
+            len_cache: Default::default(),
+        }
+    }
 
     /// Find and access a value in the map. If it doesn't exist, create it using the create function.
     ///
@@ -60,7 +69,10 @@ where
 
         // Value doesn't exist, need to create with write lock
         let mut guard = self.inner.write().unwrap();
-        let value = guard.entry(key.clone()).or_insert_with(create);
+        let value = guard.entry(key.clone()).or_insert_with(|| {
+            self.len_cache.fetch_add(1, Ordering::SeqCst);
+            (create)()
+        });
         access(value)
     }
 
@@ -74,17 +86,20 @@ where
     }
 
     /// Retain only entries that satisfy the predicate.
-    pub fn retain<F>(&self, f: F)
+    pub fn retain<F>(&self, f: F) -> usize
     where
         F: FnMut(&K, &mut V) -> bool,
     {
         let mut guard = self.inner.write().unwrap();
         guard.retain(f);
+        let len = guard.len();
+        self.len_cache.store(len, Ordering::SeqCst);
+        len
     }
 
     /// Get the number of entries in the map.
     pub fn len(&self) -> usize {
-        self.inner.read().unwrap().len()
+        self.len_cache.load(Ordering::SeqCst)
     }
 }
 
@@ -93,7 +108,8 @@ impl <K, V> Default for ConcurrentMap<K, V>
 {
     fn default() -> Self {
         Self {
-            inner: RwLock::new(Default::default())
+            inner: RwLock::new(Default::default()),
+            len_cache: Default::default(),
         }
     }
 }
@@ -119,6 +135,14 @@ where
         }
     }
 
+    /// Create a new lazy map with given capacity, and factory
+    pub fn with_capacity(cap: usize, factory: impl Fn(&K) -> V + Send + Sync + 'static) -> Self {
+        Self {
+            inner: ConcurrentMap::with_capacity(cap),
+            factory: Box::new(factory),
+        }    
+    }
+
     /// Get a value, creating it with the factory if it doesn't exist.
     pub fn get<Q, R, A>(&self, key: Q, access: A) -> R
     where
@@ -138,11 +162,11 @@ where
     }
 
     /// Retain only entries that satisfy the predicate.
-    pub fn retain<F>(&self, f: F)
+    pub fn retain<F>(&self, f: F) -> usize
     where
         F: FnMut(&K, &mut V) -> bool,
     {
-        self.inner.retain(f);
+        self.inner.retain(f)
     }
 
     /// Returns the number of entries in the map.

signal-gateway/src/gateway/log_handler.rs

@@ -209,14 +209,14 @@ impl LogHandler {
         // Send alert if we have formatted text
         if let Some((text, first_msg_len)) = formatted_text {
             let destination_override = rate_limit_result.ok().flatten();
-            if let Err(_err) = self.signal_alert_mq_tx.send(SignalAlertMessage {
+            if let Err(err) = self.signal_alert_mq_tx.send(SignalAlertMessage {
                 origin: Some(origin),
                 text,
                 attachment_paths: Default::default(),
                 summary: Summary::Prefix(first_msg_len),
                 destination_override,
             }) {
-                error!("Could not send alert message, queue is closed");
+                error!("Could not send alert message, queue is closed:\n{}", &err.0.text[0..first_msg_len]);
             }
         }
     }

signal-gateway/src/rate_limiter.rs

@@ -89,9 +89,6 @@ impl TryFrom<String> for RateThreshold {
     }
 }
 
-/// Maximum entries in a source-location rate limiter before triggering cleanup.
-const SOURCE_LOCATION_MAX_ENTRIES: usize = 2000;
-
 /// An enum over rate limiter implementations
 pub enum Limiter {
     /// Applies a threshold based on recent occurrences of the event.
@@ -128,7 +125,6 @@ impl Limiter {
     pub fn source_location(threshold: RateThreshold) -> Self {
         Limiter::SourceLocation(SourceLocationRateLimiter::new(
             threshold,
-            SOURCE_LOCATION_MAX_ENTRIES,
         ))
     }
 }
@@ -168,19 +164,19 @@ impl SimpleRateLimiter {
 /// gets its own `MultiRateLimiter` with the full threshold.
 pub struct SourceLocationRateLimiter {
     /// Maps (file, line) -> rate limiter for that location
-    limiters: LazyMap<(String, String), MultiRateLimiter>,
+    limiters: LazyMap<(Box<str>, Box<str>), MultiRateLimiter>,
     /// Threshold for creating new limiters (used for cleanup window calculation)
     threshold: RateThreshold,
     /// Maximum entries before triggering cleanup
-    max_entries: usize,
+    last_cleanup_size: Mutex<usize>,
 }
 
 impl SourceLocationRateLimiter {
-    pub fn new(threshold: RateThreshold, max_entries: usize) -> Self {
+    pub fn new(threshold: RateThreshold) -> Self {
         Self {
-            limiters: LazyMap::new(move |_key| MultiRateLimiter::from(threshold)),
+            limiters: LazyMap::with_capacity(16, move |_key| MultiRateLimiter::from(threshold)),
             threshold,
-            max_entries,
+            last_cleanup_size: Mutex::new(8),
         }
     }
 
@@ -188,26 +184,30 @@ impl SourceLocationRateLimiter {
     ///
     /// Returns true if the alert should fire (not rate-limited), false if suppressed.
     pub fn evaluate(&self, file: &str, line: &str, ts_sec: i64) -> bool {
-        let key = (file.to_owned(), line.to_owned());
+        let key = (file.to_owned().into_boxed_str(), line.to_owned().into_boxed_str());
 
         let result = self.limiters.get(&key, |limiter| limiter.evaluate(ts_sec));
 
-        // Clean up if we've exceeded max entries
+        // Opportunistically check for cleanup, but not if someone else is cleaning up
-        if self.limiters.len() > self.max_entries {
+        if let Ok(mut lk) = self.last_cleanup_size.try_lock() {
-            self.cleanup(ts_sec);
+            // 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);
+            }
         }
 
         result
     }
 
     /// Remove entries where all timestamps are older than the window
-    fn cleanup(&self, now: i64) {
+    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
-        });
+        })
     }
 }
 
@@ -438,7 +438,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold);
 
         // First alert from location A: no previous event to compare, returns false
         assert!(!limiter.evaluate("file_a.rs", "10", 1000));
@@ -466,7 +466,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold);
 
         // First three alerts: buffer filling, evicted=0, returns false
         assert!(!limiter.evaluate("file.rs", "10", 1000));
@@ -495,7 +495,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold);
 
         // First two events at each location: buffer filling, evicted=0
         assert!(!limiter.evaluate("file.rs", "10", 1000));
@@ -522,7 +522,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: false,
         };
-        let limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold);
 
         // First three alerts: evicted=0, threshold_met=false, returns true
         assert!(limiter.evaluate("file.rs", "10", 1000));
@@ -703,29 +703,25 @@ mod tests {
 
     #[test]
     fn source_location_rate_limiter_cleanup() {
-        // Use small max_entries to trigger cleanup
+        // Cleanup triggers when len >= 2 * last_cleanup_size (starts at 8, so triggers at 16)
         let threshold = RateThreshold {
             times: 1,
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let limiter = SourceLocationRateLimiter::new(threshold, 3);
+        let limiter = SourceLocationRateLimiter::new(threshold);
 
-        // Fill up the limiter (first events return false with evicted-vs-latest)
+        // Add 15 stale entries at timestamp 1000
-        assert!(!limiter.evaluate("file1.rs", "1", 1000));
+        for i in 0..15 {
-        assert!(!limiter.evaluate("file2.rs", "2", 1000));
+            limiter.evaluate(&format!("file{}.rs", i), &i.to_string(), 1000);
-        assert!(!limiter.evaluate("file3.rs", "3", 1000));
+        }
-        assert_eq!(limiter.limiters.len(), 3);
+        assert_eq!(limiter.limiters.len(), 15);
 
-        // Add one more, triggering cleanup - but all are fresh so none removed
+        // Add one more fresh entry at timestamp 2000 - this triggers cleanup (16 >= 2*8)
-        assert!(!limiter.evaluate("file4.rs", "4", 1000));
+        // cutoff = 2000 - 600 = 1400, so entries with latest <= 1400 are removed
-        // Still have 4 after cleanup since none are old enough
+        limiter.evaluate("fresh.rs", "1", 2000);
-        assert_eq!(limiter.limiters.len(), 4);
 
-        // Now add with a timestamp far in the future - old entries should be cleaned
+        // Only the fresh entry should remain
-        // (cutoff = 2000 - 600 = 1400, entries with latest <= 1400 are removed)
-        assert!(!limiter.evaluate("file5.rs", "5", 2000));
-        // Should have cleaned up entries from timestamp 1000 (older than 600 sec window)
         assert_eq!(limiter.limiters.len(), 1);
     }
 }