Shard TLS full-cert budget tracking + Bound user-labeled metrics export cardinality

src/tls_front/cache.rs

@@ -1,4 +1,6 @@
 use std::collections::HashMap;
+use std::collections::hash_map::DefaultHasher;
+use std::hash::{Hash, Hasher};
 use std::net::IpAddr;
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicU64, Ordering};
@@ -15,6 +17,7 @@ use crate::tls_front::types::{
 
 const FULL_CERT_SENT_SWEEP_INTERVAL_SECS: u64 = 30;
 const FULL_CERT_SENT_MAX_IPS: usize = 65_536;
+const FULL_CERT_SENT_SHARDS: usize = 64;
 
 static FULL_CERT_SENT_IPS_GAUGE: AtomicU64 = AtomicU64::new(0);
 static FULL_CERT_SENT_CAP_DROPS: AtomicU64 = AtomicU64::new(0);
@@ -34,7 +37,7 @@ pub(crate) fn full_cert_sent_cap_drops_for_metrics() -> u64 {
 pub struct TlsFrontCache {
     memory: RwLock<HashMap<String, Arc<CachedTlsData>>>,
     default: Arc<CachedTlsData>,
-    full_cert_sent: RwLock<HashMap<IpAddr, Instant>>,
+    full_cert_sent_shards: Vec<RwLock<HashMap<IpAddr, Instant>>>,
     full_cert_sent_last_sweep_epoch_secs: AtomicU64,
     disk_path: PathBuf,
 }
@@ -70,7 +73,9 @@ impl TlsFrontCache {
         Self {
             memory: RwLock::new(map),
             default,
-            full_cert_sent: RwLock::new(HashMap::new()),
+            full_cert_sent_shards: (0..FULL_CERT_SENT_SHARDS)
+                .map(|_| RwLock::new(HashMap::new()))
+                .collect(),
             full_cert_sent_last_sweep_epoch_secs: AtomicU64::new(0),
             disk_path: disk_path.as_ref().to_path_buf(),
         }
@@ -88,6 +93,54 @@ impl TlsFrontCache {
         self.memory.read().await.contains_key(domain)
     }
 
+    fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize {
+        let mut hasher = DefaultHasher::new();
+        client_ip.hash(&mut hasher);
+        (hasher.finish() as usize) % FULL_CERT_SENT_SHARDS
+    }
+
+    fn full_cert_sent_shard(&self, client_ip: IpAddr) -> &RwLock<HashMap<IpAddr, Instant>> {
+        &self.full_cert_sent_shards[Self::full_cert_sent_shard_index(client_ip)]
+    }
+
+    fn decrement_full_cert_sent_entries(amount: usize) {
+        if amount == 0 {
+            return;
+        }
+        let amount = amount as u64;
+        let _ =
+            FULL_CERT_SENT_IPS_GAUGE.fetch_update(Ordering::AcqRel, Ordering::Relaxed, |current| {
+                Some(current.saturating_sub(amount))
+            });
+    }
+
+    fn try_reserve_full_cert_sent_entry() -> bool {
+        let mut current = FULL_CERT_SENT_IPS_GAUGE.load(Ordering::Relaxed);
+        loop {
+            if current >= FULL_CERT_SENT_MAX_IPS as u64 {
+                return false;
+            }
+            match FULL_CERT_SENT_IPS_GAUGE.compare_exchange_weak(
+                current,
+                current.saturating_add(1),
+                Ordering::AcqRel,
+                Ordering::Relaxed,
+            ) {
+                Ok(_) => return true,
+                Err(actual) => current = actual,
+            }
+        }
+    }
+
+    async fn sweep_full_cert_sent_shards(&self, now: Instant, ttl: Duration) {
+        for shard in &self.full_cert_sent_shards {
+            let mut guard = shard.write().await;
+            let before = guard.len();
+            guard.retain(|_, seen_at| now.duration_since(*seen_at) < ttl);
+            Self::decrement_full_cert_sent_entries(before.saturating_sub(guard.len()));
+        }
+    }
+
     /// Returns true when full cert payload should be sent for client_ip
     /// according to TTL policy.
     pub async fn take_full_cert_budget_for_ip(&self, client_ip: IpAddr, ttl: Duration) -> bool {
@@ -113,11 +166,11 @@ impl TlsFrontCache {
             })
             .is_ok();
 
-        let mut guard = self.full_cert_sent.write().await;
         if should_sweep {
-            guard.retain(|_, seen_at| now.duration_since(*seen_at) < ttl);
+            self.sweep_full_cert_sent_shards(now, ttl).await;
         }
 
+        let mut guard = self.full_cert_sent_shard(client_ip).write().await;
         let allowed = match guard.get_mut(&client_ip) {
             Some(seen_at) => {
                 if now.duration_since(*seen_at) >= ttl {
@@ -128,19 +181,43 @@ impl TlsFrontCache {
                 }
             }
             None => {
-                if guard.len() >= FULL_CERT_SENT_MAX_IPS {
+                if !Self::try_reserve_full_cert_sent_entry() {
                     FULL_CERT_SENT_CAP_DROPS.fetch_add(1, Ordering::Relaxed);
-                    FULL_CERT_SENT_IPS_GAUGE.store(guard.len() as u64, Ordering::Relaxed);
                     return false;
                 }
                 guard.insert(client_ip, now);
                 true
             }
         };
-        FULL_CERT_SENT_IPS_GAUGE.store(guard.len() as u64, Ordering::Relaxed);
         allowed
     }
 
+    #[cfg(test)]
+    async fn insert_full_cert_sent_for_tests(&self, client_ip: IpAddr, seen_at: Instant) {
+        let mut guard = self.full_cert_sent_shard(client_ip).write().await;
+        if guard.insert(client_ip, seen_at).is_none() {
+            FULL_CERT_SENT_IPS_GAUGE.fetch_add(1, Ordering::Relaxed);
+        }
+    }
+
+    #[cfg(test)]
+    async fn full_cert_sent_is_empty_for_tests(&self) -> bool {
+        for shard in &self.full_cert_sent_shards {
+            if !shard.read().await.is_empty() {
+                return false;
+            }
+        }
+        true
+    }
+
+    #[cfg(test)]
+    async fn full_cert_sent_contains_for_tests(&self, client_ip: IpAddr) -> bool {
+        self.full_cert_sent_shard(client_ip)
+            .read()
+            .await
+            .contains_key(&client_ip)
+    }
+
     pub async fn set(&self, domain: &str, data: CachedTlsData) {
         let mut guard = self.memory.write().await;
         guard.insert(domain.to_string(), Arc::new(data));
@@ -381,7 +458,7 @@ mod tests {
             assert!(cache.take_full_cert_budget_for_ip(ip, ttl).await);
         }
 
-        assert!(cache.full_cert_sent.read().await.is_empty());
+        assert!(cache.full_cert_sent_is_empty_for_tests().await);
     }
 
     #[tokio::test]
@@ -395,19 +472,16 @@ mod tests {
             .unwrap_or_else(Instant::now);
 
         cache
-            .full_cert_sent
-            .write()
-            .await
-            .insert(stale_ip, stale_seen_at);
+            .insert_full_cert_sent_for_tests(stale_ip, stale_seen_at)
+            .await;
         cache
             .full_cert_sent_last_sweep_epoch_secs
             .store(0, Ordering::Relaxed);
 
         assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await);
 
-        let guard = cache.full_cert_sent.read().await;
-        assert!(!guard.contains_key(&stale_ip));
-        assert!(guard.contains_key(&new_ip));
+        assert!(!cache.full_cert_sent_contains_for_tests(stale_ip).await);
+        assert!(cache.full_cert_sent_contains_for_tests(new_ip).await);
     }
 
     #[tokio::test]
@@ -425,18 +499,15 @@ mod tests {
             .as_secs();
 
         cache
-            .full_cert_sent
-            .write()
-            .await
-            .insert(stale_ip, stale_seen_at);
+            .insert_full_cert_sent_for_tests(stale_ip, stale_seen_at)
+            .await;
         cache
             .full_cert_sent_last_sweep_epoch_secs
             .store(now_epoch_secs, Ordering::Relaxed);
 
         assert!(cache.take_full_cert_budget_for_ip(new_ip, ttl).await);
 
-        let guard = cache.full_cert_sent.read().await;
-        assert!(guard.contains_key(&stale_ip));
-        assert!(guard.contains_key(&new_ip));
+        assert!(cache.full_cert_sent_contains_for_tests(stale_ip).await);
+        assert!(cache.full_cert_sent_contains_for_tests(new_ip).await);
     }
 }