make rate limiters thread-safe with fine-grained locks, remove hard locks at higher levels

signal-gateway/src/concurrent_map.rs

@@ -4,9 +4,7 @@
 //! where values already exist, using a read lock first before falling back to a
 //! write lock for insertions.
 
-use std::{borrow::Borrow, collections::HashMap, hash::Hash};
+use std::{borrow::Borrow, collections::HashMap, hash::Hash, sync::RwLock};
-
-use tokio::sync::RwLock;
 
 /// A concurrent hash map that uses read-preferring locking.
 ///
@@ -38,10 +36,8 @@ where
     ///
     /// The key is only cloned when a new value needs to be inserted.
     ///
-    /// The lock is held while `access` runs, so `access` can safely use the reference.
+    /// The lock is held while `access` runs.
-    /// For async operations on the value, consider having `access` return a future
+    pub fn get_or_insert_with<Q, R, F, A>(&self, key: Q, create: F, access: A) -> R
-    /// that owns any data it needs.
-    pub async fn get_or_insert_with<Q, R, F, A>(&self, key: Q, create: F, access: A) -> R
     where
         Q: Borrow<K>,
         F: FnOnce() -> V,
@@ -51,14 +47,14 @@ where
 
         // Try to get existing value with read lock first
         {
-            let guard = self.inner.read().await;
+            let guard = self.inner.read().unwrap();
             if let Some(value) = guard.get(key) {
                 return access(value);
             }
         }
 
         // Value doesn't exist, need to create with write lock
-        let mut guard = self.inner.write().await;
+        let mut guard = self.inner.write().unwrap();
         // Use entry API - handles the race where another task inserted while we waited
         let value = guard.entry(key.clone()).or_insert_with(create);
         access(value)
@@ -68,13 +64,29 @@ where
     ///
     /// Acquires a read lock and calls `access` with a reference to the underlying HashMap.
     /// The lock is held while `access` runs.
-    pub async fn with_read_lock<R, A>(&self, access: A) -> R
+    pub fn with_read_lock<R, A>(&self, access: A) -> R
     where
         A: FnOnce(&HashMap<K, V>) -> R,
     {
-        let guard = self.inner.read().await;
+        let guard = self.inner.read().unwrap();
         access(&guard)
     }
+
+    /// Retain only entries that satisfy the predicate.
+    ///
+    /// Acquires a write lock and calls `retain` on the underlying HashMap.
+    pub fn retain<F>(&self, f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        let mut guard = self.inner.write().unwrap();
+        guard.retain(f);
+    }
+
+    /// Returns the number of entries in the map.
+    pub fn len(&self) -> usize {
+        self.inner.read().unwrap().len()
+    }
 }
 
 impl<K, V> Default for ConcurrentMap<K, V>
@@ -111,22 +123,33 @@ where
     ///
     /// Uses the factory provided at construction time to create new values.
     /// The key is only cloned when a new value needs to be inserted.
-    pub async fn get<Q, R, A>(&self, key: Q, access: A) -> R
+    pub fn get<Q, R, A>(&self, key: Q, access: A) -> R
     where
         Q: Borrow<K>,
         A: FnOnce(&V) -> R,
     {
-        self.inner
+        self.inner.get_or_insert_with(key, &self.factory, access)
-            .get_or_insert_with(key, &self.factory, access)
-            .await
     }
 
     /// Access all entries in the map with a read lock.
-    pub async fn with_read_lock<R, A>(&self, access: A) -> R
+    pub fn with_read_lock<R, A>(&self, access: A) -> R
     where
         A: FnOnce(&HashMap<K, V>) -> R,
     {
-        self.inner.with_read_lock(access).await
+        self.inner.with_read_lock(access)
+    }
+
+    /// Retain only entries that satisfy the predicate.
+    pub fn retain<F>(&self, f: F)
+    where
+        F: FnMut(&K, &mut V) -> bool,
+    {
+        self.inner.retain(f);
+    }
+
+    /// Returns the number of entries in the map.
+    pub fn len(&self) -> usize {
+        self.inner.len()
     }
 }
 
