use anyhow::Result;
use async_trait::async_trait;
use crate::bloom::BloomDeduplicator;
use crate::set_trait::SeenSet;
#[cfg(feature = "redis")]
use redis::AsyncCommands;
#[cfg(not(feature = "redis"))]
use std::sync::Arc;
pub struct HybridSeenSet {
bloom: BloomDeduplicator,
#[cfg(feature = "redis")]
redis: redis::aio::ConnectionManager,
#[cfg(feature = "redis")]
key_prefix: String,
#[cfg(not(feature = "redis"))]
fallback: Arc<tokio::sync::RwLock<std::collections::HashSet<String>>>,
}
impl HybridSeenSet {
#[cfg(feature = "redis")]
pub async fn new(
redis_url: &str,
key_prefix: Option<String>,
expected_urls: usize,
false_positive_rate: f64,
) -> Result<Self> {
let client = redis::Client::open(redis_url)?;
let redis = redis::aio::ConnectionManager::new(client).await?;
Ok(Self {
bloom: BloomDeduplicator::new(expected_urls, false_positive_rate),
redis,
key_prefix: key_prefix.unwrap_or_else(|| "argus:seen:".to_string()),
})
}
#[cfg(not(feature = "redis"))]
pub async fn new(
_redis_url: &str,
_key_prefix: Option<String>,
expected_urls: usize,
false_positive_rate: f64,
) -> Result<Self> {
Ok(Self {
bloom: BloomDeduplicator::new(expected_urls, false_positive_rate),
fallback: Arc::new(tokio::sync::RwLock::new(std::collections::HashSet::new())),
})
}
pub fn stats(&self) -> HybridStats {
HybridStats {
bloom_memory_bytes: self.bloom.memory_usage(),
bloom_bit_count: self.bloom.bit_count(),
bloom_hash_count: self.bloom.hash_count(),
}
}
}
#[derive(Debug, Clone)]
pub struct HybridStats {
pub bloom_memory_bytes: usize,
pub bloom_bit_count: u64,
pub bloom_hash_count: u32,
}
#[async_trait]
impl SeenSet for HybridSeenSet {
#[cfg(feature = "redis")]
async fn insert_if_new(&self, url: String) -> bool {
if !self.bloom.might_contain(&url) {
self.bloom.insert(&url);
let key = format!("{}{}", self.key_prefix, url);
let mut conn = self.redis.clone();
match conn.set_nx::<_, _, bool>(&key, 1).await {
Ok(was_new) => was_new,
Err(e) => {
eprintln!("Redis error, assuming new: {}", e);
true
}
}
} else {
let key = format!("{}{}", self.key_prefix, url);
let mut conn = self.redis.clone();
match conn.set_nx::<_, _, bool>(&key, 1).await {
Ok(was_new) => {
if was_new {
self.bloom.insert(&url);
}
was_new
}
Err(e) => {
eprintln!("Redis error, assuming duplicate: {}", e);
false
}
}
}
}
#[cfg(not(feature = "redis"))]
async fn insert_if_new(&self, url: String) -> bool {
if !self.bloom.might_contain(&url) {
self.bloom.insert(&url);
let mut set = self.fallback.write().await;
set.insert(url);
true
} else {
let mut set = self.fallback.write().await;
if set.insert(url.clone()) {
self.bloom.insert(&url);
true
} else {
false
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
#[ignore] async fn hybrid_deduplication() {
let seen = HybridSeenSet::new("redis://localhost:6379", None, 1000, 0.01)
.await
.unwrap();
assert!(
seen.insert_if_new("https://example.com/1".to_string())
.await
);
assert!(
!seen
.insert_if_new("https://example.com/1".to_string())
.await
);
assert!(
seen.insert_if_new("https://example.com/2".to_string())
.await
);
}
#[tokio::test]
#[ignore] async fn hybrid_stats() {
let seen = HybridSeenSet::new("redis://localhost:6379", None, 1_000_000, 0.01)
.await
.unwrap();
let stats = seen.stats();
assert!(stats.bloom_memory_bytes > 0);
assert!(stats.bloom_memory_bytes < 2_000_000); }
#[tokio::test]
#[ignore] async fn hybrid_large_scale() {
let seen = HybridSeenSet::new("redis://localhost:6379", None, 100_000_000, 0.01)
.await
.unwrap();
let stats = seen.stats();
let memory_mb = stats.bloom_memory_bytes as f64 / 1_048_576.0;
assert!(memory_mb < 200.0, "Memory usage: {:.2} MB", memory_mb);
}
}