Decomposing hot-path modules into focused submodules

Signed-off-by: Alexey <247128645+axkurcom@users.noreply.github.com>

src/ip_tracker/snapshot.rs

@@ -0,0 +1,309 @@
+use super::*;
+
+impl UserIpTracker {
+    pub(super) async fn maybe_compact_empty_users(&self) {
+        const COMPACT_INTERVAL_SECS: u64 = 60;
+        let now_epoch_secs = Self::now_epoch_secs();
+        let last_compact_epoch_secs = self.last_compact_epoch_secs.load(Ordering::Relaxed);
+        if now_epoch_secs.saturating_sub(last_compact_epoch_secs) < COMPACT_INTERVAL_SECS {
+            return;
+        }
+        if self
+            .last_compact_epoch_secs
+            .compare_exchange(
+                last_compact_epoch_secs,
+                now_epoch_secs,
+                Ordering::AcqRel,
+                Ordering::Relaxed,
+            )
+            .is_err()
+        {
+            return;
+        }
+
+        let window = self.limit_window();
+        let now = Instant::now();
+        for shard_lock in self.shards.iter() {
+            let mut shard = shard_lock.write().await;
+            let mut pruned_recent_entries = 0usize;
+            for user_recent in shard.recent_ips.values_mut() {
+                pruned_recent_entries = pruned_recent_entries.saturating_add(Self::prune_recent(
+                    user_recent,
+                    now,
+                    window,
+                ));
+            }
+            Self::decrement_counter(&self.recent_entry_count, pruned_recent_entries);
+
+            let mut users = Vec::<String>::with_capacity(
+                shard
+                    .active_ips
+                    .len()
+                    .saturating_add(shard.recent_ips.len()),
+            );
+            users.extend(shard.active_ips.keys().cloned());
+            for user in shard.recent_ips.keys() {
+                if !shard.active_ips.contains_key(user) {
+                    users.push(user.clone());
+                }
+            }
+
+            for user in users {
+                let active_empty = shard
+                    .active_ips
+                    .get(&user)
+                    .map(|ips| ips.is_empty())
+                    .unwrap_or(true);
+                let recent_empty = shard
+                    .recent_ips
+                    .get(&user)
+                    .map(|ips| ips.is_empty())
+                    .unwrap_or(true);
+                if active_empty && recent_empty {
+                    shard.active_ips.remove(&user);
+                    shard.recent_ips.remove(&user);
+                }
+            }
+        }
+    }
+
+    pub async fn run_periodic_maintenance(self: Arc<Self>) {
+        let mut interval = tokio::time::interval(Duration::from_secs(1));
+        loop {
+            interval.tick().await;
+            self.drain_cleanup_queue().await;
+            self.maybe_compact_empty_users().await;
+        }
+    }
+
+    pub async fn memory_stats(&self) -> UserIpTrackerMemoryStats {
+        let cleanup_queue_len = self.cleanup_queue_len.load(Ordering::Relaxed) as usize;
+        let mut active_users = 0usize;
+        let mut recent_users = 0usize;
+        let mut active_entries = 0usize;
+        let mut recent_entries = 0usize;
+        for shard_lock in self.shards.iter() {
+            let shard = shard_lock.read().await;
+            active_users = active_users.saturating_add(shard.active_ips.len());
+            recent_users = recent_users.saturating_add(shard.recent_ips.len());
+            active_entries =
+                active_entries.saturating_add(shard.active_ips.values().map(HashMap::len).sum());
+            recent_entries =
+                recent_entries.saturating_add(shard.recent_ips.values().map(HashMap::len).sum());
+        }
+
+        UserIpTrackerMemoryStats {
+            active_users,
+            recent_users,
+            active_entries,
+            recent_entries,
+            cleanup_queue_len,
+            active_cap_rejects: self.active_cap_rejects.load(Ordering::Relaxed),
+            recent_cap_rejects: self.recent_cap_rejects.load(Ordering::Relaxed),
+            cleanup_deferred_releases: self.cleanup_deferred_releases.load(Ordering::Relaxed),
+        }
+    }
+
+    pub async fn get_recent_counts_for_users(&self, users: &[String]) -> HashMap<String, usize> {
+        self.drain_cleanup_queue().await;
+        self.get_recent_counts_for_users_snapshot(users).await
+    }
+
+    pub(crate) async fn get_recent_counts_for_users_snapshot(
+        &self,
+        users: &[String],
+    ) -> HashMap<String, usize> {
+        let window = self.limit_window();
+        let now = Instant::now();
+
+        let mut counts = HashMap::with_capacity(users.len());
+        for user in users {
+            let shard_idx = Self::shard_idx(user);
+            let shard = self.shards[shard_idx].read().await;
+            let count = if let Some(user_recent) = shard.recent_ips.get(user) {
+                user_recent
+                    .values()
+                    .filter(|seen_at| now.duration_since(**seen_at) <= window)
+                    .count()
+            } else {
+                0
+            };
+            counts.insert(user.clone(), count);
+        }
+        counts
+    }
+
+    pub async fn get_active_ips_for_users(&self, users: &[String]) -> HashMap<String, Vec<IpAddr>> {
+        self.drain_cleanup_queue().await;
+        let mut out = HashMap::with_capacity(users.len());
+        for user in users {
+            let shard_idx = Self::shard_idx(user);
+            let shard = self.shards[shard_idx].read().await;
+            let mut ips = shard
+                .active_ips
+                .get(user)
+                .map(|per_ip| per_ip.keys().copied().collect::<Vec<_>>())
+                .unwrap_or_else(Vec::new);
+            ips.sort();
+            out.insert(user.clone(), ips);
+        }
+        out
+    }
+
+    pub async fn get_recent_ips_for_users(&self, users: &[String]) -> HashMap<String, Vec<IpAddr>> {
+        self.drain_cleanup_queue().await;
+        let window = self.limit_window();
+        let now = Instant::now();
+
+        let mut out = HashMap::with_capacity(users.len());
+        for user in users {
+            let shard_idx = Self::shard_idx(user);
+            let shard = self.shards[shard_idx].read().await;
+            let mut ips = if let Some(user_recent) = shard.recent_ips.get(user) {
+                user_recent
+                    .iter()
+                    .filter(|(_, seen_at)| now.duration_since(**seen_at) <= window)
+                    .map(|(ip, _)| *ip)
+                    .collect::<Vec<_>>()
+            } else {
+                Vec::new()
+            };
+            ips.sort();
+            out.insert(user.clone(), ips);
+        }
+        out
+    }
+
+    pub async fn get_active_ip_count(&self, username: &str) -> usize {
+        self.drain_cleanup_queue().await;
+        let shard_idx = Self::shard_idx(username);
+        let shard = self.shards[shard_idx].read().await;
+        shard
+            .active_ips
+            .get(username)
+            .map(|ips| ips.len())
+            .unwrap_or(0)
+    }
+
+    pub async fn get_active_ips(&self, username: &str) -> Vec<IpAddr> {
+        self.drain_cleanup_queue().await;
+        let shard_idx = Self::shard_idx(username);
+        let shard = self.shards[shard_idx].read().await;
+        shard
+            .active_ips
+            .get(username)
+            .map(|ips| ips.keys().copied().collect())
+            .unwrap_or_else(Vec::new)
+    }
+
+    pub async fn get_stats(&self) -> Vec<(String, usize, usize)> {
+        self.drain_cleanup_queue().await;
+        self.get_stats_snapshot().await
+    }
+
+    pub(crate) async fn get_stats_snapshot(&self) -> Vec<(String, usize, usize)> {
+        let mut active_counts = Vec::new();
+        for shard_lock in self.shards.iter() {
+            let shard = shard_lock.read().await;
+            active_counts.extend(
+                shard
+                    .active_ips
+                    .iter()
+                    .map(|(username, user_ips)| (username.clone(), user_ips.len())),
+            );
+        }
+
+        let mut stats = Vec::with_capacity(active_counts.len());
+        for (username, active_count) in active_counts {
+            let limit = self.user_limit(&username).unwrap_or(0);
+            stats.push((username, active_count, limit));
+        }
+
+        stats.sort_by(|a, b| a.0.cmp(&b.0));
+        stats
+    }
+
+    pub async fn clear_user_ips(&self, username: &str) {
+        let shard_idx = Self::shard_idx(username);
+        let mut shard = self.shards[shard_idx].write().await;
+        let removed_active_entries = shard
+            .active_ips
+            .remove(username)
+            .map(|ips| ips.len())
+            .unwrap_or(0);
+        Self::decrement_counter(&self.active_entry_count, removed_active_entries);
+
+        let removed_recent_entries = shard
+            .recent_ips
+            .remove(username)
+            .map(|ips| ips.len())
+            .unwrap_or(0);
+        Self::decrement_counter(&self.recent_entry_count, removed_recent_entries);
+    }
+
+    pub async fn clear_all(&self) {
+        for shard_lock in self.shards.iter() {
+            let mut shard = shard_lock.write().await;
+            shard.active_ips.clear();
+            shard.recent_ips.clear();
+        }
+        self.active_entry_count.store(0, Ordering::Relaxed);
+        self.recent_entry_count.store(0, Ordering::Relaxed);
+        for cleanup_shard in self.cleanup_shards.iter() {
+            match cleanup_shard.queue.lock() {
+                Ok(mut queue) => queue.clear(),
+                Err(poisoned) => {
+                    poisoned.into_inner().clear();
+                    cleanup_shard.queue.clear_poison();
+                }
+            }
+        }
+        self.cleanup_queue_len.store(0, Ordering::Relaxed);
+    }
+
+    pub async fn is_ip_active(&self, username: &str, ip: IpAddr) -> bool {
+        self.drain_cleanup_queue().await;
+        let shard_idx = Self::shard_idx(username);
+        let shard = self.shards[shard_idx].read().await;
+        shard
+            .active_ips
+            .get(username)
+            .map(|ips| ips.contains_key(&ip))
+            .unwrap_or(false)
+    }
+
+    pub async fn get_user_limit(&self, username: &str) -> Option<usize> {
+        self.user_limit(username)
+    }
+
+    pub async fn format_stats(&self) -> String {
+        let stats = self.get_stats().await;
+
+        if stats.is_empty() {
+            return String::from("No active users");
+        }
+
+        let mut output = String::from("User IP Statistics:\n");
+        output.push_str("==================\n");
+
+        for (username, active_count, limit) in stats {
+            output.push_str(&format!(
+                "User: {:<20} Active IPs: {}/{}\n",
+                username,
+                active_count,
+                if limit > 0 {
+                    limit.to_string()
+                } else {
+                    "unlimited".to_string()
+                }
+            ));
+
+            let ips = self.get_active_ips(&username).await;
+            for ip in ips {
+                output.push_str(&format!("  - {}\n", ip));
+            }
+        }
+
+        output
+    }
+}