forked from wrenn/wrenn
Rafeed M. Bhuiyan
05ddf62399
Co-authored-by: Tasnim Kabir Sadik <tksadik@omukk.dev> Reviewed-on: wrenn/wrenn#50
71 lines
1.6 KiB
Rust
71 lines
1.6 KiB
Rust
use std::collections::HashSet;
|
|
use std::sync::Mutex;
|
|
|
|
/// Tracks active TCP connections.
|
|
pub struct ConnTracker {
|
|
inner: Mutex<ConnTrackerInner>,
|
|
}
|
|
|
|
struct ConnTrackerInner {
|
|
active: HashSet<u64>,
|
|
next_id: u64,
|
|
}
|
|
|
|
impl ConnTracker {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
inner: Mutex::new(ConnTrackerInner {
|
|
active: HashSet::new(),
|
|
next_id: 0,
|
|
}),
|
|
}
|
|
}
|
|
|
|
pub fn register_connection(&self) -> u64 {
|
|
let mut inner = self.inner.lock().unwrap();
|
|
let id = inner.next_id;
|
|
inner.next_id += 1;
|
|
inner.active.insert(id);
|
|
id
|
|
}
|
|
|
|
pub fn remove_connection(&self, id: u64) {
|
|
let mut inner = self.inner.lock().unwrap();
|
|
inner.active.remove(&id);
|
|
}
|
|
|
|
#[cfg(test)]
|
|
fn active_count(&self) -> usize {
|
|
self.inner.lock().unwrap().active.len()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn register_assigns_sequential_ids() {
|
|
let ct = ConnTracker::new();
|
|
assert_eq!(ct.register_connection(), 0);
|
|
assert_eq!(ct.register_connection(), 1);
|
|
assert_eq!(ct.register_connection(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn remove_clears_active() {
|
|
let ct = ConnTracker::new();
|
|
let id = ct.register_connection();
|
|
assert_eq!(ct.active_count(), 1);
|
|
ct.remove_connection(id);
|
|
assert_eq!(ct.active_count(), 0);
|
|
}
|
|
|
|
#[test]
|
|
fn remove_nonexistent_is_noop() {
|
|
let ct = ConnTracker::new();
|
|
ct.remove_connection(999);
|
|
assert_eq!(ct.active_count(), 0);
|
|
}
|
|
}
|