Add helper methods to connect queues/publishers

Author: danyalprout

crates/audit/src/archiver.rs

@@ -5,7 +5,7 @@ use std::time::Duration;
 use tokio::time::sleep;
 use tracing::{error, info};
 
-pub struct KafkaMempoolArchiver<R, W>
+pub struct KafkaAuditArchiver<R, W>
 where
     R: EventReader,
     W: EventWriter,
@@ -14,7 +14,7 @@ where
     writer: W,
 }
 
-impl<R, W> KafkaMempoolArchiver<R, W>
+impl<R, W> KafkaAuditArchiver<R, W>
 where
     R: EventReader,
     W: EventWriter,

crates/audit/src/bin/main.rs

@@ -5,7 +5,7 @@ use aws_sdk_s3::{Client as S3Client, config::Builder as S3ConfigBuilder};
 use clap::{Parser, ValueEnum};
 use rdkafka::consumer::Consumer;
 use tips_audit::{
-    KafkaAuditLogReader, KafkaMempoolArchiver, S3EventReaderWriter, create_kafka_consumer,
+    KafkaAuditArchiver, KafkaAuditLogReader, S3EventReaderWriter, create_kafka_consumer,
 };
 use tips_core::logger::init_logger;
 use tracing::info;
@@ -71,7 +71,7 @@ async fn main() -> Result<()> {
     let s3_bucket = args.s3_bucket.clone();
     let writer = S3EventReaderWriter::new(s3_client, s3_bucket);
 
-    let mut archiver = KafkaMempoolArchiver::new(reader, writer);
+    let mut archiver = KafkaAuditArchiver::new(reader, writer);
 
     info!("Audit archiver initialized, starting main loop");
 

crates/audit/src/lib.rs

@@ -4,8 +4,25 @@ pub mod reader;
 pub mod storage;
 pub mod types;
 
+use tokio::sync::mpsc;
+use tracing::error;
+
 pub use archiver::*;
 pub use publisher::*;
 pub use reader::*;
 pub use storage::*;
 pub use types::*;
+
+pub fn connect_audit_to_publisher<P>(event_rx: mpsc::UnboundedReceiver<BundleEvent>, publisher: P)
+where
+    P: BundleEventPublisher + 'static,
+{
+    tokio::spawn(async move {
+        let mut event_rx = event_rx;
+        while let Some(event) = event_rx.recv().await {
+            if let Err(e) = publisher.publish(event).await {
+                error!(error = %e, "Failed to publish bundle event");
+            }
+        }
+    });
+}

crates/audit/tests/integration_tests.rs

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

crates/bundle-pool/src/lib.rs

@@ -1,6 +1,35 @@
 pub mod pool;
 pub mod source;
 
+use source::BundleSource;
+use std::sync::{Arc, Mutex};
+use tokio::sync::mpsc;
+use tracing::error;
+
 pub use pool::{BundleStore, InMemoryBundlePool};
 pub use source::KafkaBundleSource;
 pub use tips_core::{Bundle, BundleHash, BundleWithMetadata, CancelBundle};
+
+pub fn connect_sources_to_pool<S, P>(
+    sources: Vec<S>,
+    bundle_rx: mpsc::UnboundedReceiver<BundleWithMetadata>,
+    pool: Arc<Mutex<P>>,
+) where
+    S: BundleSource + Send + 'static,
+    P: BundleStore + Send + 'static,
+{
+    for source in sources {
+        tokio::spawn(async move {
+            if let Err(e) = source.run().await {
+                error!(error = %e, "Bundle source failed");
+            }
+        });
+    }
+
+    tokio::spawn(async move {
+        let mut bundle_rx = bundle_rx;
+        while let Some(bundle) = bundle_rx.recv().await {
+            pool.lock().unwrap().add_bundle(bundle);
+        }
+    });
+}

crates/bundle-pool/src/source.rs

@@ -1,10 +1,16 @@
 use anyhow::Result;
+use async_trait::async_trait;
 use rdkafka::consumer::{Consumer, StreamConsumer};
 use rdkafka::{ClientConfig, Message};
 use tips_core::{Bundle, BundleWithMetadata};
 use tokio::sync::mpsc;
 use tracing::{debug, error};
 
+#[async_trait]
+pub trait BundleSource {
+    async fn run(&self) -> Result<()>;
+}
+
 pub struct KafkaBundleSource {
     queue_consumer: StreamConsumer,
     publisher: mpsc::UnboundedSender<BundleWithMetadata>,
@@ -23,8 +29,11 @@ impl KafkaBundleSource {
             publisher,
         })
     }
