worker pool

Author: wlawt

crates/audit/src/archiver.rs

@@ -1,19 +1,27 @@
 use crate::metrics::Metrics;
-use crate::reader::EventReader;
+use crate::reader::{Event, EventReader};
 use crate::storage::EventWriter;
 use anyhow::Result;
+use std::marker::PhantomData;
+use std::sync::Arc;
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
+use tokio::sync::{Mutex, mpsc};
 use tokio::time::sleep;
 use tracing::{error, info};
 
+// TODO: make this configurable
+const WORKER_POOL_SIZE: usize = 80;
+const CHANNEL_BUFFER_SIZE: usize = 500;
+
 pub struct KafkaAuditArchiver<R, W>
 where
     R: EventReader,
     W: EventWriter + Clone + Send + 'static,
 {
     reader: R,
-    writer: W,
+    event_tx: mpsc::Sender<Event>,
     metrics: Metrics,
+    _phantom: PhantomData<W>,
 }
 
 impl<R, W> KafkaAuditArchiver<R, W>
@@ -22,16 +30,58 @@ where
     W: EventWriter + Clone + Send + 'static,
 {
     pub fn new(reader: R, writer: W) -> Self {
+        let (event_tx, event_rx) = mpsc::channel(CHANNEL_BUFFER_SIZE);
+        let metrics = Metrics::default();
+
+        Self::spawn_workers(writer, event_rx, metrics.clone());
+
         Self {
             reader,
-            writer,
-            metrics: Metrics::default(),
+            event_tx,
+            metrics,
+            _phantom: PhantomData,
         }
     }
 
-    pub async fn run(&mut self) -> Result<()> {
-        info!("Starting Kafka bundle archiver");
+    fn spawn_workers(writer: W, event_rx: mpsc::Receiver<Event>, metrics: Metrics) {
+        let event_rx = Arc::new(Mutex::new(event_rx));
+
+        for worker_id in 0..WORKER_POOL_SIZE {
+            let writer = writer.clone();
+            let metrics = metrics.clone();
+            let event_rx = event_rx.clone();
+
+            tokio::spawn(async move {
+                loop {
+                    let event = {
+                        let mut rx = event_rx.lock().await;
+                        rx.recv().await
+                    };
+
+                    match event {
+                        Some(event) => {
+                            let archive_start = Instant::now();
+                            if let Err(e) = writer.archive_event(event).await {
+                                error!(worker_id, error = %e, "Failed to write event");
+                            } else {
+                                metrics
+                                    .archive_event_duration
+                                    .record(archive_start.elapsed().as_secs_f64());
+                                metrics.events_processed.increment(1);
+                            }
+                            metrics.in_flight_archive_tasks.decrement(1.0);
+                        }
+                        None => {
+                            info!(worker_id, "Worker stopped - channel closed");
+                            break;
+                        }
+                    }
+                }
+            });
+        }
+    }
 
+    pub async fn run(&mut self) -> Result<()> {
         loop {
             let read_start = Instant::now();
             match self.reader.read_event().await {
@@ -47,22 +97,11 @@ where
                     let event_age_ms = now_ms.saturating_sub(event.timestamp);
                     self.metrics.event_age.record(event_age_ms as f64);
 
-                    // TODO: the integration test breaks because Minio doesn't support etag
-                    let writer = self.writer.clone();
-                    let metrics = self.metrics.clone();
                     self.metrics.in_flight_archive_tasks.increment(1.0);
-                    tokio::spawn(async move {
-                        let archive_start = Instant::now();
-                        if let Err(e) = writer.archive_event(event).await {
-                            error!(error = %e, "Failed to write event");
-                        } else {
-                            metrics
-                                .archive_event_duration
-                                .record(archive_start.elapsed().as_secs_f64());
-                            metrics.events_processed.increment(1);
-                        }
-                        metrics.in_flight_archive_tasks.decrement(1.0);
-                    });
+                    if let Err(e) = self.event_tx.send(event).await {
+                        error!(error = %e, "Failed to send event to worker pool");
+                        self.metrics.in_flight_archive_tasks.decrement(1.0);
+                    }
 
                     let commit_start = Instant::now();
                     if let Err(e) = self.reader.commit().await {

crates/audit/src/storage.rs

@@ -15,7 +15,7 @@ use std::fmt;
 use std::fmt::Debug;
 use std::time::Instant;
 use tips_core::AcceptedBundle;
-use tracing::{error, info};
+use tracing::info;
 
 #[derive(Debug)]
 pub enum S3Key {
@@ -470,36 +470,35 @@ impl EventWriter for S3EventReaderWriter {
         let bundle_id = event.event.bundle_id();
         let transaction_ids = event.event.transaction_ids();
 
-        let bundle_writer = self.clone();
-        tokio::spawn(async move {
-            let start = Instant::now();
-            if let Err(e) = bundle_writer.update_bundle_history(event).await {
-                error!(error = %e, "Failed to update bundle history");
-            } else {
-                bundle_writer
-                    .metrics
-                    .update_bundle_history_duration
-                    .record(start.elapsed().as_secs_f64());
+        let bundle_start = Instant::now();
+        let bundle_future = self.update_bundle_history(event);
+
+        let tx_start = Instant::now();
+        let tx_futures: Vec<_> = transaction_ids
+            .into_iter()
+            .map(|tx_id| async move {
+                self.update_transaction_by_hash_index(&tx_id, bundle_id)
+                    .await
+            })
+            .collect();
+
+        let (bundle_result, tx_results) = tokio::join!(bundle_future, async {
+            let mut results = Vec::new();
+            for fut in tx_futures {
+                results.push(fut.await);
             }
+            results.into_iter().collect::<Result<Vec<_>>>()
         });
 
-        for tx_id in transaction_ids {
-            let start = Instant::now();
-            let tx_writer = self.clone();
-            tokio::spawn(async move {
-                if let Err(e) = tx_writer
-                    .update_transaction_by_hash_index(&tx_id, bundle_id)
-                    .await
-                {
-                    error!(error = %e, "Failed to update transaction by hash index");
-                } else {
-                    tx_writer
-                        .metrics
-                        .update_tx_indexes_duration
-                        .record(start.elapsed().as_secs_f64());
-                }
-            });
-        }
+        bundle_result?;
+        tx_results?;
+
+        self.metrics
+            .update_bundle_history_duration
+            .record(bundle_start.elapsed().as_secs_f64());
+        self.metrics
+            .update_tx_indexes_duration
+            .record(tx_start.elapsed().as_secs_f64());
 
         Ok(())
     }