diff --git a/signal-gateway/src/concurrent_map.rs b/signal-gateway/src/concurrent_map.rs index 98f6dc1..54c8a95 100644 --- a/signal-gateway/src/concurrent_map.rs +++ b/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 tokio::sync::RwLock; +use std::{borrow::Borrow, collections::HashMap, hash::Hash, 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. - /// For async operations on the value, consider having `access` return a future - /// that owns any data it needs. - pub async fn get_or_insert_with(&self, key: Q, create: F, access: A) -> R + /// The lock is held while `access` runs. + pub fn get_or_insert_with(&self, key: Q, create: F, access: A) -> R where Q: Borrow, 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(&self, access: A) -> R + pub fn with_read_lock(&self, access: A) -> R where A: FnOnce(&HashMap) -> 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(&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 Default for ConcurrentMap @@ -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(&self, key: Q, access: A) -> R + pub fn get(&self, key: Q, access: A) -> R where Q: Borrow, A: FnOnce(&V) -> R, { - self.inner - .get_or_insert_with(key, &self.factory, access) - .await + self.inner.get_or_insert_with(key, &self.factory, access) } /// Access all entries in the map with a read lock. - pub async fn with_read_lock(&self, access: A) -> R + pub fn with_read_lock(&self, access: A) -> R where A: FnOnce(&HashMap) -> R, { - self.inner.with_read_lock(access).await + self.inner.with_read_lock(access) + } + + /// Retain only entries that satisfy the predicate. + pub fn retain(&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 std::fmt::Debug for LazyMap { mod tests { use super::*; - #[tokio::test] - async fn test_get_or_insert_new_key() { + #[test] + fn test_get_or_insert_new_key() { let map: ConcurrentMap = ConcurrentMap::new(); - let result = map - .get_or_insert_with("key1".to_string(), || 42, |v| *v) - .await; + let result = map.get_or_insert_with("key1".to_string(), || 42, |v| *v); assert_eq!(result, 42); } - #[tokio::test] - async fn test_get_or_insert_existing_key() { + #[test] + fn test_get_or_insert_existing_key() { let map: ConcurrentMap = ConcurrentMap::new(); // Insert first time - map.get_or_insert_with("key1".to_string(), || 42, |_| ()) - .await; + map.get_or_insert_with("key1".to_string(), || 42, |_| ()); // Access again - should get existing value, not call create let mut create_called = false; - let result = map - .get_or_insert_with( - "key1".to_string(), - || { - create_called = true; - 100 - }, - |v| *v, - ) - .await; + let result = map.get_or_insert_with( + "key1".to_string(), + || { + create_called = true; + 100 + }, + |v| *v, + ); assert_eq!(result, 42); assert!(!create_called); } - #[tokio::test] - async fn test_get_or_insert_multiple_keys() { + #[test] + fn test_get_or_insert_multiple_keys() { let map: ConcurrentMap = ConcurrentMap::new(); - map.get_or_insert_with("a".to_string(), || 1, |_| ()).await; - map.get_or_insert_with("b".to_string(), || 2, |_| ()).await; - map.get_or_insert_with("c".to_string(), || 3, |_| ()).await; + map.get_or_insert_with("a".to_string(), || 1, |_| ()); + map.get_or_insert_with("b".to_string(), || 2, |_| ()); + map.get_or_insert_with("c".to_string(), || 3, |_| ()); - let a = map.get_or_insert_with("a".to_string(), || 0, |v| *v).await; - let b = map.get_or_insert_with("b".to_string(), || 0, |v| *v).await; - let c = map.get_or_insert_with("c".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); + let c = map.get_or_insert_with("c".to_string(), || 0, |v| *v); assert_eq!(a, 1); assert_eq!(b, 2); diff --git a/signal-gateway/src/gateway/log_handler.rs b/signal-gateway/src/gateway/log_handler.rs index d772d56..7418ef5 100644 --- a/signal-gateway/src/gateway/log_handler.rs +++ b/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, /// Routes with their associated limiter sets. - routes: Vec<(Route, Mutex)>, + routes: Vec<(Route, LimiterSet)>, /// Overall rate limiters applied after route checks pass. - overall_limits: Vec>, + overall_limits: Vec, } 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))); } } diff --git a/signal-gateway/src/gateway/rate_limiter_set.rs b/signal-gateway/src/gateway/rate_limiter_set.rs index d1ae3d9..4cff05d 100644 --- a/signal-gateway/src/gateway/rate_limiter_set.rs +++ b/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 Vec + Send + Sync>, /// Per-origin rate limiters, keyed by origin. Lazily created. - limiters: HashMap>, + limiters: LazyMap>, /// Global rate limiters (shared across all origins). global_limiters: Vec, } @@ -34,8 +32,7 @@ impl LimiterSet { global_limiters: Vec, ) -> Self { Self { - make_limiters: Box::new(make_limiters), - limiters: HashMap::new(), + limiters: LazyMap::new(make_limiters), 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 { - // Get or create limiters for this origin - let origin_limiters = self - .limiters - .entry(origin.clone()) - .or_insert_with(&self.make_limiters); - - for (i, limiter) in origin_limiters.iter_mut().enumerate() { - if !limiter.evaluate(log_msg, ts_sec) { - return LimitResult::Limiter(i); + pub fn evaluate(&self, log_msg: &LogMessage, origin: &Origin, ts_sec: i64) -> LimitResult { + // Check per-origin limiters + let origin_result = self.limiters.get(origin, |origin_limiters| { + for (i, limiter) in origin_limiters.iter().enumerate() { + if !limiter.evaluate(log_msg, ts_sec) { + return Some(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); } diff --git a/signal-gateway/src/rate_limiter.rs b/signal-gateway/src/rate_limiter.rs index 7ff54a4..b6b3de3 100644 --- a/signal-gateway/src/rate_limiter.rs +++ b/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 - // 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); + self.last_timestamp.fetch_max(ts_sec, Ordering::SeqCst); } !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>, - /// Threshold for creating new limiters + limiters: LazyMap<(String, String), MultiRateLimiter>, + /// 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 - .limiters - .entry(key) - .or_insert_with(|| MultiRateLimiter::from(self.threshold)); - - let result = limiter.evaluate(ts_sec); + let result = self.limiters.get(&key, |limiter| 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, /// 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 { - let (earliest, latest) = self.inner.insert_and_pop(new_timestamp); - self.latest = self.latest.max(latest); + pub fn evaluate(&self, new_timestamp: i64) -> bool { + let (earliest, latest) = self.inner.lock().unwrap().insert_and_pop(new_timestamp); + 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));