From 4584ae23bd0b2a6cb55cf0b7196d598533d95633 Mon Sep 17 00:00:00 2001
From: Chris Beck <beck.ct@gmail.com>
Date: Mon, 15 Dec 2025 10:26:49 -0700
Subject: [PATCH] make concurrent map able to give default values which depend
 on key

this allows to uphold more complex invariants
---
 signal-gateway/src/concurrent_map.rs           | 13 ++++++++++---
 signal-gateway/src/gateway/log_handler.rs      |  2 +-
 signal-gateway/src/gateway/rate_limiter_set.rs |  2 +-
 signal-gateway/src/rate_limiter.rs             |  2 +-
 4 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/signal-gateway/src/concurrent_map.rs b/signal-gateway/src/concurrent_map.rs
index c8a1432..b4ff4c0 100644
--- a/signal-gateway/src/concurrent_map.rs
+++ b/signal-gateway/src/concurrent_map.rs
@@ -11,6 +11,12 @@
 //! sources, and insertions occur only a few times at the beginning of the process.
 //! Then almost all accesses are to existing elements, and from that point on,
 //! only read locks are taken when using this API.
+//!
+//! The API also allows to call "retain" if the map gets too big and it needs
+//! to be pruned in some manner.
+//!
+//! This is used instead of dash_map and once_map to avoid unnecessary complexity
+//! and dependencies, and give exactly the API needed in our application.
 
 use std::{borrow::Borrow, collections::HashMap, hash::Hash, sync::RwLock};
 
@@ -98,7 +104,7 @@ impl <K, V> Default for ConcurrentMap<K, V>
 /// to pass the creation function on every access.
 pub struct LazyMap<K, V> {
     inner: ConcurrentMap<K, V>,
-    factory: Box<dyn Fn() -> V + Send + Sync>,
+    factory: Box<dyn Fn(&K) -> V + Send + Sync>,
 }
 
 impl<K, V> LazyMap<K, V>
@@ -106,7 +112,7 @@ where
     K: Eq + Hash + Clone,
 {
     /// Create a new lazy map with the given factory for creating initial values.
-    pub fn new(factory: impl Fn() -> V + Send + Sync + 'static) -> Self {
+    pub fn new(factory: impl Fn(&K) -> V + Send + Sync + 'static) -> Self {
         Self {
             inner: ConcurrentMap::new(),
             factory: Box::new(factory),
@@ -119,7 +125,8 @@ where
         Q: Borrow<K>,
         A: FnOnce(&V) -> R,
     {
-        self.inner.get_or_insert_with(key, &self.factory, access)
+        let key = key.borrow();
+        self.inner.get_or_insert_with(key, || (self.factory)(key), access)
     }
 
     /// Access all entries in the map with a read lock.
diff --git a/signal-gateway/src/gateway/log_handler.rs b/signal-gateway/src/gateway/log_handler.rs
index 5bb3e5a..ced7ea3 100644
--- a/signal-gateway/src/gateway/log_handler.rs
+++ b/signal-gateway/src/gateway/log_handler.rs
@@ -115,7 +115,7 @@ impl LogHandler {
         Self {
             config,
             signal_alert_mq_tx,
-            log_buffers: LazyMap::new(move || LogBuffer::new(buffer_size)),
+            log_buffers: LazyMap::new(move |_key| LogBuffer::new(buffer_size)),
             routes,
             overall_limits,
         }
diff --git a/signal-gateway/src/gateway/rate_limiter_set.rs b/signal-gateway/src/gateway/rate_limiter_set.rs
index f9c679f..af98d4c 100644
--- a/signal-gateway/src/gateway/rate_limiter_set.rs
+++ b/signal-gateway/src/gateway/rate_limiter_set.rs
@@ -36,7 +36,7 @@ impl LimiterSet {
         global_limiters: Vec<(LogFilter, Limiter)>,
     ) -> Self {
         Self {
-            limiters: LazyMap::new(make_limiters),
+            limiters: LazyMap::new(move |_key| make_limiters()),
             global_limiters,
         }
     }
diff --git a/signal-gateway/src/rate_limiter.rs b/signal-gateway/src/rate_limiter.rs
index 175e7ae..e7dee0a 100644
--- a/signal-gateway/src/rate_limiter.rs
+++ b/signal-gateway/src/rate_limiter.rs
@@ -178,7 +178,7 @@ pub struct SourceLocationRateLimiter {
 impl SourceLocationRateLimiter {
     pub fn new(threshold: RateThreshold, max_entries: usize) -> Self {
         Self {
-            limiters: LazyMap::new(move || MultiRateLimiter::from(threshold)),
+            limiters: LazyMap::new(move |_key| MultiRateLimiter::from(threshold)),
             threshold,
             max_entries,
         }