use async_trait::async_trait;
use redis::{aio::ConnectionManager, AsyncCommands, Client, RedisResult};
use serde::{Deserialize, Serialize};
use std::fmt::Debug;
use std::hash::{Hash, Hasher};
use std::marker::PhantomData;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, error, trace, warn};
use uuid::Uuid;
use crate::error::{StateError, StateResult};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DeduplicationConfig {
pub redis_url: String,
pub key_prefix: String,
pub default_ttl: Option<Duration>,
pub connect_timeout: Duration,
pub read_timeout: Duration,
pub write_timeout: Duration,
pub max_retries: u32,
pub retry_base_delay: Duration,
pub retry_max_delay: Duration,
pub pipeline_batch_size: usize,
pub compress_event_ids: bool,
}
impl Default for DeduplicationConfig {
fn default() -> Self {
Self {
redis_url: "redis://localhost:6379".to_string(),
key_prefix: "dedup:".to_string(),
default_ttl: Some(Duration::from_secs(3600)), connect_timeout: Duration::from_secs(5),
read_timeout: Duration::from_secs(3),
write_timeout: Duration::from_secs(3),
max_retries: 3,
retry_base_delay: Duration::from_millis(100),
retry_max_delay: Duration::from_secs(5),
pipeline_batch_size: 100,
compress_event_ids: false,
}
}
}
impl DeduplicationConfig {
pub fn builder() -> DeduplicationConfigBuilder {
DeduplicationConfigBuilder::default()
}
}
#[derive(Debug, Default)]
pub struct DeduplicationConfigBuilder {
config: DeduplicationConfig,
}
impl DeduplicationConfigBuilder {
pub fn redis_url(mut self, url: impl Into<String>) -> Self {
self.config.redis_url = url.into();
self
}
pub fn key_prefix(mut self, prefix: impl Into<String>) -> Self {
self.config.key_prefix = prefix.into();
self
}
pub fn default_ttl(mut self, ttl: Duration) -> Self {
self.config.default_ttl = Some(ttl);
self
}
pub fn connect_timeout(mut self, timeout: Duration) -> Self {
self.config.connect_timeout = timeout;
self
}
pub fn retries(mut self, max: u32, base_delay: Duration, max_delay: Duration) -> Self {
self.config.max_retries = max;
self.config.retry_base_delay = base_delay;
self.config.retry_max_delay = max_delay;
self
}
pub fn pipeline_batch_size(mut self, size: usize) -> Self {
self.config.pipeline_batch_size = size;
self
}
pub fn compress_event_ids(mut self, enabled: bool) -> Self {
self.config.compress_event_ids = enabled;
self
}
pub fn build(self) -> StateResult<DeduplicationConfig> {
if self.config.redis_url.is_empty() {
return Err(StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: "Redis URL cannot be empty".to_string(),
});
}
if self.config.pipeline_batch_size == 0 {
return Err(StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: "Pipeline batch size must be greater than 0".to_string(),
});
}
Ok(self.config)
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct DeduplicationStats {
pub events_checked: u64,
pub duplicates_found: u64,
pub new_events: u64,
pub cache_hits: u64,
pub cache_misses: u64,
pub batch_operations: u64,
pub retry_count: u64,
pub error_count: u64,
pub avg_latency_us: u64,
pub bytes_written: u64,
pub bytes_read: u64,
}
impl DeduplicationStats {
pub fn duplicate_rate(&self) -> f64 {
if self.events_checked == 0 {
0.0
} else {
self.duplicates_found as f64 / self.events_checked as f64
}
}
pub fn cache_hit_rate(&self) -> f64 {
let total = self.cache_hits + self.cache_misses;
if total == 0 {
0.0
} else {
self.cache_hits as f64 / total as f64
}
}
pub fn error_rate(&self) -> f64 {
if self.events_checked == 0 {
0.0
} else {
self.error_count as f64 / self.events_checked as f64
}
}
pub fn total_operations(&self) -> u64 {
self.events_checked + self.new_events
}
}
pub enum EventIdExtractor<T> {
Uuid(Box<dyn Fn(&T) -> Uuid + Send + Sync>),
String(Box<dyn Fn(&T) -> String + Send + Sync>),
Hash(Box<dyn Fn(&T) -> u64 + Send + Sync>),
Custom(Box<dyn Fn(&T) -> String + Send + Sync>),
}
impl<T> EventIdExtractor<T> {
pub fn uuid<F>(f: F) -> Self
where
F: Fn(&T) -> Uuid + Send + Sync + 'static,
{
Self::Uuid(Box::new(f))
}
pub fn string<F>(f: F) -> Self
where
F: Fn(&T) -> String + Send + Sync + 'static,
{
Self::String(Box::new(f))
}
pub fn hash<F>(f: F) -> Self
where
F: Fn(&T) -> u64 + Send + Sync + 'static,
{
Self::Hash(Box::new(f))
}
pub fn custom<F>(f: F) -> Self
where
F: Fn(&T) -> String + Send + Sync + 'static,
{
Self::Custom(Box::new(f))
}
pub fn extract_id(&self, event: &T) -> String {
match self {
Self::Uuid(f) => f(event).to_string(),
Self::String(f) => f(event),
Self::Hash(f) => format!("{:016x}", f(event)),
Self::Custom(f) => f(event),
}
}
}
impl<T> Debug for EventIdExtractor<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Uuid(_) => write!(f, "EventIdExtractor::Uuid"),
Self::String(_) => write!(f, "EventIdExtractor::String"),
Self::Hash(_) => write!(f, "EventIdExtractor::Hash"),
Self::Custom(_) => write!(f, "EventIdExtractor::Custom"),
}
}
}
struct RedisDeduplicationBackend {
connection: Arc<ConnectionManager>,
config: DeduplicationConfig,
stats: Arc<RwLock<DeduplicationStats>>,
}
impl RedisDeduplicationBackend {
async fn new(config: DeduplicationConfig) -> StateResult<Self> {
let client = Self::create_client(&config)?;
let connection = Self::create_connection_manager(client, &config).await?;
Ok(Self {
connection: Arc::new(connection),
config,
stats: Arc::new(RwLock::new(DeduplicationStats::default())),
})
}
fn create_client(config: &DeduplicationConfig) -> StateResult<Client> {
Client::open(config.redis_url.as_str()).map_err(|e| StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: format!("Failed to create Redis client: {}", e),
})
}
async fn create_connection_manager(
client: Client,
config: &DeduplicationConfig,
) -> StateResult<ConnectionManager> {
let mut retries = 0;
let mut delay = config.retry_base_delay;
loop {
match ConnectionManager::new(client.clone()).await {
Ok(conn) => {
debug!("Redis deduplication backend connected");
return Ok(conn);
}
Err(e) => {
retries += 1;
if retries >= config.max_retries {
error!("Failed to connect to Redis after {} retries: {}", retries, e);
return Err(StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: format!("Connection failed after {} retries: {}", retries, e),
});
}
warn!(
"Redis connection attempt {} failed: {}, retrying in {:?}",
retries, e, delay
);
tokio::time::sleep(delay).await;
delay = std::cmp::min(delay * 2, config.retry_max_delay);
}
}
}
}
fn build_key(&self, event_id: &str) -> String {
format!("{}{}", self.config.key_prefix, event_id)
}
async fn with_retry<F, Fut, T>(&self, mut operation: F) -> StateResult<T>
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = RedisResult<T>>,
{
let mut retries = 0;
let mut delay = self.config.retry_base_delay;
loop {
match operation().await {
Ok(result) => return Ok(result),
Err(e) => {
retries += 1;
let mut stats = self.stats.write().await;
stats.error_count += 1;
stats.retry_count += 1;
drop(stats);
if retries >= self.config.max_retries {
error!("Deduplication operation failed after {} retries: {}", retries, e);
return Err(StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: format!("Operation failed after {} retries: {}", retries, e),
});
}
warn!(
"Deduplication operation attempt {} failed: {}, retrying in {:?}",
retries, e, delay
);
tokio::time::sleep(delay).await;
delay = std::cmp::min(delay * 2, self.config.retry_max_delay);
}
}
}
}
async fn exists(&self, event_id: &str) -> StateResult<bool> {
let start = std::time::Instant::now();
let key = self.build_key(event_id);
trace!("Checking existence of key: {}", key);
let conn = self.connection.as_ref().clone();
let key_clone = key.clone();
let exists: bool = self
.with_retry(move || {
let mut conn = conn.clone();
let key = key_clone.clone();
async move { conn.exists(&key).await }
})
.await?;
let mut stats = self.stats.write().await;
stats.events_checked += 1;
if exists {
stats.duplicates_found += 1;
stats.cache_hits += 1;
} else {
stats.cache_misses += 1;
}
let latency_us = start.elapsed().as_micros() as u64;
stats.avg_latency_us = if stats.events_checked == 1 {
latency_us
} else {
(stats.avg_latency_us * 9 + latency_us) / 10
};
Ok(exists)
}
async fn mark_seen(&self, event_id: &str) -> StateResult<()> {
let key = self.build_key(event_id);
trace!("Marking key as seen: {}", key);
let conn = self.connection.as_ref().clone();
let key_clone = key.clone();
let ttl = self.config.default_ttl;
self.with_retry(move || {
let mut conn = conn.clone();
let key = key_clone.clone();
async move {
if let Some(ttl_duration) = ttl {
let ttl_secs = ttl_duration.as_secs();
conn.set_ex(&key, "1", ttl_secs).await
} else {
conn.set(&key, "1").await
}
}
})
.await?;
let mut stats = self.stats.write().await;
stats.new_events += 1;
stats.bytes_written += event_id.len() as u64;
Ok(())
}
async fn batch_exists(&self, event_ids: &[String]) -> StateResult<Vec<bool>> {
if event_ids.is_empty() {
return Ok(Vec::new());
}
let keys: Vec<String> = event_ids.iter().map(|id| self.build_key(id)).collect();
let conn = self.connection.as_ref().clone();
let keys_clone = keys.clone();
let results: Vec<bool> = self
.with_retry(move || {
let mut conn = conn.clone();
let keys = keys_clone.clone();
async move {
let mut pipe = redis::pipe();
for key in &keys {
pipe.exists(key);
}
pipe.query_async(&mut conn).await
}
})
.await?;
let mut stats = self.stats.write().await;
stats.events_checked += event_ids.len() as u64;
stats.batch_operations += 1;
for &exists in &results {
if exists {
stats.duplicates_found += 1;
stats.cache_hits += 1;
} else {
stats.cache_misses += 1;
}
}
Ok(results)
}
async fn batch_mark_seen(&self, event_ids: &[String]) -> StateResult<()> {
if event_ids.is_empty() {
return Ok(());
}
for chunk in event_ids.chunks(self.config.pipeline_batch_size) {
let commands: Vec<(String, Option<u64>)> = chunk
.iter()
.map(|id| {
let key = self.build_key(id);
let ttl = self.config.default_ttl.map(|d| d.as_secs());
(key, ttl)
})
.collect();
let conn = self.connection.as_ref().clone();
self.with_retry(move || {
let mut conn = conn.clone();
let commands = commands.clone();
async move {
let mut pipe = redis::pipe();
for (key, ttl) in &commands {
if let Some(ttl_secs) = ttl {
pipe.set_ex(key, "1", *ttl_secs);
} else {
pipe.set(key, "1");
}
}
pipe.query_async(&mut conn).await
}
})
.await?;
let mut stats = self.stats.write().await;
stats.new_events += chunk.len() as u64;
stats.batch_operations += 1;
stats.bytes_written += chunk.iter().map(|id| id.len() as u64).sum::<u64>();
}
Ok(())
}
async fn stats(&self) -> DeduplicationStats {
self.stats.read().await.clone()
}
async fn health_check(&self) -> StateResult<bool> {
let mut conn = self.connection.as_ref().clone();
let result: RedisResult<String> = redis::cmd("PING").query_async(&mut conn).await;
match result {
Ok(response) => Ok(response == "PONG"),
Err(e) => {
error!("Deduplication backend health check failed: {}", e);
Err(StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: format!("Health check failed: {}", e),
})
}
}
}
async fn clear(&self) -> StateResult<usize> {
let pattern = format!("{}*", self.config.key_prefix);
let mut conn = self.connection.as_ref().clone();
let mut cursor = 0u64;
let mut total_deleted = 0;
loop {
let (new_cursor, keys): (u64, Vec<String>) = redis::cmd("SCAN")
.arg(cursor)
.arg("MATCH")
.arg(&pattern)
.arg("COUNT")
.arg(1000)
.query_async(&mut conn)
.await
.map_err(|e| StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: format!("SCAN failed: {}", e),
})?;
if !keys.is_empty() {
let deleted: usize = redis::cmd("DEL")
.arg(&keys)
.query_async(&mut conn)
.await
.map_err(|e| StateError::StorageError {
backend_type: "redis-dedup".to_string(),
details: format!("DEL failed: {}", e),
})?;
total_deleted += deleted;
}
cursor = new_cursor;
if cursor == 0 {
break;
}
}
debug!("Cleared {} deduplication entries", total_deleted);
Ok(total_deleted)
}
}
pub struct EventDeduplicator<T> {
backend: Arc<RedisDeduplicationBackend>,
extractor: EventIdExtractor<T>,
_phantom: PhantomData<T>,
}
impl<T> EventDeduplicator<T> {
pub async fn new(
config: DeduplicationConfig,
extractor: EventIdExtractor<T>,
) -> StateResult<Self> {
let backend = RedisDeduplicationBackend::new(config).await?;
Ok(Self {
backend: Arc::new(backend),
extractor,
_phantom: PhantomData,
})
}
pub async fn is_duplicate(&self, event: &T) -> StateResult<bool> {
let event_id = self.extractor.extract_id(event);
self.backend.exists(&event_id).await
}
pub async fn mark_seen(&self, event: &T) -> StateResult<()> {
let event_id = self.extractor.extract_id(event);
self.backend.mark_seen(&event_id).await
}
pub async fn check_and_mark(&self, event: &T) -> StateResult<bool> {
let event_id = self.extractor.extract_id(event);
let was_duplicate = self.backend.exists(&event_id).await?;
if !was_duplicate {
self.backend.mark_seen(&event_id).await?;
}
Ok(was_duplicate)
}
pub async fn batch_check_duplicates(&self, events: &[T]) -> StateResult<Vec<bool>> {
let event_ids: Vec<String> = events.iter().map(|e| self.extractor.extract_id(e)).collect();
self.backend.batch_exists(&event_ids).await
}
pub async fn batch_mark_seen(&self, events: &[&T]) -> StateResult<()> {
let event_ids: Vec<String> = events.iter().map(|e| self.extractor.extract_id(e)).collect();
self.backend.batch_mark_seen(&event_ids).await
}
pub async fn stats(&self) -> DeduplicationStats {
self.backend.stats().await
}
pub async fn health_check(&self) -> StateResult<bool> {
self.backend.health_check().await
}
pub async fn clear(&self) -> StateResult<usize> {
self.backend.clear().await
}
}
impl<T> Clone for EventDeduplicator<T> {
fn clone(&self) -> Self {
Self {
backend: Arc::clone(&self.backend),
extractor: EventIdExtractor::Custom(Box::new(|_| {
panic!("Cannot clone EventDeduplicator with extractor. Create a new instance instead.")
})),
_phantom: PhantomData,
}
}
}
impl<T> Debug for EventDeduplicator<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EventDeduplicator")
.field("extractor", &self.extractor)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
struct TestEvent {
id: String,
data: String,
}
async fn create_test_deduplicator() -> Option<EventDeduplicator<TestEvent>> {
let config = DeduplicationConfig::builder()
.redis_url("redis://localhost:6379")
.key_prefix("test:dedup:")
.default_ttl(Duration::from_secs(60))
.build()
.ok()?;
let extractor = EventIdExtractor::string(|event: &TestEvent| event.id.clone());
match EventDeduplicator::new(config, extractor).await {
Ok(dedup) => {
let _ = dedup.clear().await;
Some(dedup)
}
Err(_) => None,
}
}
#[tokio::test]
async fn test_basic_deduplication() {
if let Some(dedup) = create_test_deduplicator().await {
let event1 = TestEvent {
id: "event-1".to_string(),
data: "data1".to_string(),
};
assert!(!dedup.is_duplicate(&event1).await.unwrap());
dedup.mark_seen(&event1).await.unwrap();
assert!(dedup.is_duplicate(&event1).await.unwrap());
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[tokio::test]
async fn test_check_and_mark() {
if let Some(dedup) = create_test_deduplicator().await {
let event = TestEvent {
id: "event-check-mark".to_string(),
data: "data".to_string(),
};
let was_dup1 = dedup.check_and_mark(&event).await.unwrap();
assert!(!was_dup1);
let was_dup2 = dedup.check_and_mark(&event).await.unwrap();
assert!(was_dup2);
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[tokio::test]
async fn test_batch_operations() {
if let Some(dedup) = create_test_deduplicator().await {
let events = vec![
TestEvent {
id: "batch-1".to_string(),
data: "data1".to_string(),
},
TestEvent {
id: "batch-2".to_string(),
data: "data2".to_string(),
},
TestEvent {
id: "batch-3".to_string(),
data: "data3".to_string(),
},
];
let results = dedup.batch_check_duplicates(&events).await.unwrap();
assert_eq!(results, vec![false, false, false]);
let refs: Vec<&TestEvent> = events.iter().collect();
dedup.batch_mark_seen(&refs).await.unwrap();
let results2 = dedup.batch_check_duplicates(&events).await.unwrap();
assert_eq!(results2, vec![true, true, true]);
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[tokio::test]
async fn test_stats() {
if let Some(dedup) = create_test_deduplicator().await {
let event1 = TestEvent {
id: "stats-1".to_string(),
data: "data1".to_string(),
};
let event2 = TestEvent {
id: "stats-2".to_string(),
data: "data2".to_string(),
};
dedup.is_duplicate(&event1).await.unwrap();
dedup.mark_seen(&event1).await.unwrap();
dedup.is_duplicate(&event1).await.unwrap();
dedup.is_duplicate(&event2).await.unwrap();
let stats = dedup.stats().await;
assert_eq!(stats.events_checked, 3);
assert_eq!(stats.duplicates_found, 1);
assert_eq!(stats.new_events, 1);
assert_eq!(stats.cache_hits, 1);
assert_eq!(stats.cache_misses, 2);
assert!(stats.duplicate_rate() > 0.0);
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[tokio::test]
async fn test_uuid_extractor() {
if let Some(_) = create_test_deduplicator().await {
#[derive(Debug, Clone)]
struct UuidEvent {
uuid: Uuid,
}
let config = DeduplicationConfig::builder()
.redis_url("redis://localhost:6379")
.key_prefix("test:uuid:")
.build()
.unwrap();
let extractor = EventIdExtractor::uuid(|event: &UuidEvent| event.uuid);
if let Ok(dedup) = EventDeduplicator::new(config, extractor).await {
let uuid = Uuid::new_v4();
let event = UuidEvent { uuid };
assert!(!dedup.is_duplicate(&event).await.unwrap());
dedup.mark_seen(&event).await.unwrap();
assert!(dedup.is_duplicate(&event).await.unwrap());
dedup.clear().await.unwrap();
}
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[tokio::test]
async fn test_hash_extractor() {
if let Some(_) = create_test_deduplicator().await {
use std::collections::hash_map::DefaultHasher;
#[derive(Debug, Clone)]
struct HashEvent {
data: String,
}
let config = DeduplicationConfig::builder()
.redis_url("redis://localhost:6379")
.key_prefix("test:hash:")
.build()
.unwrap();
let extractor = EventIdExtractor::hash(|event: &HashEvent| {
let mut hasher = DefaultHasher::new();
event.data.hash(&mut hasher);
hasher.finish()
});
if let Ok(dedup) = EventDeduplicator::new(config, extractor).await {
let event1 = HashEvent {
data: "same-data".to_string(),
};
let event2 = HashEvent {
data: "same-data".to_string(),
};
assert!(!dedup.is_duplicate(&event1).await.unwrap());
dedup.mark_seen(&event1).await.unwrap();
assert!(dedup.is_duplicate(&event2).await.unwrap());
dedup.clear().await.unwrap();
}
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[tokio::test]
async fn test_health_check() {
if let Some(dedup) = create_test_deduplicator().await {
let healthy = dedup.health_check().await.unwrap();
assert!(healthy);
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[tokio::test]
async fn test_clear() {
if let Some(dedup) = create_test_deduplicator().await {
let events: Vec<TestEvent> = (0..5)
.map(|i| TestEvent {
id: format!("clear-{}", i),
data: "data".to_string(),
})
.collect();
let refs: Vec<&TestEvent> = events.iter().collect();
dedup.batch_mark_seen(&refs).await.unwrap();
let deleted = dedup.clear().await.unwrap();
assert_eq!(deleted, 5);
let results = dedup.batch_check_duplicates(&events).await.unwrap();
assert_eq!(results, vec![false, false, false, false, false]);
} else {
eprintln!("Skipping test: Redis not available");
}
}
#[test]
fn test_config_builder() {
let config = DeduplicationConfig::builder()
.redis_url("redis://localhost:6379")
.key_prefix("myapp:dedup:")
.default_ttl(Duration::from_secs(7200))
.pipeline_batch_size(200)
.compress_event_ids(true)
.build()
.unwrap();
assert_eq!(config.redis_url, "redis://localhost:6379");
assert_eq!(config.key_prefix, "myapp:dedup:");
assert_eq!(config.default_ttl, Some(Duration::from_secs(7200)));
assert_eq!(config.pipeline_batch_size, 200);
assert!(config.compress_event_ids);
}
#[test]
fn test_stats_calculations() {
let mut stats = DeduplicationStats::default();
assert_eq!(stats.duplicate_rate(), 0.0);
assert_eq!(stats.cache_hit_rate(), 0.0);
stats.events_checked = 100;
stats.duplicates_found = 25;
stats.cache_hits = 25;
stats.cache_misses = 75;
assert_eq!(stats.duplicate_rate(), 0.25);
assert_eq!(stats.cache_hit_rate(), 0.25);
assert_eq!(stats.total_operations(), 100);
}
}