rejigger concurrent map api into lazy map

signal-gateway/src/concurrent_map.rs

@@ -4,7 +4,8 @@
 //! where values already exist, using a read lock first before falling back to a
 //! write lock for insertions.
 
-use std::{collections::HashMap, hash::Hash};
+use std::{borrow::Borrow, collections::HashMap, hash::Hash};
+
 use tokio::sync::RwLock;
 
 /// A concurrent hash map that uses read-preferring locking.
@@ -35,18 +36,23 @@ where
     /// 2. If found, calls `access` with a reference to the value
     /// 3. If not found, acquires a write lock, inserts using `create`, then calls `access`
     ///
+    /// The key is only cloned when a new value needs to be inserted.
+    ///
     /// The lock is held while `access` runs, so `access` can safely use the reference.
     /// For async operations on the value, consider having `access` return a future
     /// that owns any data it needs.
-    pub async fn get_or_insert_with<R, F, A>(&self, key: K, create: F, access: A) -> R
+    pub async fn get_or_insert_with<Q, R, F, A>(&self, key: Q, create: F, access: A) -> R
     where
+        Q: Borrow<K>,
         F: FnOnce() -> V,
         A: FnOnce(&V) -> R,
     {
+        let key = key.borrow();
+
         // Try to get existing value with read lock first
         {
             let guard = self.inner.read().await;
-            if let Some(value) = guard.get(&key) {
+            if let Some(value) = guard.get(key) {
                 return access(value);
             }
         }
@@ -54,7 +60,7 @@ where
         // Value doesn't exist, need to create with write lock
         let mut guard = self.inner.write().await;
         // Use entry API - handles the race where another task inserted while we waited
-        let value = guard.entry(key).or_insert_with(create);
+        let value = guard.entry(key.clone()).or_insert_with(create);
         access(value)
     }
 
@@ -80,6 +86,56 @@ where
     }
 }
 
+/// A concurrent hash map with a default factory for creating new values.
+///
+/// Wraps a [`ConcurrentMap`] and stores a factory closure, so callers don't need
+/// to pass the creation function on every access.
+pub struct LazyMap<K, V> {
+    inner: ConcurrentMap<K, V>,
+    factory: Box<dyn Fn() -> V + Send + Sync>,
+}
+
+impl<K, V> LazyMap<K, V>
+where
+    K: Eq + Hash + Clone,
+{
+    /// Create a new lazy map with the given factory for creating values.
+    pub fn new(factory: impl Fn() -> V + Send + Sync + 'static) -> Self {
+        Self {
+            inner: ConcurrentMap::new(),
+            factory: Box::new(factory),
+        }
+    }
+
+    /// Get a value, creating it with the factory if it doesn't exist.
+    ///
+    /// Uses the factory provided at construction time to create new values.
+    /// The key is only cloned when a new value needs to be inserted.
+    pub async fn get<Q, R, A>(&self, key: Q, access: A) -> R
+    where
+        Q: Borrow<K>,
+        A: FnOnce(&V) -> R,
+    {
+        self.inner
+            .get_or_insert_with(key, &self.factory, access)
+            .await
+    }
+
+    /// Access all entries in the map with a read lock.
+    pub async fn with_read_lock<R, A>(&self, access: A) -> R
+    where
+        A: FnOnce(&HashMap<K, V>) -> R,
+    {
+        self.inner.with_read_lock(access).await
+    }
+}
+
+impl<K, V> std::fmt::Debug for LazyMap<K, V> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LazyMap").finish_non_exhaustive()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

signal-gateway/src/gateway/log_handler.rs

@@ -4,7 +4,7 @@ use super::{
     route::{Destination, Limit, Route},
 };
 use crate::{
-    concurrent_map::ConcurrentMap,
+    concurrent_map::LazyMap,
     log_format::LogFormatConfig,
     log_message::{LogMessage, Origin},
 };
@@ -78,7 +78,7 @@ pub struct LogHandler {
     config: LogHandlerConfig,
     signal_alert_mq_tx: UnboundedSender<SignalAlertMessage>,
     /// Log buffers keyed by origin (app + host). Lazily created.
-    log_buffers: ConcurrentMap<Origin, LogBuffer>,
+    log_buffers: LazyMap<Origin, LogBuffer>,
     /// Routes with their associated limiter sets.
     routes: Vec<(Route, Mutex<LimiterSet>)>,
     /// Overall rate limiters applied after route checks pass.
@@ -103,10 +103,12 @@ impl LogHandler {
             .map(|limit| Mutex::new(limit.make_limiter()))
             .collect();
 
+        let buffer_size = config.log_buffer_size;
+
         Self {
             config,
             signal_alert_mq_tx,
-            log_buffers: ConcurrentMap::new(),
+            log_buffers: LazyMap::new(move || LogBuffer::new(buffer_size)),
             routes,
             overall_limits,
         }
@@ -168,40 +170,34 @@ impl LogHandler {
             info!("Suppressed {sev} ({reason}):\n{}", log_msg.msg);
         }
 
-        let buffer_size = self.config.log_buffer_size;
-
         // Get or create the buffer for this origin, then record the message
         let formatted_text = self
             .log_buffers
-            .get_or_insert_with(
+            .get(&origin, |buffer| {
-                origin.clone(),
+                if rate_limit_result.is_err() {
-                || LogBuffer::new(buffer_size),
+                    buffer.push_back(log_msg);
-                |buffer| {
+                    None
-                    if rate_limit_result.is_err() {
+                } else {
-                        buffer.push_back(log_msg);
+                    // Guess at capacity, it will be faster to use too much memory than too little
-                        None
+                    // signal-cli JVM is a hog anyways.
-                    } else {
+                    let mut text = String::with_capacity(4096);
-                        // Guess at capacity, it will be faster to use too much memory than too little
+                    let mut first_msg_len = 0;
-                        // signal-cli JVM is a hog anyways.
+                    let mut is_first = true;
-                        let mut text = String::with_capacity(4096);
+                    let now = Utc::now();
-                        let mut first_msg_len = 0;
-                        let mut is_first = true;
-                        let now = Utc::now();
 
-                        buffer.push_back_and_drain(log_msg, |log_msg| {
+                    buffer.push_back_and_drain(log_msg, |log_msg| {
-                            self.config
+                        self.config
-                                .log_format
+                            .log_format
-                                .write_log_msg(&mut text, log_msg, now);
+                            .write_log_msg(&mut text, log_msg, now);
-                            if is_first {
+                        if is_first {
-                                first_msg_len = text.len();
+                            first_msg_len = text.len();
-                                is_first = false;
+                            is_first = false;
-                            }
+                        }
-                        });
+                    });
 
-                        Some((text, first_msg_len))
+                    Some((text, first_msg_len))
-                    }
+                }
-                },
+            })
-            )
             .await;
 
         // Send alert if we have formatted text