systematically use lazy map cleaner on lazy maps

signal-gateway/src/circular_buffer.rs

@@ -38,6 +38,11 @@ impl<T> CircularBuffer<T> {
         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 {

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<R>(&self, f: impl FnOnce(&LogMessage) -> R) -> Option<R> {
+        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)
+    }
 }

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<SignalAlertMessage>,
     /// Log buffers keyed by origin (app + host). Lazily created.
     log_buffers: LazyMap<Origin, LogBuffer>,
+    /// 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.

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<Origin, LimiterSequence>,
     /// 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;
         }

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,
+        )
     }
 }
 

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<K, V>(&self, now: i64, map: &LazyMap<K, V>)
-        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;

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<T>(iter: T) -> Self
     where

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<str>, Box<str>), 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'