1.1.0.0

src/stats/mod.rs

@@ -4,9 +4,11 @@ use std::sync::atomic::{AtomicU64, Ordering};
 use std::sync::Arc;
 use std::time::Instant;
 use dashmap::DashMap;
-use parking_lot::RwLock;
+use parking_lot::{RwLock, Mutex};
 use lru::LruCache;
 use std::num::NonZeroUsize;
+use std::hash::{Hash, Hasher};
+use std::collections::hash_map::DefaultHasher;
 
 /// Thread-safe statistics
 #[derive(Default)]
@@ -141,37 +143,57 @@ impl Stats {
     }
 }
 
-// Arc<Stats> Hightech Stats :D
-
-/// Replay attack checker using LRU cache
+/// Sharded Replay attack checker using LRU cache
+/// Uses multiple independent LRU caches to reduce lock contention
 pub struct ReplayChecker {
-    handshakes: RwLock<LruCache<Vec<u8>, ()>>,
-    tls_digests: RwLock<LruCache<Vec<u8>, ()>>,
+    shards: Vec<Mutex<LruCache<Vec<u8>, ()>>>,
+    shard_mask: usize,
 }
 
 impl ReplayChecker {
-    pub fn new(capacity: usize) -> Self {
-        let cap = NonZeroUsize::new(capacity.max(1)).unwrap();
+    /// Create new replay checker with specified capacity per shard
+    /// Total capacity = capacity * num_shards
+    pub fn new(total_capacity: usize) -> Self {
+        // Use 64 shards for good concurrency
+        let num_shards = 64;
+        let shard_capacity = (total_capacity / num_shards).max(1);
+        let cap = NonZeroUsize::new(shard_capacity).unwrap();
+        
+        let mut shards = Vec::with_capacity(num_shards);
+        for _ in 0..num_shards {
+            shards.push(Mutex::new(LruCache::new(cap)));
+        }
+        
         Self {
-            handshakes: RwLock::new(LruCache::new(cap)),
-            tls_digests: RwLock::new(LruCache::new(cap)),
+            shards,
+            shard_mask: num_shards - 1,
         }
     }
     
+    fn get_shard(&self, key: &[u8]) -> usize {
+        let mut hasher = DefaultHasher::new();
+        key.hash(&mut hasher);
+        (hasher.finish() as usize) & self.shard_mask
+    }
+    
     pub fn check_handshake(&self, data: &[u8]) -> bool {
-        self.handshakes.read().contains(&data.to_vec())
+        let shard_idx = self.get_shard(data);
+        self.shards[shard_idx].lock().contains(&data.to_vec())
     }
     
     pub fn add_handshake(&self, data: &[u8]) {
-        self.handshakes.write().put(data.to_vec(), ());
+        let shard_idx = self.get_shard(data);
+        self.shards[shard_idx].lock().put(data.to_vec(), ());
     }
     
     pub fn check_tls_digest(&self, data: &[u8]) -> bool {
-        self.tls_digests.read().contains(&data.to_vec())
+        let shard_idx = self.get_shard(data);
+        self.shards[shard_idx].lock().contains(&data.to_vec())
     }
     
     pub fn add_tls_digest(&self, data: &[u8]) {
-        self.tls_digests.write().put(data.to_vec(), ());
+        let shard_idx = self.get_shard(data);
+        self.shards[shard_idx].lock().put(data.to_vec(), ());
     }
 }
 
@@ -183,7 +205,6 @@ mod tests {
     fn test_stats_shared_counters() {
         let stats = Arc::new(Stats::new());
         
-        // Симулируем использование из разных "задач"
         let stats1 = Arc::clone(&stats);
         let stats2 = Arc::clone(&stats);
         
@@ -191,33 +212,20 @@ mod tests {
         stats2.increment_connects_all();
         stats1.increment_connects_all();
         
-        // Все инкременты должны быть видны
         assert_eq!(stats.get_connects_all(), 3);
     }
     
     #[test]
-    fn test_user_stats_shared() {
-        let stats = Arc::new(Stats::new());
+    fn test_replay_checker_sharding() {
+        let checker = ReplayChecker::new(100);
+        let data1 = b"test1";
+        let data2 = b"test2";
         
-        let stats1 = Arc::clone(&stats);
-        let stats2 = Arc::clone(&stats);
+        checker.add_handshake(data1);
+        assert!(checker.check_handshake(data1));
+        assert!(!checker.check_handshake(data2));
         
-        stats1.add_user_octets_from("user1", 100);
-        stats2.add_user_octets_from("user1", 200);
-        stats1.add_user_octets_to("user1", 50);
-        
-        assert_eq!(stats.get_user_total_octets("user1"), 350);
-    }
-    
-    #[test]
-    fn test_concurrent_user_connects() {
-        let stats = Arc::new(Stats::new());
-        
-        stats.increment_user_curr_connects("user1");
-        stats.increment_user_curr_connects("user1");
-        assert_eq!(stats.get_user_curr_connects("user1"), 2);
-        
-        stats.decrement_user_curr_connects("user1");
-        assert_eq!(stats.get_user_curr_connects("user1"), 1);
+        checker.add_handshake(data2);
+        assert!(checker.check_handshake(data2));
     }
 }