Labyricorn/signal-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rejigger concurrent map api into lazy map

Browse Source
This commit is contained in:
Chris Beck
2025-12-05 23:39:45 -07:00
parent d2d9754df0
commit 4f84cacc7b
2 changed files with 88 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
signal-gateway/src/concurrent_map.rs
+60 -4
View File
@@ -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
+28 -32
View File
@@ -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(
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