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());
if let Some(mut entry) = self.inner.get_mut(&k) {
if *entry.value() <= now {
*entry.value_mut() = expires;
true
} else {
false
}
} else {
self.inner.insert(k, expires);
true
}
}
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);
}
}