chore: kafka bundle consumer / clean up

Author: danyalprout

Cargo.lock

@@ -5,10 +5,10 @@ use aws_sdk_s3::{Client as S3Client, config::Builder as S3ConfigBuilder};
 use clap::{Parser, ValueEnum};
 use rdkafka::consumer::Consumer;
 use tips_audit::{
-    KafkaMempoolArchiver, KafkaMempoolReader, S3EventReaderWriter, create_kafka_consumer,
+    KafkaAuditLogReader, KafkaMempoolArchiver, S3EventReaderWriter, create_kafka_consumer,
 };
-use tracing::{info, warn};
-use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+use tips_core::logger::init_logger;
+use tracing::info;
 
 #[derive(Debug, Clone, ValueEnum)]
 enum S3ConfigType {
@@ -53,28 +53,7 @@ async fn main() -> Result<()> {
 
     let args = Args::parse();
 
-    let log_level = match args.log_level.to_lowercase().as_str() {
-        "trace" => tracing::Level::TRACE,
-        "debug" => tracing::Level::DEBUG,
-        "info" => tracing::Level::INFO,
-        "warn" => tracing::Level::WARN,
-        "error" => tracing::Level::ERROR,
-        _ => {
-            warn!(
-                "Invalid log level '{}', defaulting to 'info'",
-                args.log_level
-            );
-            tracing::Level::INFO
-        }
-    };
-
-    tracing_subscriber::registry()
-        .with(
-            tracing_subscriber::EnvFilter::try_from_default_env()
-                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(log_level.to_string())),
-        )
-        .with(tracing_subscriber::fmt::layer())
-        .init();
+    init_logger(&args.log_level);
 
     info!(
         kafka_properties_file = %args.kafka_properties_file,
@@ -86,7 +65,7 @@ async fn main() -> Result<()> {
     let consumer = create_kafka_consumer(&args.kafka_properties_file)?;
     consumer.subscribe(&[&args.kafka_topic])?;
 
-    let reader = KafkaMempoolReader::new(consumer, args.kafka_topic.clone())?;
+    let reader = KafkaAuditLogReader::new(consumer, args.kafka_topic.clone())?;
 
     let s3_client = create_s3_client(&args).await?;
     let s3_bucket = args.s3_bucket.clone();

crates/audit/src/bin/main.rs

@@ -7,36 +7,18 @@ use rdkafka::{
     consumer::{Consumer, StreamConsumer},
     message::Message,
 };
-use std::fs;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use tips_core::kafka::load_kafka_config_from_file;
 use tokio::time::sleep;
-use tracing::{debug, error, info};
+use tracing::{debug, error};
 
 pub fn create_kafka_consumer(kafka_properties_file: &str) -> Result<StreamConsumer> {
-    let client_config = load_kafka_config_from_file(kafka_properties_file)?;
+    let client_config =
+        ClientConfig::from_iter(load_kafka_config_from_file(kafka_properties_file)?);
     let consumer: StreamConsumer = client_config.create()?;
     Ok(consumer)
 }
 
-fn load_kafka_config_from_file(properties_file_path: &str) -> Result<ClientConfig> {
-    let kafka_properties = fs::read_to_string(properties_file_path)?;
-    info!("Kafka properties:\n{}", kafka_properties);
-
-    let mut client_config = ClientConfig::new();
-
-    for line in kafka_properties.lines() {
-        let line = line.trim();
-        if line.is_empty() || line.starts_with('#') {
-            continue;
-        }
-        if let Some((key, value)) = line.split_once('=') {
-            client_config.set(key.trim(), value.trim());
-        }
-    }
-
-    Ok(client_config)
-}
-
 pub fn assign_topic_partition(consumer: &StreamConsumer, topic: &str) -> Result<()> {
     let mut tpl = TopicPartitionList::new();
     tpl.add_partition(topic, 0);
@@ -57,14 +39,14 @@ pub trait EventReader {
     async fn commit(&mut self) -> Result<()>;
 }
 
-pub struct KafkaMempoolReader {
+pub struct KafkaAuditLogReader {
     consumer: StreamConsumer,
     topic: String,
     last_message_offset: Option<i64>,
     last_message_partition: Option<i32>,
 }
 
-impl KafkaMempoolReader {
+impl KafkaAuditLogReader {
     pub fn new(consumer: StreamConsumer, topic: String) -> Result<Self> {
         consumer.subscribe(&[&topic])?;
         Ok(Self {
@@ -77,7 +59,7 @@ impl KafkaMempoolReader {
 }
 
 #[async_trait]
-impl EventReader for KafkaMempoolReader {
+impl EventReader for KafkaAuditLogReader {
     async fn read_event(&mut self) -> Result<Event> {
         match self.consumer.recv().await {
             Ok(message) => {
@@ -143,7 +125,7 @@ impl EventReader for KafkaMempoolReader {
     }
 }
 
-impl KafkaMempoolReader {
+impl KafkaAuditLogReader {
     pub fn topic(&self) -> &str {
         &self.topic
     }

crates/audit/src/reader.rs

@@ -1,6 +1,6 @@
 use std::time::Duration;
 use tips_audit::{
-    KafkaMempoolArchiver, KafkaMempoolReader,
+    KafkaAuditLogReader, KafkaMempoolArchiver,
     publisher::{BundleEventPublisher, KafkaBundleEventPublisher},
     storage::{BundleEventS3Reader, S3EventReaderWriter},
     types::{BundleEvent, DropReason},
@@ -38,7 +38,7 @@ async fn test_kafka_publisher_s3_archiver_integration()
     }
 
     let mut consumer = KafkaMempoolArchiver::new(
-        KafkaMempoolReader::new(harness.kafka_consumer, topic.to_string())?,
+        KafkaAuditLogReader::new(harness.kafka_consumer, topic.to_string())?,
         s3_writer.clone(),
     );
 

crates/audit/tests/integration_tests.rs

@@ -17,6 +17,8 @@ tracing.workspace = true
 tokio.workspace = true
 anyhow.workspace = true
 async-trait.workspace = true
+rdkafka.workspace = true
+serde_json.workspace = true
 
 [dev-dependencies]
 tips-core = { workspace = true, features = ["test-utils"] }

crates/bundle-pool/Cargo.toml

@@ -1,4 +1,6 @@
 pub mod pool;
+pub mod source;
 
 pub use pool::{BundleStore, InMemoryBundlePool};
+pub use source::KafkaBundleSource;
 pub use tips_core::{Bundle, BundleHash, BundleWithMetadata, CancelBundle};

crates/bundle-pool/src/lib.rs

@@ -0,0 +1,72 @@
+use anyhow::Result;
+use rdkafka::consumer::{Consumer, StreamConsumer};
+use rdkafka::{ClientConfig, Message};
+use tips_core::{Bundle, BundleWithMetadata};
+use tokio::sync::mpsc;
+use tracing::{debug, error};
+
+pub struct KafkaBundleSource {
+    queue_consumer: StreamConsumer,
+    publisher: mpsc::UnboundedSender<BundleWithMetadata>,
+}
+
+impl KafkaBundleSource {
+    pub fn new(
+        client_config: ClientConfig,
+        topic: String,
+        publisher: mpsc::UnboundedSender<BundleWithMetadata>,
+    ) -> Result<Self> {
+        let queue_consumer: StreamConsumer = client_config.create()?;
+        queue_consumer.subscribe(&[topic.as_str()])?;
+        Ok(Self {
+            queue_consumer,
+            publisher,
+        })
+    }
+
+    pub async fn run(&self) -> Result<()> {
+        loop {
+            match self.queue_consumer.recv().await {
+                Ok(message) => {
+                    let payload = match message.payload() {
+                        Some(p) => p,
+                        None => {
+                            error!("Message has no payload");
+                            continue;
+                        }
+                    };
+
+                    let bundle: Bundle = match serde_json::from_slice(payload) {
+                        Ok(b) => b,
+                        Err(e) => {
+                            error!(error = %e, "Failed to deserialize bundle");
+                            continue;
+                        }
+                    };
+
+                    debug!(
+                        bundle = ?bundle,
+                        offset = message.offset(),
+                        partition = message.partition(),
+                        "Received bundle from Kafka"
+                    );
+
+                    let bundle_with_metadata = match BundleWithMetadata::load(bundle) {
+                        Ok(b) => b,
+                        Err(e) => {
+                            error!(error = %e, "Failed to load bundle");
+                            continue;
+                        }
+                    };
+
+                    if let Err(e) = self.publisher.send(bundle_with_metadata) {
+                        error!(error = ?e, "Failed to publish bundle to queue");
+                    }
+                }
+                Err(e) => {
+                    error!(error = %e, "Error receiving message from Kafka");
+                }
+            }
+        }
+    }
+}

crates/bundle-pool/src/source.rs

@@ -21,6 +21,7 @@ alloy-serde = { version = "1.0.41", default-features = false }
 alloy-signer-local = { workspace = true, optional = true }
 op-alloy-rpc-types = { workspace = true, optional = true }
 tracing.workspace = true
+tracing-subscriber.workspace = true
 
 [dev-dependencies]
 alloy-signer-local.workspace = true

crates/core/Cargo.toml

@@ -0,0 +1,22 @@
+use std::collections::HashMap;
+use std::fs;
+
+pub fn load_kafka_config_from_file(
+    properties_file_path: &str,
+) -> Result<HashMap<String, String>, std::io::Error> {
+    let kafka_properties = fs::read_to_string(properties_file_path)?;
+
+    let mut config = HashMap::new();
+
+    for line in kafka_properties.lines() {
+        let line = line.trim();
+        if line.is_empty() || line.starts_with('#') {
+            continue;
+        }
+        if let Some((key, value)) = line.split_once('=') {
+            config.insert(key.trim().to_string(), value.trim().to_string());
+        }
+    }
+
+    Ok(config)
+}

crates/core/src/kafka.rs

@@ -1,3 +1,5 @@
+pub mod kafka;
+pub mod logger;
 pub mod types;
 
 #[cfg(any(test, feature = "test-utils"))]