From 2306364bbb80145dc20924f9e6370dca137350b5 Mon Sep 17 00:00:00 2001 From: Chris Beck Date: Mon, 15 Dec 2025 13:32:19 -0700 Subject: [PATCH] add and use lazy map cleaner --- signal-gateway/src/lazy_map_cleaner.rs | 64 ++++++++++++++++++++++++++ signal-gateway/src/lib.rs | 1 + signal-gateway/src/rate_limiter.rs | 39 ++++++---------- 3 files changed, 78 insertions(+), 26 deletions(-) create mode 100644 signal-gateway/src/lazy_map_cleaner.rs diff --git a/signal-gateway/src/lazy_map_cleaner.rs b/signal-gateway/src/lazy_map_cleaner.rs new file mode 100644 index 0000000..4801eee --- /dev/null +++ b/signal-gateway/src/lazy_map_cleaner.rs @@ -0,0 +1,64 @@ +use crate::concurrent_map::LazyMap; +use std::{sync::Mutex}; + +/// Trait for types that expose a timestamp in seconds. +/// For containers, this should be the oldest timestamp in the container. +pub trait TsSecs { + fn ts_secs(&self) -> i64; +} + +#[derive(Debug, Default)] +struct CleanerState { + last_cleanup_size: usize, + last_checkup: i64, +} + +/// A policy object that can cleanup entries from a LazyMap, if the value type has timestamps. +/// The lazy map is cleaned whenever it doubles in size, or whenever enough time has passed +/// since the last cleanup. This can prevent lazy maps from growing without bound. +pub struct LazyMapCleaner { + max_age: i64, + min_checkup_period: i64, + state: Mutex, +} + +impl LazyMapCleaner { + /// Create a new lazy map cleaner, with given max age of items (in seconds) + pub fn new(max_age: i64) -> Self { + Self { + max_age, + min_checkup_period: 3600, + state: Default::default(), + } + } + + /// Set the min checkup period (in seconds) + /// Defaults to 1 hour if unset + #[allow(dead_code)] + pub fn with_min_checkup_period(self, min_checkup_period: i64) -> Self { + Self { + max_age: self.max_age, + min_checkup_period, + state: self.state, + } + } + + /// 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 + { + // 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 { + 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; + lk.last_checkup = now; + } + } + } +} diff --git a/signal-gateway/src/lib.rs b/signal-gateway/src/lib.rs index 1d3681d..92715ec 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 lazy_map_cleaner; pub(crate) mod limiter_sequence; pub(crate) mod log_format; pub(crate) mod log_message; diff --git a/signal-gateway/src/rate_limiter.rs b/signal-gateway/src/rate_limiter.rs index 0fd1e2d..b4090eb 100644 --- a/signal-gateway/src/rate_limiter.rs +++ b/signal-gateway/src/rate_limiter.rs @@ -1,6 +1,6 @@ //! Rate limiting for log alerts. -use crate::{concurrent_map::LazyMap, log_message::LogMessage}; +use crate::{concurrent_map::LazyMap, lazy_map_cleaner::{LazyMapCleaner, TsSecs}, log_message::LogMessage}; use serde::Deserialize; use std::{ str::FromStr, @@ -163,18 +163,15 @@ impl SimpleRateLimiter { pub struct SourceLocationRateLimiter { /// Maps (file, line) -> rate limiter for that location limiters: LazyMap<(Box, Box), MultiRateLimiter>, - /// Threshold for creating new limiters (used for cleanup window calculation) - threshold: RateThreshold, - /// Maximum entries before triggering cleanup - last_cleanup_size: Mutex, + /// Manages cleanup of the lazy map + lazy_map_cleaner: LazyMapCleaner, } impl SourceLocationRateLimiter { pub fn new(threshold: RateThreshold) -> Self { Self { - limiters: LazyMap::with_capacity(16, move |_key| MultiRateLimiter::from(threshold)), - threshold, - last_cleanup_size: Mutex::new(8), + limiters: LazyMap::with_capacity(8, move |_key| MultiRateLimiter::from(threshold)), + lazy_map_cleaner: LazyMapCleaner::new(threshold.duration.as_secs() as i64), } } @@ -189,27 +186,11 @@ impl SourceLocationRateLimiter { let result = self.limiters.get(&key, |limiter| limiter.evaluate(ts_sec)); - // Opportunistically check for cleanup, but not if someone else is cleaning up - if let Ok(mut lk) = self.last_cleanup_size.try_lock() { - // If we're twice as large as the ending count from the last time we - // cleaned up, then let's try to clean up again. - if self.limiters.len() >= 2 * *lk { - *lk = self.cleanup(ts_sec); - } - } + // Maybe prune the lazy map + self.lazy_map_cleaner.maybe_clean(ts_sec, &self.limiters); result } - - /// Remove entries where all timestamps are older than the window - fn cleanup(&self, now: i64) -> usize { - let window = self.threshold.duration.as_secs() as i64; - let cutoff = now - window; - self.limiters.retain(|_, limiter| { - // Keep if any timestamp is recent enough - limiter.get_latest() > cutoff - }) - } } /// Implements rate-limiting criteria such as 'at least n in the last w seconds' @@ -258,6 +239,12 @@ impl MultiRateLimiter { } } +impl TsSecs for MultiRateLimiter { + fn ts_secs(&self) -> i64 { + self.get_latest() + } +} + struct MultiRateLimiterInner { /// Records the last n events as a ring buffer. /// Initialized to 0, representing "very distant past".