diff --git a/signal-gateway/src/concurrent_map.rs b/signal-gateway/src/concurrent_map.rs index fc0022f..64dea32 100644 --- a/signal-gateway/src/concurrent_map.rs +++ b/signal-gateway/src/concurrent_map.rs @@ -18,7 +18,15 @@ //! 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::{atomic::{Ordering, AtomicUsize}, RwLock}}; +use std::{ + borrow::Borrow, + collections::HashMap, + hash::Hash, + sync::{ + RwLock, + atomic::{AtomicUsize, Ordering}, + }, +}; /// A concurrent hash map that uses read-preferring locking. /// @@ -39,7 +47,7 @@ where pub fn new() -> Self { Self::default() } - + /// Initialize a map with a given capacity pub fn with_capacity(cap: usize) -> Self { Self { @@ -103,8 +111,9 @@ where } } -impl Default for ConcurrentMap - where K: Eq + Hash +impl Default for ConcurrentMap +where + K: Eq + Hash, { fn default() -> Self { Self { @@ -140,7 +149,7 @@ where Self { inner: ConcurrentMap::with_capacity(cap), factory: Box::new(factory), - } + } } /// Get a value, creating it with the factory if it doesn't exist. @@ -150,7 +159,8 @@ where A: FnOnce(&V) -> R, { let key = key.borrow(); - self.inner.get_or_insert_with(key, || (self.factory)(key), access) + self.inner + .get_or_insert_with(key, || (self.factory)(key), access) } /// Access all entries in the map with a read lock. diff --git a/signal-gateway/src/gateway/log_handler.rs b/signal-gateway/src/gateway/log_handler.rs index 61cbaf5..e2e1083 100644 --- a/signal-gateway/src/gateway/log_handler.rs +++ b/signal-gateway/src/gateway/log_handler.rs @@ -1,13 +1,14 @@ use super::{ - LimitResult, Limiter, LimiterSet, SignalAlertMessage, Summary, evaluate_limiter_sequence, + LimitResult, LimiterSet, SignalAlertMessage, Summary, log_buffer::LogBuffer, - route::{Destination, Limit, Route}, + route::{Destination, Route}, }; use crate::{ assistant::{Tool, ToolExecutor, ToolResult}, concurrent_map::LazyMap, + limiter_sequence::{Limit, LimiterSequence}, log_format::LogFormatConfig, - log_message::{LogFilter, LogMessage, Origin}, + log_message::{LogMessage, Origin}, }; use async_trait::async_trait; use chrono::Utc; @@ -88,8 +89,7 @@ pub struct LogHandler { /// Routes with their associated limiter sets. routes: Vec<(Route, LimiterSet)>, /// Overall rate limiters applied after route checks pass. - /// Each entry is a (filter, limiter) pair. - overall_limits: Vec<(LogFilter, Limiter)>, + overall_limits: LimiterSequence, } impl LogHandler { @@ -104,11 +104,7 @@ impl LogHandler { .map(|route| (route.clone(), route.make_limiter_set())) .collect(); - let overall_limits = config - .overall_limits - .iter() - .map(|limit| limit.make_limiter()) - .collect(); + let overall_limits = config.overall_limits.iter().collect(); let buffer_size = config.log_buffer_size; @@ -216,7 +212,10 @@ impl LogHandler { summary: Summary::Prefix(first_msg_len), destination_override, }) { - error!("Could not send alert message, queue is closed:\n{}", &err.0.text[0..first_msg_len]); + error!( + "Could not send alert message, queue is closed:\n{}", + &err.0.text[0..first_msg_len] + ); } } } @@ -276,7 +275,7 @@ impl LogHandler { }; // At least one route passed, now check overall limits - if let Err(i) = evaluate_limiter_sequence(&self.overall_limits, log_msg) { + if let Err(i) = self.overall_limits.evaluate(log_msg) { return Err(SuppressionReason::OverallLimiter(i)); } diff --git a/signal-gateway/src/gateway/mod.rs b/signal-gateway/src/gateway/mod.rs index b1cabe3..b83386d 100644 --- a/signal-gateway/src/gateway/mod.rs +++ b/signal-gateway/src/gateway/mod.rs @@ -48,7 +48,7 @@ mod rate_limiter_set; pub use rate_limiter_set::{LimitResult, LimiterSet}; mod route; -pub use route::{Destination, Limit, Route, evaluate_limiter_sequence}; +pub use route::{Destination, Route}; pub use crate::rate_limiter::{Limiter, RateThreshold}; diff --git a/signal-gateway/src/gateway/rate_limiter_set.rs b/signal-gateway/src/gateway/rate_limiter_set.rs index af98d4c..5a2dac8 100644 --- a/signal-gateway/src/gateway/rate_limiter_set.rs +++ b/signal-gateway/src/gateway/rate_limiter_set.rs @@ -1,10 +1,9 @@ //! Rate limiter set for managing per-route rate limiting. -use super::evaluate_limiter_sequence; use crate::{ concurrent_map::LazyMap, - log_message::{LogFilter, LogMessage, Origin}, - rate_limiter::Limiter, + limiter_sequence::LimiterSequence, + log_message::{LogMessage, Origin}, }; /// Result of evaluating a limiter set. @@ -18,22 +17,19 @@ pub enum LimitResult { GlobalLimiter(usize), } -/// A set of limiters for a route, containing both per-origin and global limiters. -/// Each limiter is paired with a filter that must match before the limiter is evaluated. +/// A set of limiter sequences for a route, containing both per-origin and global limiters. pub struct LimiterSet { /// Per-origin rate limiters, keyed by origin. Lazily created. - /// Each entry is a (filter, limiter) pair. - limiters: LazyMap>, + limiters: LazyMap, /// Global rate limiters (shared across all origins). - /// Each entry is a (filter, limiter) pair. - global_limiters: Vec<(LogFilter, Limiter)>, + global_limiters: LimiterSequence, } impl LimiterSet { /// Create a new limiter set with factories for per-origin and global limiters. pub fn new( - make_limiters: impl Fn() -> Vec<(LogFilter, Limiter)> + Send + Sync + 'static, - global_limiters: Vec<(LogFilter, Limiter)>, + make_limiters: impl Fn() -> LimiterSequence + Send + Sync + 'static, + global_limiters: LimiterSequence, ) -> Self { Self { limiters: LazyMap::new(move |_key| make_limiters()), @@ -53,7 +49,8 @@ impl LimiterSet { pub fn evaluate(&self, log_msg: &LogMessage, origin: &Origin) -> LimitResult { // Check per-origin limiters let origin_result = self.limiters.get(origin, |origin_limiters| { - evaluate_limiter_sequence(origin_limiters, log_msg) + origin_limiters + .evaluate(log_msg) .err() .map(LimitResult::Limiter) }); @@ -63,7 +60,7 @@ impl LimiterSet { } // Check global limiters - if let Err(i) = evaluate_limiter_sequence(&self.global_limiters, log_msg) { + if let Err(i) = self.global_limiters.evaluate(log_msg) { return LimitResult::GlobalLimiter(i); } diff --git a/signal-gateway/src/gateway/route.rs b/signal-gateway/src/gateway/route.rs index 18a6a87..da465de 100644 --- a/signal-gateway/src/gateway/route.rs +++ b/signal-gateway/src/gateway/route.rs @@ -5,57 +5,11 @@ use super::LimiterSet; use crate::{ - log_message::{Level, LogFilter, LogMessage}, - rate_limiter::{Limiter, RateThreshold}, + limiter_sequence::Limit, + log_message::{Level, LogFilter}, }; use serde::Deserialize; -/// Evaluate a sequence of (filter, limiter) pairs against a log message. -/// -/// Returns `Ok(())` if no limiter blocks the message. -/// Returns `Err(index)` if the limiter at `index` blocked the message. -pub fn evaluate_limiter_sequence( - seq: &[(LogFilter, Limiter)], - log_msg: &LogMessage, -) -> Result<(), usize> { - for (i, (filter, limiter)) in seq.iter().enumerate() { - if filter.matches(log_msg) && !limiter.evaluate(log_msg) { - return Err(i); - } - } - Ok(()) -} - -/// A rate limit rule for suppressing repeated alerts. -/// -/// Combines a filter to match specific log messages with a rate threshold. -#[derive(Clone, Debug, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Limit { - /// Filter criteria for messages this limit applies to. - #[serde(flatten)] - pub filter: LogFilter, - /// Rate threshold for suppressing alerts. - pub threshold: RateThreshold, - /// If true, rate limit independently per source location (file:line). - /// If false (default), count all matching events together. - #[serde(default)] - pub by_source_location: bool, -} - -impl Limit { - /// Create the appropriate limiter for this limit configuration. - /// Returns a (filter, limiter) pair so the filter can be checked before rate limiting. - pub fn make_limiter(&self) -> (LogFilter, Limiter) { - let limiter = if self.by_source_location { - Limiter::source_location(self.threshold) - } else { - Limiter::multi(self.threshold) - }; - (self.filter.clone(), limiter) - } -} - /// A route configuration for processing log messages. /// /// Routes match incoming log messages based on an optional filter, then apply @@ -99,15 +53,8 @@ impl Route { /// Create a limiter set from this route's limit configurations. pub fn make_limiter_set(&self) -> LimiterSet { let limits = self.limits.clone(); - let global_limiters = self - .global_limits - .iter() - .map(|l| l.make_limiter()) - .collect(); - LimiterSet::new( - move || limits.iter().map(|l| l.make_limiter()).collect(), - global_limiters, - ) + let global_limiters = self.global_limits.iter().collect(); + LimiterSet::new(move || limits.iter().collect(), global_limiters) } } diff --git a/signal-gateway/src/lib.rs b/signal-gateway/src/lib.rs index b350a56..1d3681d 100644 --- a/signal-gateway/src/lib.rs +++ b/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 limiter_sequence; pub(crate) mod log_format; pub(crate) mod log_message; pub(crate) mod prometheus; @@ -20,6 +21,7 @@ pub(crate) mod signal_jsonrpc; pub(crate) mod transports; pub use gateway::{CommandRouter, CommandRouterBuilder, Gateway, GatewayConfig, Handling}; +pub use limiter_sequence::Limit; pub use log_message::{Level, LogFilter, LogMessage, LogMessageBuilder}; pub use message_handler::{ AdminMessage, AdminMessageResponse, Context, MessageHandler, MessageHandlerResult, diff --git a/signal-gateway/src/limiter_sequence.rs b/signal-gateway/src/limiter_sequence.rs new file mode 100644 index 0000000..251de6d --- /dev/null +++ b/signal-gateway/src/limiter_sequence.rs @@ -0,0 +1,75 @@ +//! Configuration for a sequence of rate limiters, each with a filter criteria. +//! When a sequence is evaluated, all rate limiters are evaluated against the message, +//! without stopping early. But if any of them blocks the message, it is blocked, +//! and the first one to do so is returned. + +use crate::{ + log_message::{LogFilter, LogMessage}, + rate_limiter::{Limiter, RateThreshold}, +}; +use serde::Deserialize; + +/// A rate limit rule for suppressing repeated alerts. +/// +/// Combines a filter to match specific log messages with a rate threshold. +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Limit { + /// Filter criteria for messages this limit applies to. + #[serde(flatten)] + pub filter: LogFilter, + /// Rate threshold for suppressing alerts. + pub threshold: RateThreshold, + /// If true, rate limit independently per source location (file:line). + /// If false (default), count all matching events together. + #[serde(default)] + pub by_source_location: bool, +} + +impl Limit { + /// Create the appropriate limiter for this limit configuration. + /// Returns a (filter, limiter) pair so the filter can be checked before rate limiting. + pub fn make_limiter(&self) -> (LogFilter, Limiter) { + let limiter = if self.by_source_location { + Limiter::source_location(self.threshold) + } else { + Limiter::multi(self.threshold) + }; + (self.filter.clone(), limiter) + } +} + +/// A sequence of limits each applied in parallel to incoming messages. +/// If any of them blocks a message, it is blocked. +pub struct LimiterSequence(Vec<(LogFilter, Limiter)>); + +impl LimiterSequence { + /// Evaluate a sequence of (filter, limiter) pairs against a log message. + /// + /// Returns `Ok(())` if no limiter blocks the message. + /// Returns `Err(index)` if the limiter at `index` blocked the message. + /// + /// Semantic: + /// All limiters are evaluated regardless of if an earlier limiter blocks. + pub fn evaluate(&self, log_msg: &LogMessage) -> Result<(), usize> { + let mut maybe_idx = None; + for (i, (filter, limiter)) in self.0.iter().enumerate() { + if filter.matches(log_msg) && !limiter.evaluate(log_msg) && maybe_idx.is_none() { + maybe_idx = Some(i); + } + } + match maybe_idx { + None => Ok(()), + Some(i) => Err(i), + } + } +} + +impl<'a> FromIterator<&'a Limit> for LimiterSequence { + fn from_iter(iter: T) -> Self + where + T: IntoIterator, + { + Self(iter.into_iter().map(Limit::make_limiter).collect()) + } +} diff --git a/signal-gateway/src/rate_limiter.rs b/signal-gateway/src/rate_limiter.rs index 73df634..0fd1e2d 100644 --- a/signal-gateway/src/rate_limiter.rs +++ b/signal-gateway/src/rate_limiter.rs @@ -123,9 +123,7 @@ impl Limiter { /// /// Each source location (file:line) gets its own rate limiter with the full threshold. pub fn source_location(threshold: RateThreshold) -> Self { - Limiter::SourceLocation(SourceLocationRateLimiter::new( - threshold, - )) + Limiter::SourceLocation(SourceLocationRateLimiter::new(threshold)) } } @@ -184,7 +182,10 @@ 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().into_boxed_str(), line.to_owned().into_boxed_str()); + 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));