feat: kafka partition key and headers (#1821)

Author: dinhani-cw

src/infra/kafka/kafka.rs

@@ -1,12 +1,13 @@
 use anyhow::Result;
 use clap::ValueEnum;
-use ethereum_types::H256;
+use rdkafka::message::Header;
+use rdkafka::message::OwnedHeaders;
 use rdkafka::producer::FutureProducer;
 use rdkafka::producer::FutureRecord;
 use rdkafka::ClientConfig;
 
-use crate::eth::primitives::Hash;
 use crate::ledger::events::Event;
+use crate::log_and_err;
 
 #[derive(Clone, serde::Serialize, serde::Deserialize)]
 pub struct KafkaConfig {
@@ -104,22 +105,24 @@ impl KafkaConnector {
     }
 
     pub async fn send_event<T: Event>(&self, event: T) -> Result<()> {
-        let payload = serde_json::to_string(&event)?;
+        // prepare base payload
+        let headers = event.event_headers()?;
+        let key = event.event_key()?;
+        let payload = event.event_payload()?;
 
-        match self
-            .producer
-            .send(
-                FutureRecord::to(&self.topic).payload(&payload).key(&Hash(H256::random()).to_string()),
-                std::time::Duration::from_secs(0),
-            )
-            .await
-        {
-            Ok(_) => {
-                tracing::info!(payload = payload, "event sent to kafka");
-                println!("event sent to kafka {:?}", payload);
-                Ok(())
-            }
-            Err(e) => Err(anyhow::anyhow!("failed to send event to kafka: {:?}", e)),
+        // prepare kafka payload
+        let mut kafka_headers = OwnedHeaders::new_with_capacity(headers.len());
+        for (key, value) in headers.iter() {
+            let header = Header { key, value: Some(value) };
+            kafka_headers = kafka_headers.insert(header);
         }
+        let kafka_record = FutureRecord::to(&self.topic).payload(&payload).key(&key).headers(kafka_headers);
+
+        // publis and handle response
+        tracing::info!(%key, %payload, ?headers, "publishing kafka event");
+        if let Err((e, _)) = self.producer.send(kafka_record, std::time::Duration::from_secs(0)).await {
+            return log_and_err!(reason = e, "failed to publish kafka event");
+        }
+        Ok(())
     }
 }

src/ledger/events.rs

@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::collections::HashMap;
 use std::collections::HashSet;
 
 use chrono::DateTime;
@@ -200,10 +201,31 @@ impl Serialize for AccountTransferDirection {
 // Marker Trait
 // -----------------------------------------------------------------------------
 
-/// Marker trait for events
-pub trait Event: Serialize {}
+/// Struct is an event that can be published to external systems.
+pub trait Event: Serialize + Sized {
+    /// Returns the partition key component of the event.
+    fn event_key(&self) -> anyhow::Result<String>;
 
-impl Event for AccountTransfers {}
+    /// Returns the headers component of the event.
+    ///
+    /// By default, it returns empty headers.
+    fn event_headers(&self) -> anyhow::Result<HashMap<String, String>> {
+        Ok(HashMap::default())
+    }
+
+    /// Returns the payload component of the event.
+    ///
+    /// By default, it serializes the implementing struct as JSON.
+    fn event_payload(&self) -> anyhow::Result<String> {
+        Ok(serde_json::to_string(self)?)
+    }
+}
+
+impl Event for AccountTransfers {
+    fn event_key(&self) -> anyhow::Result<String> {
+        Ok(self.account_address.to_string())
+    }
+}
 
 // -----------------------------------------------------------------------------
 // Conversions