use std::time::Duration;
use dashmap::DashMap;
use crate::now_ms;
#[derive(Debug, Default)]
pub struct DedupeStore {
inner: DashMap<(String, String), i64>,
}
impl DedupeStore {
pub fn new() -> Self {
Self::default()
}
pub fn check_and_record(&self, topic: &str, key: &str, window: Duration) -> bool {
let now = now_ms();
let expires = now + window.as_millis() as i64;
let k = (topic.to_string(), key.to_string());
let mut is_fresh = false;
self.inner
.entry(k)
.and_modify(|v| {
if *v <= now {
*v = expires;
is_fresh = true;
}
})
.or_insert_with(|| {
is_fresh = true;
expires
});
is_fresh
}
pub fn sweep(&self) -> usize {
let now = now_ms();
let expired: Vec<(String, String)> = self
.inner
.iter()
.filter(|kv| *kv.value() <= now)
.map(|kv| (kv.key().0.clone(), kv.key().1.clone()))
.collect();
let mut removed = 0;
for k in expired {
if self.inner.remove(&k).is_some() {
removed += 1;
}
}
removed
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fresh_then_duplicate() {
let d = DedupeStore::new();
assert!(d.check_and_record("t", "k", Duration::from_secs(60)));
assert!(!d.check_and_record("t", "k", Duration::from_secs(60)));
}
#[test]
fn per_topic() {
let d = DedupeStore::new();
assert!(d.check_and_record("t1", "k", Duration::from_secs(60)));
assert!(d.check_and_record("t2", "k", Duration::from_secs(60)));
}
#[test]
fn sweep_drops_expired() {
let d = DedupeStore::new();
d.check_and_record("t", "k", Duration::from_millis(0));
d.inner.insert(("t".into(), "k".into()), 0);
assert_eq!(d.sweep(), 1);
}
#[test]
fn concurrent_dedup_returns_one_fresh() {
use std::sync::Arc;
use std::thread;
let store = Arc::new(DedupeStore::new());
let w = Duration::from_secs(60);
let handles: Vec<_> = (0..8)
.map(|_| {
let s = store.clone();
thread::spawn(move || s.check_and_record("t", "k", w))
})
.collect();
let results: Vec<bool> = handles.into_iter().map(|h| h.join().unwrap()).collect();
let fresh_count = results.iter().filter(|&&b| b).count();
assert_eq!(fresh_count, 1, "exactly one thread should see fresh");
}
}