+}
 
-    pub async fn run(&self) -> Result<()> {
+#[async_trait]
+impl BundleSource for KafkaBundleSource {
+    async fn run(&self) -> Result<()> {
         loop {
             match self.queue_consumer.recv().await {
                 Ok(message) => {

crates/bundle-pool/tests/integration_tests.rs

@@ -4,17 +4,20 @@ use rdkafka::producer::{FutureProducer, FutureRecord};
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 use testcontainers::runners::AsyncRunner;
+use testcontainers_modules::testcontainers::ContainerAsync;
 use testcontainers_modules::{kafka, kafka::Kafka};
 use tips_audit::BundleEvent;
-use tips_bundle_pool::{BundleStore, InMemoryBundlePool, KafkaBundleSource};
+use tips_bundle_pool::{
+    BundleStore, InMemoryBundlePool, KafkaBundleSource, connect_sources_to_pool,
+};
 use tips_core::{
     BundleWithMetadata,
     test_utils::{create_test_bundle, create_transaction},
 };
 use tokio::sync::mpsc;
 
-#[tokio::test]
-async fn test_kafka_bundle_source_to_pool_integration() -> Result<(), Box<dyn std::error::Error>> {
+async fn setup_kafka()
+-> Result<(ContainerAsync<Kafka>, FutureProducer, ClientConfig), Box<dyn std::error::Error>> {
     let kafka_container = Kafka::default().start().await?;
     let bootstrap_servers = format!(
         "127.0.0.1:{}",
@@ -23,15 +26,11 @@ async fn test_kafka_bundle_source_to_pool_integration() -> Result<(), Box<dyn st
             .await?
     );
 
-    let topic = "test-bundles";
-
     let kafka_producer = ClientConfig::new()
         .set("bootstrap.servers", &bootstrap_servers)
         .set("message.timeout.ms", "5000")
         .create::<FutureProducer>()?;
 
-
-    let (bundle_tx, mut bundle_rx) = mpsc::unbounded_channel::<BundleWithMetadata>();
     let mut kafka_consumer_config = ClientConfig::new();
     kafka_consumer_config
         .set("group.id", "bundle-pool-test-source")
@@ -40,23 +39,25 @@ async fn test_kafka_bundle_source_to_pool_integration() -> Result<(), Box<dyn st
         .set("enable.auto.commit", "false")
         .set("auto.offset.reset", "earliest");
 
+    Ok((kafka_container, kafka_producer, kafka_consumer_config))
+}
+
+#[tokio::test]
+async fn test_kafka_bundle_source_to_pool_integration() -> Result<(), Box<dyn std::error::Error>> {
+    let topic = "test-bundles";
+    let (_kafka_container, kafka_producer, kafka_consumer_config) = setup_kafka().await?;
+
+    let (bundle_tx, bundle_rx) = mpsc::unbounded_channel::<BundleWithMetadata>();
+
     let kafka_source = KafkaBundleSource::new(kafka_consumer_config, topic.to_string(), bundle_tx)?;
-    tokio::spawn(async move {
-        kafka_source.run().await.expect("Kafka source failed");
-    });
 
     let (audit_tx, _audit_rx) = mpsc::unbounded_channel::<BundleEvent>();
     let pool = Arc::new(Mutex::new(InMemoryBundlePool::new(
         audit_tx,
         "test-builder".to_string(),
     )));
 
-    let pool_clone = pool.clone();
-    tokio::spawn(async move {
-        while let Some(bundle) = bundle_rx.recv().await {
-            pool_clone.lock().unwrap().add_bundle(bundle);
-        }
-    });
+    connect_sources_to_pool(vec![kafka_source], bundle_rx, pool.clone());
 
     let alice = PrivateKeySigner::random();
     let bob = PrivateKeySigner::random();
@@ -79,11 +80,9 @@ async fn test_kafka_bundle_source_to_pool_integration() -> Result<(), Box<dyn st
     let mut counter = 0;
     loop {
         counter += 1;
-        if counter > 10 {
-            panic!("Bundle was not added to pool within timeout");
-        }
+        assert!(counter < 10);
 
-        tokio::time::sleep(Duration::from_secs(1)).await;
+        tokio::time::sleep(Duration::from_millis(500)).await;
 
         let bundles = pool.lock().unwrap().get_bundles();
         if bundles.is_empty() {