From 4f84cacc7b53a7b3b0413294a045297d1174ad99 Mon Sep 17 00:00:00 2001 From: Chris Beck Date: Fri, 5 Dec 2025 23:39:45 -0700 Subject: [PATCH] rejigger concurrent map api into lazy map --- signal-gateway/src/concurrent_map.rs | 64 +++++++++++++++++++++-- signal-gateway/src/gateway/log_handler.rs | 60 ++++++++++----------- 2 files changed, 88 insertions(+), 36 deletions(-) diff --git a/signal-gateway/src/concurrent_map.rs b/signal-gateway/src/concurrent_map.rs index 942432f..98f6dc1 100644 --- a/signal-gateway/src/concurrent_map.rs +++ b/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(&self, key: K, create: F, access: A) -> R + pub async fn get_or_insert_with(&self, key: Q, create: F, access: A) -> R where + Q: Borrow, 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 { + inner: ConcurrentMap, + factory: Box V + Send + Sync>, +} + +impl LazyMap +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(&self, key: Q, access: A) -> R + where + Q: Borrow, + 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(&self, access: A) -> R + where + A: FnOnce(&HashMap) -> R, + { + self.inner.with_read_lock(access).await + } +} + +impl std::fmt::Debug for LazyMap { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LazyMap").finish_non_exhaustive() + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/signal-gateway/src/gateway/log_handler.rs b/signal-gateway/src/gateway/log_handler.rs index 73321f2..dbfc258 100644 --- a/signal-gateway/src/gateway/log_handler.rs +++ b/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, /// Log buffers keyed by origin (app + host). Lazily created. - log_buffers: ConcurrentMap, + log_buffers: LazyMap, /// Routes with their associated limiter sets. routes: Vec<(Route, Mutex)>, /// 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( - origin.clone(), - || LogBuffer::new(buffer_size), - |buffer| { - if rate_limit_result.is_err() { - buffer.push_back(log_msg); - None - } else { - // Guess at capacity, it will be faster to use too much memory than too little - // signal-cli JVM is a hog anyways. - let mut text = String::with_capacity(4096); - let mut first_msg_len = 0; - let mut is_first = true; - let now = Utc::now(); + .get(&origin, |buffer| { + if rate_limit_result.is_err() { + buffer.push_back(log_msg); + None + } else { + // Guess at capacity, it will be faster to use too much memory than too little + // signal-cli JVM is a hog anyways. + 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 - .log_format - .write_log_msg(&mut text, log_msg, now); - if is_first { - first_msg_len = text.len(); - is_first = false; - } - }); + buffer.push_back_and_drain(log_msg, |log_msg| { + self.config + .log_format + .write_log_msg(&mut text, log_msg, now); + if is_first { + first_msg_len = text.len(); + is_first = false; + } + }); - Some((text, first_msg_len)) - } - }, - ) + Some((text, first_msg_len)) + } + }) .await; // Send alert if we have formatted text