diff --git a/signal-gateway/src/circular_buffer.rs b/signal-gateway/src/circular_buffer.rs index 4c630aa..e8e02ce 100644 --- a/signal-gateway/src/circular_buffer.rs +++ b/signal-gateway/src/circular_buffer.rs @@ -38,6 +38,11 @@ impl CircularBuffer { self.buf.push_back(value); } + /// Returns the back element of the buffer + pub fn back(&self) -> Option<&T> { + self.buf.back() + } + /// Returns the number of elements in the buffer. #[allow(unused)] pub fn len(&self) -> usize { diff --git a/signal-gateway/src/gateway/log_buffer.rs b/signal-gateway/src/gateway/log_buffer.rs index a01654f..2f29345 100644 --- a/signal-gateway/src/gateway/log_buffer.rs +++ b/signal-gateway/src/gateway/log_buffer.rs @@ -2,7 +2,7 @@ //! //! Wraps a circular buffer with a synchronous mutex for fast, blocking access. -use crate::{circular_buffer::CircularBuffer, log_message::LogMessage}; +use crate::{circular_buffer::CircularBuffer, lazy_map_cleaner::TsSecs, log_message::LogMessage}; use std::sync::Mutex; /// A thread-safe circular buffer for log messages. @@ -58,4 +58,17 @@ impl LogBuffer { let mut iter = buf.iter().rev(); f(&mut iter) } + + /// Access the back of the buffer, if it exists. Returns None if not. + pub fn peek_back(&self, f: impl FnOnce(&LogMessage) -> R) -> Option { + let buf = self.buf.lock().unwrap(); + Some(f(buf.back()?)) + } +} + +impl TsSecs for LogBuffer { + fn ts_secs(&self) -> i64 { + self.peek_back(|log| log.get_timestamp_or_fallback()) + .unwrap_or(0) + } } diff --git a/signal-gateway/src/gateway/log_handler.rs b/signal-gateway/src/gateway/log_handler.rs index e2e1083..b12f0b3 100644 --- a/signal-gateway/src/gateway/log_handler.rs +++ b/signal-gateway/src/gateway/log_handler.rs @@ -6,6 +6,7 @@ use super::{ use crate::{ assistant::{Tool, ToolExecutor, ToolResult}, concurrent_map::LazyMap, + lazy_map_cleaner::LazyMapCleaner, limiter_sequence::{Limit, LimiterSequence}, log_format::LogFormatConfig, log_message::{LogMessage, Origin}, @@ -13,7 +14,7 @@ use crate::{ use async_trait::async_trait; use chrono::Utc; use conf::Conf; -use std::fmt; +use std::{fmt, time::Duration}; use tokio::sync::mpsc::UnboundedSender; use tracing::{error, info}; @@ -57,6 +58,9 @@ pub struct LogHandlerConfig { /// Number of recent log messages to buffer per origin #[conf(long, env, default_value = "64")] pub log_buffer_size: usize, + /// Max age of an origin, after this we remove it to reclaim memory + #[conf(long, env, value_parser = conf_extra::parse_duration, default_value = "3d")] + pub max_origin_age: Duration, /// Debug logging level for suppressed messages. /// 0 = no logging, 1 = log only overall limiter, 2 = log routes + overall, 3 = log all. #[conf(long, env, default_value = "0")] @@ -86,6 +90,8 @@ pub struct LogHandler { signal_alert_mq_tx: UnboundedSender, /// Log buffers keyed by origin (app + host). Lazily created. log_buffers: LazyMap, + /// Clean up log buffers as origins get old + lazy_map_cleaner: LazyMapCleaner, /// Routes with their associated limiter sets. routes: Vec<(Route, LimiterSet)>, /// Overall rate limiters applied after route checks pass. @@ -101,17 +107,19 @@ impl LogHandler { let routes = config .routes .iter() - .map(|route| (route.clone(), route.make_limiter_set())) + .map(|route| (route.clone(), route.make_limiter_set(config.max_origin_age))) .collect(); let overall_limits = config.overall_limits.iter().collect(); let buffer_size = config.log_buffer_size; + let max_age = config.max_origin_age.as_secs() as i64; Self { config, signal_alert_mq_tx, log_buffers: LazyMap::new(move |_key| LogBuffer::new(buffer_size)), + lazy_map_cleaner: LazyMapCleaner::new(max_age), routes, overall_limits, } @@ -175,6 +183,8 @@ impl LogHandler { } } + let now = Utc::now(); + // Get or create the buffer for this origin, then record the message let formatted_text = self.log_buffers.get(&origin, |buffer| { if rate_limit_result.is_err() { @@ -186,7 +196,6 @@ impl LogHandler { let mut text = String::with_capacity(4096); let mut first_msg_len = 0; let mut is_first = true; - let now = Utc::now(); buffer.push_back_and_drain(log_msg, |log_msg| { self.config @@ -218,6 +227,10 @@ impl LogHandler { ); } } + + // Maybe cleanup old log buffers + self.lazy_map_cleaner + .maybe_clean(now.timestamp_millis(), &self.log_buffers); } /// Check if a log message passes all rate limiters. diff --git a/signal-gateway/src/gateway/rate_limiter_set.rs b/signal-gateway/src/gateway/rate_limiter_set.rs index 5a2dac8..65500a9 100644 --- a/signal-gateway/src/gateway/rate_limiter_set.rs +++ b/signal-gateway/src/gateway/rate_limiter_set.rs @@ -2,9 +2,11 @@ use crate::{ concurrent_map::LazyMap, + lazy_map_cleaner::LazyMapCleaner, limiter_sequence::LimiterSequence, log_message::{LogMessage, Origin}, }; +use std::time::Duration; /// Result of evaluating a limiter set. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -23,6 +25,8 @@ pub struct LimiterSet { limiters: LazyMap, /// Global rate limiters (shared across all origins). global_limiters: LimiterSequence, + /// Lazy map cleaner + lazy_map_cleaner: LazyMapCleaner, } impl LimiterSet { @@ -30,10 +34,12 @@ impl LimiterSet { pub fn new( make_limiters: impl Fn() -> LimiterSequence + Send + Sync + 'static, global_limiters: LimiterSequence, + max_origin_age: Duration, ) -> Self { Self { limiters: LazyMap::new(move |_key| make_limiters()), global_limiters, + lazy_map_cleaner: LazyMapCleaner::new(max_origin_age.as_secs() as i64), } } @@ -55,6 +61,9 @@ impl LimiterSet { .map(LimitResult::Limiter) }); + self.lazy_map_cleaner + .maybe_clean(log_msg.get_timestamp_or_fallback(), &self.limiters); + if let Some(result) = origin_result { return result; } diff --git a/signal-gateway/src/gateway/route.rs b/signal-gateway/src/gateway/route.rs index da465de..aa223d7 100644 --- a/signal-gateway/src/gateway/route.rs +++ b/signal-gateway/src/gateway/route.rs @@ -9,6 +9,7 @@ use crate::{ log_message::{Level, LogFilter}, }; use serde::Deserialize; +use std::time::Duration; /// A route configuration for processing log messages. /// @@ -51,10 +52,14 @@ fn default_alert_level() -> Level { impl Route { /// Create a limiter set from this route's limit configurations. - pub fn make_limiter_set(&self) -> LimiterSet { + pub fn make_limiter_set(&self, max_origin_age: Duration) -> LimiterSet { let limits = self.limits.clone(); let global_limiters = self.global_limits.iter().collect(); - LimiterSet::new(move || limits.iter().collect(), global_limiters) + LimiterSet::new( + move || limits.iter().collect(), + global_limiters, + max_origin_age, + ) } } diff --git a/signal-gateway/src/lazy_map_cleaner.rs b/signal-gateway/src/lazy_map_cleaner.rs index 4801eee..90fdddf 100644 --- a/signal-gateway/src/lazy_map_cleaner.rs +++ b/signal-gateway/src/lazy_map_cleaner.rs @@ -1,5 +1,5 @@ use crate::concurrent_map::LazyMap; -use std::{sync::Mutex}; +use std::sync::Mutex; /// Trait for types that expose a timestamp in seconds. /// For containers, this should be the oldest timestamp in the container. @@ -27,11 +27,11 @@ impl LazyMapCleaner { pub fn new(max_age: i64) -> Self { Self { max_age, - min_checkup_period: 3600, + min_checkup_period: 3600, state: Default::default(), } } - + /// Set the min checkup period (in seconds) /// Defaults to 1 hour if unset #[allow(dead_code)] @@ -45,15 +45,17 @@ impl LazyMapCleaner { /// Maybe clean a lazy map, if we have met conditions to do so pub fn maybe_clean(&self, now: i64, map: &LazyMap) - where K: Eq + std::hash::Hash + Clone, - V: TsSecs + where + K: Eq + std::hash::Hash + Clone, + V: TsSecs, { // Only lock if we can do so without contention, if someone else is checking to clean, // we don't have to do so ourselves. if let Ok(mut lk) = self.state.try_lock() { // If the last cleanup size is 0, we'll still treat it as 1 and only look to clean if we reach 2. if map.len() >= lk.last_cleanup_size.max(1) * 2 - || lk.last_checkup + self.min_checkup_period < now { + || lk.last_checkup + self.min_checkup_period < now + { let cutoff = now - self.max_age; let new_cleanup_size = map.retain(|_key, val| val.ts_secs() >= cutoff); lk.last_cleanup_size = new_cleanup_size; diff --git a/signal-gateway/src/limiter_sequence.rs b/signal-gateway/src/limiter_sequence.rs index 251de6d..507a761 100644 --- a/signal-gateway/src/limiter_sequence.rs +++ b/signal-gateway/src/limiter_sequence.rs @@ -4,6 +4,7 @@ //! and the first one to do so is returned. use crate::{ + lazy_map_cleaner::TsSecs, log_message::{LogFilter, LogMessage}, rate_limiter::{Limiter, RateThreshold}, }; @@ -65,6 +66,12 @@ impl LimiterSequence { } } +impl TsSecs for LimiterSequence { + fn ts_secs(&self) -> i64 { + self.0.iter().map(|(_f, l)| l.ts_secs()).max().unwrap_or(0) + } +} + impl<'a> FromIterator<&'a Limit> for LimiterSequence { fn from_iter(iter: T) -> Self where diff --git a/signal-gateway/src/rate_limiter.rs b/signal-gateway/src/rate_limiter.rs index b4090eb..874cf3a 100644 --- a/signal-gateway/src/rate_limiter.rs +++ b/signal-gateway/src/rate_limiter.rs @@ -1,6 +1,10 @@ //! Rate limiting for log alerts. -use crate::{concurrent_map::LazyMap, lazy_map_cleaner::{LazyMapCleaner, TsSecs}, log_message::LogMessage}; +use crate::{ + concurrent_map::LazyMap, + lazy_map_cleaner::{LazyMapCleaner, TsSecs}, + log_message::LogMessage, +}; use serde::Deserialize; use std::{ str::FromStr, @@ -127,6 +131,15 @@ impl Limiter { } } +impl TsSecs for Limiter { + fn ts_secs(&self) -> i64 { + match self { + Self::Multi(m) => m.get_latest(), + Self::SourceLocation(s) => s.get_latest(), + } + } +} + /// A rate limiter containing a single counter, and a minimum time window for the next event to pass #[allow(dead_code)] #[derive(Debug, Default)] @@ -165,6 +178,8 @@ pub struct SourceLocationRateLimiter { limiters: LazyMap<(Box, Box), MultiRateLimiter>, /// Manages cleanup of the lazy map lazy_map_cleaner: LazyMapCleaner, + /// Max evaluated timestamp + latest_ts_sec: AtomicI64, } impl SourceLocationRateLimiter { @@ -172,6 +187,7 @@ impl SourceLocationRateLimiter { Self { limiters: LazyMap::with_capacity(8, move |_key| MultiRateLimiter::from(threshold)), lazy_map_cleaner: LazyMapCleaner::new(threshold.duration.as_secs() as i64), + latest_ts_sec: Default::default(), } } @@ -184,6 +200,8 @@ impl SourceLocationRateLimiter { line.to_owned().into_boxed_str(), ); + self.latest_ts_sec.fetch_max(ts_sec, Ordering::SeqCst); + let result = self.limiters.get(&key, |limiter| limiter.evaluate(ts_sec)); // Maybe prune the lazy map @@ -191,6 +209,10 @@ impl SourceLocationRateLimiter { result } + + pub fn get_latest(&self) -> i64 { + self.latest_ts_sec.load(Ordering::SeqCst) + } } /// Implements rate-limiting criteria such as 'at least n in the last w seconds'