Labyricorn/signal-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Chris Beck
2025-12-06 02:34:36 -07:00
parent 9b47963f16
commit ae3761bffe
4 changed files with 113 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
signal-gateway/src/concurrent_map.rs
+62 -44
View File
@@ -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<Q, R, F, A>(&self, key: Q, create: F, access: A) -> R
/// The lock is held while `access` runs.
pub 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
.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<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]
async fn test_get_or_insert_new_key() {
#[test]
fn test_get_or_insert_new_key() {
let map: ConcurrentMap<String, i32> = 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<String, i32> = 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<String, i32> = 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);
signal-gateway/src/gateway/log_handler.rs
+8 -10
View File
@@ -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
+17 -17
View File
@@ -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: 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);
}
signal-gateway/src/rate_limiter.rs
+26 -40
View File
@@ -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<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 {
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));