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

factor out limit sequence object, simplify code, fmt

Browse Source
This commit is contained in:
Chris Beck
2025-12-15 12:32:12 -07:00
parent 3ae7030404
commit 2670104b49
8 changed files with 124 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
signal-gateway/src/concurrent_map.rs
+16 -6
View File
@@ -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 <K, V> Default for ConcurrentMap<K, V>
where K: Eq + Hash
impl<K, V> Default for ConcurrentMap<K, V>
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.
signal-gateway/src/gateway/log_handler.rs
+11 -12
View File
@@ -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));
}
signal-gateway/src/gateway/mod.rs
+1 -1
View File
@@ -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};
signal-gateway/src/gateway/rate_limiter_set.rs
+10 -13
View File
@@ -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<Origin, Vec<(LogFilter, Limiter)>>,
limiters: LazyMap<Origin, LimiterSequence>,
/// 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);
}
signal-gateway/src/gateway/route.rs
+4 -57
View File
@@ -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)
}
}
signal-gateway/src/lib.rs
+2
View File
@@ -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,
signal-gateway/src/limiter_sequence.rs
+75
View File
@@ -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<T>(iter: T) -> Self
where
T: IntoIterator<Item = &'a Limit>,
{
Self(iter.into_iter().map(Limit::make_limiter).collect())
}
}
signal-gateway/src/rate_limiter.rs
+5 -4
View File
@@ -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));