chore(kafka-dedup): refactor deduplication logic (#38924)

Author: jose-sequeira

rust/common/types/src/event.rs

@@ -7,6 +7,13 @@ use serde_json::Value;
 use time::{format_description::well_known::Rfc3339, OffsetDateTime};
 use uuid::Uuid;
 
+/// Information about the library/SDK that sent an event
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct LibraryInfo {
+    pub name: String,
+    pub version: Option<String>,
+}
+
 #[derive(Default, Debug, Deserialize, Serialize)]
 pub struct RawEvent {
     #[serde(
@@ -238,6 +245,24 @@ impl RawEvent {
             *value = f(value.take());
         }
     }
+
+    /// Extract library information from the event properties
+    /// Returns None if $lib property is not present
+    pub fn extract_library_info(&self) -> Option<LibraryInfo> {
+        let name = self
+            .properties
+            .get("$lib")
+            .and_then(|v| v.as_str())
+            .map(String::from)?;
+
+        let version = self
+            .properties
+            .get("$lib_version")
+            .and_then(|v| v.as_str())
+            .map(String::from);
+
+        Some(LibraryInfo { name, version })
+    }
 }
 
 impl CapturedEvent {

rust/kafka-deduplicator/src/checkpoint/worker.rs

@@ -335,7 +335,9 @@ mod tests {
 
     use super::*;
     use crate::checkpoint::CheckpointConfig;
-    use crate::store::{DeduplicationStore, DeduplicationStoreConfig};
+    use crate::store::{
+        DeduplicationStore, DeduplicationStoreConfig, TimestampKey, TimestampMetadata,
+    };
 
     use common_types::RawEvent;
     use tempfile::TempDir;
@@ -394,7 +396,9 @@ mod tests {
 
         // Add an event to the store
         let event = create_test_event();
-        store.handle_event_with_raw(&event).unwrap();
+        let key = TimestampKey::from(&event);
+        let metadata = TimestampMetadata::new(&event);
+        store.put_timestamp_record(&key, &metadata).unwrap();
 
         let tmp_checkpoint_dir = TempDir::new().unwrap();
         let config = CheckpointConfig {
@@ -452,7 +456,9 @@ mod tests {
 
         // Add an event to the store
         let event = create_test_event();
-        store.handle_event_with_raw(&event).unwrap();
+        let key = TimestampKey::from(&event);
+        let metadata = TimestampMetadata::new(&event);
+        store.put_timestamp_record(&key, &metadata).unwrap();
 
         let partition = Partition::new("some_test_topic".to_string(), 0);
 

rust/kafka-deduplicator/src/checkpoint_manager.rs

@@ -643,7 +643,9 @@ impl Drop for CheckpointManager {
 mod tests {
     use super::*;
     use crate::checkpoint::worker::CheckpointTarget;
-    use crate::store::{DeduplicationStore, DeduplicationStoreConfig};
+    use crate::store::{
+        DeduplicationStore, DeduplicationStoreConfig, TimestampKey, TimestampMetadata,
+    };
     use common_types::RawEvent;
     use std::{collections::HashMap, path::PathBuf, time::Duration};
     use tempfile::TempDir;
@@ -775,8 +777,11 @@ mod tests {
 
         // Add events to the stores
         let event = create_test_event();
-        store1.handle_event_with_raw(&event).unwrap();
-        store2.handle_event_with_raw(&event).unwrap();
+        // Add test data directly to stores
+        let key = TimestampKey::from(&event);
+        let metadata = TimestampMetadata::new(&event);
+        store1.put_timestamp_record(&key, &metadata).unwrap();
+        store2.put_timestamp_record(&key, &metadata).unwrap();
 
         // add dedup stores to manager
         let stores = store_manager.stores();
@@ -840,9 +845,13 @@ mod tests {
 
         // Add an event
         let event1 = create_test_event();
-        store.handle_event_with_raw(&event1).unwrap();
+        let key1 = TimestampKey::from(&event1);
+        let metadata1 = TimestampMetadata::new(&event1);
+        store.put_timestamp_record(&key1, &metadata1).unwrap();
         let event2 = create_test_event();
-        store.handle_event_with_raw(&event2).unwrap();
+        let key2 = TimestampKey::from(&event2);
+        let metadata2 = TimestampMetadata::new(&event2);
+        store.put_timestamp_record(&key2, &metadata2).unwrap();
 
         // Create manager with short interval for testing
         let tmp_checkpoint_dir = TempDir::new().unwrap();
@@ -971,8 +980,11 @@ mod tests {
 
         // Add events to the stores
         let event = create_test_event();
-        store1.handle_event_with_raw(&event).unwrap();
-        store2.handle_event_with_raw(&event).unwrap();
+        // Add test data directly to stores
+        let key = TimestampKey::from(&event);
+        let metadata = TimestampMetadata::new(&event);
+        store1.put_timestamp_record(&key, &metadata).unwrap();
+        store2.put_timestamp_record(&key, &metadata).unwrap();
 
         // add dedup stores to manager
         let stores = store_manager.stores();
@@ -1036,8 +1048,11 @@ mod tests {
 
         // Add events to the stores
         let event = create_test_event();
-        store1.handle_event_with_raw(&event).unwrap();
-        store2.handle_event_with_raw(&event).unwrap();
+        // Add test data directly to stores
+        let key = TimestampKey::from(&event);
+        let metadata = TimestampMetadata::new(&event);
+        store1.put_timestamp_record(&key, &metadata).unwrap();
+        store2.put_timestamp_record(&key, &metadata).unwrap();
 
         // add dedup stores to manager
         let stores = store_manager.stores();
@@ -1102,7 +1117,9 @@ mod tests {
             let event = create_test_event();
             let part = Partition::new("max_inflight_checkpoints".to_string(), i);
             let store = create_test_store(part.topic(), part.partition_number());
-            store.handle_event_with_raw(&event).unwrap();
+            let key = TimestampKey::from(&event);
+            let metadata = TimestampMetadata::new(&event);
+            store.put_timestamp_record(&key, &metadata).unwrap();
             stores.insert(part, store);
         }
 
@@ -1172,9 +1189,11 @@ mod tests {
 
         // Add events to the stores
         let event = create_test_event();
-        store1.handle_event_with_raw(&event).unwrap();
-        store2.handle_event_with_raw(&event).unwrap();
-        store3.handle_event_with_raw(&event).unwrap();
+        let key = crate::store::keys::TimestampKey::from(&event);
+        let metadata = crate::store::metadata::TimestampMetadata::new(&event);
+        store1.put_timestamp_record(&key, &metadata).unwrap();
+        store2.put_timestamp_record(&key, &metadata).unwrap();
+        store3.put_timestamp_record(&key, &metadata).unwrap();
 
         // add dedup stores to manager
         let stores = store_manager.stores();