@@ -140,53 +163,48 @@ impl<K, V> std::fmt::Debug for LazyMap<K, V> {
 mod tests {
     use super::*;
 
-    #[tokio::test]
+    #[test]
-    async fn test_get_or_insert_new_key() {
+    fn test_get_or_insert_new_key() {
         let map: ConcurrentMap<String, i32> = ConcurrentMap::new();
 
-        let result = map
+        let result = map.get_or_insert_with("key1".to_string(), || 42, |v| *v);
-            .get_or_insert_with("key1".to_string(), || 42, |v| *v)
-            .await;
 
         assert_eq!(result, 42);
     }
 
-    #[tokio::test]
+    #[test]
-    async fn test_get_or_insert_existing_key() {
+    fn test_get_or_insert_existing_key() {
         let map: ConcurrentMap<String, i32> = ConcurrentMap::new();
 
         // Insert first time
-        map.get_or_insert_with("key1".to_string(), || 42, |_| ())
+        map.get_or_insert_with("key1".to_string(), || 42, |_| ());
-            .await;
 
         // Access again - should get existing value, not call create
         let mut create_called = false;
-        let result = map
+        let result = map.get_or_insert_with(
-            .get_or_insert_with(
+            "key1".to_string(),
-                "key1".to_string(),
+            || {
-                || {
+                create_called = true;
-                    create_called = true;
+                100
-                    100
+            },
-                },
+            |v| *v,
-                |v| *v,
+        );
-            )
-            .await;
 
         assert_eq!(result, 42);
         assert!(!create_called);
     }
 
-    #[tokio::test]
+    #[test]
-    async fn test_get_or_insert_multiple_keys() {
+    fn test_get_or_insert_multiple_keys() {
         let map: ConcurrentMap<String, i32> = ConcurrentMap::new();
 
-        map.get_or_insert_with("a".to_string(), || 1, |_| ()).await;
+        map.get_or_insert_with("a".to_string(), || 1, |_| ());
-        map.get_or_insert_with("b".to_string(), || 2, |_| ()).await;
+        map.get_or_insert_with("b".to_string(), || 2, |_| ());
-        map.get_or_insert_with("c".to_string(), || 3, |_| ()).await;
+        map.get_or_insert_with("c".to_string(), || 3, |_| ());
 
-        let a = map.get_or_insert_with("a".to_string(), || 0, |v| *v).await;
+        let a = map.get_or_insert_with("a".to_string(), || 0, |v| *v);
-        let b = map.get_or_insert_with("b".to_string(), || 0, |v| *v).await;
+        let b = map.get_or_insert_with("b".to_string(), || 0, |v| *v);
-        let c = map.get_or_insert_with("c".to_string(), || 0, |v| *v).await;
+        let c = map.get_or_insert_with("c".to_string(), || 0, |v| *v);
 
         assert_eq!(a, 1);
         assert_eq!(b, 2);

signal-gateway/src/gateway/log_handler.rs

@@ -11,7 +11,7 @@ use crate::{
 use chrono::Utc;
 use conf::Conf;
 use std::fmt;
-use tokio::sync::{Mutex, mpsc::UnboundedSender};
+use tokio::sync::mpsc::UnboundedSender;
 use tracing::{error, info};
 
 /// Reason why an alert was suppressed by rate limiting.
@@ -79,9 +79,9 @@ pub struct LogHandler {
     /// Log buffers keyed by origin (app + host). Lazily created.
     log_buffers: LazyMap<Origin, LogBuffer>,
     /// Routes with their associated limiter sets.
-    routes: Vec<(Route, Mutex<LimiterSet>)>,
+    routes: Vec<(Route, LimiterSet)>,
     /// Overall rate limiters applied after route checks pass.
-    overall_limits: Vec<Mutex<Limiter>>,
+    overall_limits: Vec<Limiter>,
 }
 
 impl LogHandler {
@@ -93,13 +93,13 @@ impl LogHandler {
         let routes = config
             .routes
             .iter()
-            .map(|route| (route.clone(), Mutex::new(route.make_limiter_set())))
+            .map(|route| (route.clone(), route.make_limiter_set()))
             .collect();
 
         let overall_limits = config
             .overall_limits
             .iter()
-            .map(|limit| Mutex::new(limit.make_limiter()))
+            .map(|limit| limit.make_limiter())
             .collect();
 
         let buffer_size = config.log_buffer_size;
@@ -153,7 +153,6 @@ impl LogHandler {
                     text
                 }
             })
-            .await
     }
 
     /// Consume a new log message from the given origin
@@ -196,8 +195,7 @@ impl LogHandler {
 
                     Some((text, first_msg_len))
                 }
-            })
+            });
-            .await;
 
         // Send alert if we have formatted text
         if let Some((text, first_msg_len)) = formatted_text {
@@ -246,7 +244,7 @@ impl LogHandler {
             }
 
             // Filter matched, evaluate the limiter set
-            let result = limiter_set.lock().await.evaluate(log_msg, origin, ts_sec);
+            let result = limiter_set.evaluate(log_msg, origin, ts_sec);
 
             match result {
                 LimitResult::Passed => {
@@ -271,7 +269,7 @@ impl LogHandler {
 
         // At least one route passed, now check overall limits
         for (idx, limiter) in self.overall_limits.iter().enumerate() {
-            if !limiter.lock().await.evaluate(log_msg, ts_sec) {
+            if !limiter.evaluate(log_msg, ts_sec) {
                 return Err(SuppressionReason::Overall(idx, LimitResult::Limiter(0)));
             }
         }

signal-gateway/src/gateway/rate_limiter_set.rs

@@ -1,10 +1,10 @@
 //! Rate limiter set for managing per-route rate limiting.
 
 use crate::{
+    concurrent_map::LazyMap,
     log_message::{LogMessage, Origin},
     rate_limiter::Limiter,
 };
-use std::collections::HashMap;
 
 /// Result of evaluating a limiter set.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -19,10 +19,8 @@ pub enum LimitResult {
 
 /// A set of limiters for a route, containing both per-origin and global limiters.
 pub struct LimiterSet {
-    /// Factory to create limiters for new origins.
-    make_limiters: Box<dyn Fn() -> Vec<Limiter> + Send + Sync>,
     /// Per-origin rate limiters, keyed by origin. Lazily created.
-    limiters: HashMap<Origin, Vec<Limiter>>,
+    limiters: LazyMap<Origin, Vec<Limiter>>,
     /// Global rate limiters (shared across all origins).
     global_limiters: Vec<Limiter>,
 }
@@ -34,8 +32,7 @@ impl LimiterSet {
         global_limiters: Vec<Limiter>,
     ) -> Self {
         Self {
-            make_limiters: Box::new(make_limiters),
+            limiters: LazyMap::new(make_limiters),
-            limiters: HashMap::new(),
             global_limiters,
         }
     }
@@ -45,20 +42,23 @@ impl LimiterSet {
     /// Returns [`LimitResult::Passed`] if the event passes all limits.
     /// Returns [`LimitResult::Limiter(i)`] if blocked by per-origin limiter at index `i`.
     /// Returns [`LimitResult::GlobalLimiter(i)`] if blocked by global limiter at index `i`.
-    pub fn evaluate(&mut self, log_msg: &LogMessage, origin: &Origin, ts_sec: i64) -> LimitResult {
+    pub fn evaluate(&self, log_msg: &LogMessage, origin: &Origin, ts_sec: i64) -> LimitResult {
-        // Get or create limiters for this origin
+        // Check per-origin limiters
-        let origin_limiters = self
+        let origin_result = self.limiters.get(origin, |origin_limiters| {
-            .limiters
+            for (i, limiter) in origin_limiters.iter().enumerate() {
-            .entry(origin.clone())
+                if !limiter.evaluate(log_msg, ts_sec) {
-            .or_insert_with(&self.make_limiters);
+                    return Some(LimitResult::Limiter(i));
-
+                }
-        for (i, limiter) in origin_limiters.iter_mut().enumerate() {
-            if !limiter.evaluate(log_msg, ts_sec) {
-                return LimitResult::Limiter(i);
             }
+            None
+        });
+
+        if let Some(result) = origin_result {
+            return result;
         }
 
-        for (i, limiter) in self.global_limiters.iter_mut().enumerate() {
+        // Check global limiters
+        for (i, limiter) in self.global_limiters.iter().enumerate() {
             if !limiter.evaluate(log_msg, ts_sec) {
                 return LimitResult::GlobalLimiter(i);
             }

signal-gateway/src/rate_limiter.rs

@@ -1,11 +1,13 @@
 //! Rate limiting for log alerts.
 
-use crate::log_message::LogMessage;
+use crate::{concurrent_map::LazyMap, log_message::LogMessage};
 use serde::Deserialize;
 use std::{
-    collections::HashMap,
     str::FromStr,
-    sync::atomic::{AtomicI64, Ordering},
+    sync::{
+        Mutex,
+        atomic::{AtomicI64, Ordering},
+    },
     time::Duration,
 };
 
@@ -105,7 +107,7 @@ impl Limiter {
     ///
     /// Returns `true` if the event should be allowed (not rate-limited),
     /// `false` if it should be suppressed.
-    pub fn evaluate(&mut self, log_msg: &LogMessage, ts_sec: i64) -> bool {
+    pub fn evaluate(&self, log_msg: &LogMessage, ts_sec: i64) -> bool {
         match self {
             Limiter::Multi(limiter) => limiter.evaluate(ts_sec),
             Limiter::SourceLocation(limiter) => {
@@ -154,23 +156,12 @@ impl SimpleRateLimiter {
         let last_ts = self.last_timestamp.load(Ordering::SeqCst);
         let rate_limited = ts_sec - last_ts < self.window;
         if !rate_limited && ts_sec > last_ts {
-            // If this is called concurrently, guarantee that we keep going
+            self.last_timestamp.fetch_max(ts_sec, Ordering::SeqCst);
-            // until the max value is stored at self.last_timestamp,
-            // so self.last_timestamp is "eventually" only monotonically increasing.
-            store_max(ts_sec, &self.last_timestamp);
         }
         !rate_limited
     }
 }
 
-#[allow(dead_code)]
-fn store_max(val: i64, at: &AtomicI64) {
-    let prev = at.swap(val, Ordering::SeqCst);
-    if prev > val {
-        store_max(prev, at)
-    }
-}
-
 /// A rate limiter that tracks alerts per source location (file:line).
 ///
 /// This allows different error locations to alert independently, preventing one noisy
@@ -178,8 +169,8 @@ fn store_max(val: i64, at: &AtomicI64) {
 /// gets its own `MultiRateLimiter` with the full threshold.
 pub struct SourceLocationRateLimiter {
     /// Maps (file, line) -> rate limiter for that location
-    limiters: HashMap<(String, String), MultiRateLimiter>,
+    limiters: LazyMap<(String, String), MultiRateLimiter>,
-    /// Threshold for creating new limiters
+    /// Threshold for creating new limiters (used for cleanup window calculation)
     threshold: RateThreshold,
     /// Maximum entries before triggering cleanup
     max_entries: usize,
@@ -188,7 +179,7 @@ pub struct SourceLocationRateLimiter {
 impl SourceLocationRateLimiter {
     pub fn new(threshold: RateThreshold, max_entries: usize) -> Self {
         Self {
-            limiters: HashMap::new(),
+            limiters: LazyMap::new(move || MultiRateLimiter::from(threshold)),
             threshold,
             max_entries,
         }
@@ -197,15 +188,10 @@ impl SourceLocationRateLimiter {
     /// Check if an error from this source location should trigger an alert.
     ///
     /// Returns true if the alert should fire (not rate-limited), false if suppressed.
-    pub fn evaluate(&mut self, file: &str, line: &str, ts_sec: i64) -> bool {
+    pub fn evaluate(&self, file: &str, line: &str, ts_sec: i64) -> bool {
         let key = (file.to_owned(), line.to_owned());
 
-        let limiter = self
+        let result = self.limiters.get(&key, |limiter| limiter.evaluate(ts_sec));
-            .limiters
-            .entry(key)
-            .or_insert_with(|| MultiRateLimiter::from(self.threshold));
-
-        let result = limiter.evaluate(ts_sec);
 
         // Clean up if we've exceeded max entries
         if self.limiters.len() > self.max_entries {
@@ -216,7 +202,7 @@ impl SourceLocationRateLimiter {
     }
 
     /// Remove entries where all timestamps are older than the window
-    fn cleanup(&mut self, now: i64) {
+    fn cleanup(&self, now: i64) {
         let window = self.threshold.duration.as_secs() as i64;
         let cutoff = now - window;
         self.limiters.retain(|_, limiter| {
@@ -230,35 +216,35 @@ impl SourceLocationRateLimiter {
 ///
 /// Uses a ring buffer to track the N most recent timestamps.
 pub struct MultiRateLimiter {
-    inner: MultiRateLimiterInner,
+    inner: Mutex<MultiRateLimiterInner>,
     /// The length of the window (in seconds)
     window: i64,
     /// If true, returns true when rate >= threshold (burst detection).
     /// If false, returns true when rate < threshold (suppression).
     comparator_is_ge: bool,
     /// Cache of latest timestamp recorded
-    latest: i64,
+    latest: AtomicI64,
 }
 
 impl MultiRateLimiter {
     pub fn new(num: usize, window: Duration, comparator_is_ge: bool) -> Self {
         Self {
-            inner: MultiRateLimiterInner::new(num),
+            inner: Mutex::new(MultiRateLimiterInner::new(num)),
             window: window.as_secs().try_into().unwrap(),
             comparator_is_ge,
-            latest: 0,
+            latest: AtomicI64::new(0),
         }
     }
 
     /// Get the latest timestamp recorded
     pub fn get_latest(&self) -> i64 {
-        self.latest
+        self.latest.load(Ordering::SeqCst)
     }
 
     /// Check if a particular new timestamp passes the limit. This also updates the internal state.
-    pub fn evaluate(&mut self, new_timestamp: i64) -> bool {
+    pub fn evaluate(&self, new_timestamp: i64) -> bool {
-        let (earliest, latest) = self.inner.insert_and_pop(new_timestamp);
+        let (earliest, latest) = self.inner.lock().unwrap().insert_and_pop(new_timestamp);
-        self.latest = self.latest.max(latest);
+        self.latest.fetch_max(latest, Ordering::SeqCst);
 
         let threshold_met = earliest + self.window >= latest;
 
@@ -454,7 +440,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let mut limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold, 100);
 
         // First alert from location A: no previous event to compare, returns false
         assert!(!limiter.evaluate("file_a.rs", "10", 1000));
@@ -482,7 +468,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let mut limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold, 100);
 
         // First three alerts: buffer filling, evicted=0, returns false
         assert!(!limiter.evaluate("file.rs", "10", 1000));
@@ -511,7 +497,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let mut limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold, 100);
 
         // First two events at each location: buffer filling, evicted=0
         assert!(!limiter.evaluate("file.rs", "10", 1000));
@@ -538,7 +524,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: false,
         };
-        let mut limiter = SourceLocationRateLimiter::new(threshold, 100);
+        let limiter = SourceLocationRateLimiter::new(threshold, 100);
 
         // First three alerts: evicted=0, threshold_met=false, returns true
         assert!(limiter.evaluate("file.rs", "10", 1000));
@@ -725,7 +711,7 @@ mod tests {
             duration: Duration::from_secs(600),
             comparator_is_ge: true,
         };
-        let mut limiter = SourceLocationRateLimiter::new(threshold, 3);
+        let limiter = SourceLocationRateLimiter::new(threshold, 3);
 
         // Fill up the limiter (first events return false with evicted-vs-latest)
         assert!(!limiter.evaluate("file1.rs", "1", 1000));