chore(event_ingestion): Try to adapt to new design

eppo_core/src/event_ingestion/auto_flusher.rs

@@ -0,0 +1,142 @@
+use std::time::Duration;
+
+use tokio::{sync::mpsc, time::Instant};
+
+use super::BatchedMessage;
+
+/// Auto-flusher forwards all messages from `uplink` to `downlink` unchanged and inserts extra flush
+/// requests if it hasn't seen one within the given `period`. In other words, it makes sure that the
+/// channel is flushed at least every `period`.
+pub(super) async fn auto_flusher<T>(
+    mut uplink: mpsc::Receiver<BatchedMessage<T>>,
+    downlink: mpsc::Sender<BatchedMessage<T>>,
+    period: Duration,
+) -> Option<()> {
+    'flushed: loop {
+        // Process first message.
+        let msg = uplink.recv().await?;
+        let flushed = msg.flush.is_some();
+        downlink.send(msg).await.ok()?;
+
+        // No need to time if we just flushed.
+        if flushed {
+            continue;
+        }
+
+        let flush_at = Instant::now() + period;
+        // loop till we reach flush_at or see a flushed message.
+        loop {
+            tokio::select! {
+                _ =  tokio::time::sleep_until(flush_at) =>  {
+                    downlink.send(BatchedMessage { batch: Vec::new(), flush: Some(()) }).await.ok()?;
+                    continue 'flushed;
+                },
+                msg = uplink.recv() => {
+                    let msg = msg?;
+                    let flushed = msg.flush.is_some();
+                    downlink.send(msg).await.ok()?;
+                    if flushed {
+                        continue 'flushed;
+                    }
+                }
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::event_ingestion::auto_flusher;
+    use crate::event_ingestion::batched_message::BatchedMessage;
+    use tokio::sync::mpsc;
+    use tokio::time::{self, Duration};
+
+    #[tokio::test]
+    async fn test_auto_flusher() {
+        let (uplink_tx, uplink_rx) = mpsc::channel(10);
+        let (downlink_tx, mut downlink_rx) = mpsc::channel(10);
+        let flush_period = Duration::from_millis(100);
+        tokio::spawn(auto_flusher::auto_flusher(
+            uplink_rx,
+            downlink_tx,
+            flush_period,
+        ));
+
+        uplink_tx
+            .send(BatchedMessage {
+                batch: vec![1, 2, 3],
+                flush: None,
+            })
+            .await
+            .unwrap();
+        uplink_tx
+            .send(BatchedMessage {
+                batch: vec![4, 5, 6],
+                flush: None,
+            })
+            .await
+            .unwrap();
+
+        // Verify that the messages are forwarded to downlink
+        assert_eq!(
+            downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: vec![1, 2, 3],
+                flush: None
+            })
+        );
+        assert_eq!(
+            downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: vec![4, 5, 6],
+                flush: None
+            })
+        );
+
+        // Wait for the flush period to trigger an auto-flush
+        time::sleep(flush_period * 2).await;
+
+        // Verify the auto-flush behavior
+        assert_eq!(
+            downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: Vec::new(),
+                flush: Some(())
+            })
+        );
+
+        // Send a flushed message explicitly
+        uplink_tx
+            .send(BatchedMessage {
+                batch: vec![],
+                flush: Some(()),
+            })
+            .await
+            .unwrap();
+
+        // Verify that the flushed message is forwarded immediately
+        assert_eq!(
+            downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: vec![],
+                flush: Some(())
+            })
+        );
+
+        // Ensure the loop continues and processes further messages
+        uplink_tx
+            .send(BatchedMessage {
+                batch: vec![7, 8, 9],
+                flush: None,
+            })
+            .await
+            .unwrap();
+        assert_eq!(
+            downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: vec![7, 8, 9],
+                flush: None
+            })
+        );
+    }
+}

eppo_core/src/event_ingestion/batched_message.rs

@@ -0,0 +1,31 @@
+/// Batched message contain a batch of data and may optionally require processors to flush any processing.
+#[derive(Debug, PartialEq)]
+pub(super) struct BatchedMessage<T> {
+    pub batch: Vec<T>,
+    /// `None` means the message does not require a flush.
+    /// `Some` contains a list of watchers.
+    pub flush: Option<()>,
+}
+
+impl<T> BatchedMessage<T> {
+    /// Create a new empty message.
+    pub fn empty() -> BatchedMessage<T> {
+        BatchedMessage {
+            batch: Vec::new(),
+            flush: None,
+        }
+    }
+
+    pub fn requires_flush(&self) -> bool {
+        self.flush.is_some()
+    }
+
+    // Mark the message as successfully flushed, consuming it and notifying any interested parties.
+    // pub fn flushed(self) {
+    //     if let Some(flush) = self.flush {
+    //         for f in flush {
+    //             f.send(());
+    //         }
+    //     }
+    // }
+}

eppo_core/src/event_ingestion/batcher.rs

@@ -0,0 +1,145 @@
+use tokio::sync::mpsc;
+
+use super::BatchedMessage;
+
+/// Batch messages, so they are at least `min_batch_size` size. Push incomplete batch down if flush
+/// is received.
+///
+/// If uplink is closed, send all buffered data downstream and exit.
+///
+/// If downlink is closed, just exit.
+pub(super) async fn batcher<T>(
+    mut uplink: mpsc::Receiver<BatchedMessage<T>>,
+    downlink: mpsc::Sender<BatchedMessage<T>>,
+    min_batch_size: usize,
+) -> Option<()> {
+    let mut uplink_alive = true;
+    while uplink_alive {
+        let mut batch = BatchedMessage::empty();
+
+        while uplink_alive && batch.batch.len() < min_batch_size && batch.flush.is_none() {
+            match uplink.recv().await {
+                None => {
+                    uplink_alive = false;
+                }
+                Some(BatchedMessage {
+                    batch: events,
+                    flush,
+                }) => {
+                    batch.batch.extend(events);
+                    batch.flush = flush;
+                }
+            }
+        }
+
+        downlink.send(batch).await.ok()?;
+    }
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::event_ingestion::batched_message::BatchedMessage;
+    use crate::event_ingestion::{auto_flusher, batcher};
+
+    #[tokio::test]
+    async fn test_auto_flusher_and_batcher_pipeline() {
+        use tokio::sync::mpsc;
+        use tokio::time::{self, Duration};
+
+        // Define test parameters
+        let flush_period = Duration::from_millis(100);
+        let min_batch_size = 5;
+
+        // Channels for the pipeline
+        let (flusher_uplink_tx, flusher_uplink_rx) = mpsc::channel(10);
+        let (flusher_downlink_tx, flusher_downlink_rx) = mpsc::channel(10);
+        let (batcher_downlink_tx, mut batcher_downlink_rx) = mpsc::channel(10);
+
+        // Spawn the auto_flusher and batcher
+        tokio::spawn(auto_flusher::auto_flusher(
+            flusher_uplink_rx,
+            flusher_downlink_tx,
+            flush_period,
+        ));
+        tokio::spawn(batcher::batcher(
+            flusher_downlink_rx,
+            batcher_downlink_tx,
+            min_batch_size,
+        ));
+
+        // Send some messages to the flusher uplink
+        flusher_uplink_tx
+            .send(BatchedMessage {
+                batch: vec![1, 2, 3],
+                flush: None,
+            })
+            .await
+            .unwrap();
+        flusher_uplink_tx
+            .send(BatchedMessage {
+                batch: vec![4],
+                flush: None,
+            })
+            .await
+            .unwrap();
+
+        // Verify that the batcher does not output anything until the minimum batch size is met
+        assert_eq!(
+            batcher_downlink_rx.try_recv(),
+            Err(mpsc::error::TryRecvError::Empty)
+        );
+
+        // Wait for the auto_flusher to send a flush message
+        time::sleep(flush_period * 2).await;
+
+        // Verify that the batcher outputs a batch due to the flush
+        assert_eq!(
+            batcher_downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: vec![1, 2, 3, 4],
+                flush: Some(()),
+            })
+        );
+
+        // Send additional messages to the flusher
+        flusher_uplink_tx
+            .send(BatchedMessage {
+                batch: vec![5, 6, 7, 8, 9],
+                flush: None,
+            })
+            .await
+            .unwrap();
+
+        // Verify that the batcher outputs a batch once the minimum batch size is met
+        assert_eq!(
+            batcher_downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: vec![5, 6, 7, 8, 9],
+                flush: None,
+            })
+        );
+
+        // Simulate uplink closure
+        flusher_uplink_tx
+            .send(BatchedMessage {
+                batch: vec![10],
+                flush: None,
+            })
+            .await
+            .unwrap();
+        drop(flusher_uplink_tx);
+
+        // Verify that the batcher flushes the remaining data on uplink closure
+        assert_eq!(
+            batcher_downlink_rx.recv().await,
+            Some(BatchedMessage {
+                batch: vec![10],
+                flush: None,
+            })
+        );
+
+        // Verify that the batcher exits cleanly
+        assert_eq!(batcher_downlink_rx.recv().await, None);
+    }
+}

eppo_core/src/event_ingestion/delivery.rs

@@ -0,0 +1,70 @@
+use super::{BatchedMessage, Event};
+use crate::event_ingestion::event_delivery::{EventDelivery, EventDeliveryError};
+use log::warn;
+use tokio::sync::mpsc;
+
+pub(super) struct DeliveryStatus {
+    success: Vec<Event>,
+    failure: Vec<Event>,
+}
+
+pub(super) async fn delivery(
+    mut uplink: mpsc::Receiver<BatchedMessage<Event>>,
+    delivery_status: mpsc::Sender<DeliveryStatus>,
+    event_delivery: EventDelivery,
+) -> Option<()> {
+    loop {
+        let event_delivery = event_delivery.clone();
+        match uplink.recv().await {
+            None => {
+                return None;
+            }
+            Some(BatchedMessage {
+                batch,
+                flush: _flush,
+            }) => {
+                let result = event_delivery.deliver(batch.clone()).await;
+                match result {
+                    Ok(response) => {
+                        let failed_event_uuids = response.failed_events;
+                        if !failed_event_uuids.is_empty() {
+                            warn!("Failed to deliver {} events", failed_event_uuids.len());
+                            let mut success = Vec::new();
+                            let mut failure = Vec::new();
+                            batch.into_iter().for_each(|queued_event| {
+                                if failed_event_uuids.contains(&queued_event.uuid) {
+                                    failure.push(queued_event);
+                                } else {
+                                    success.push(queued_event);
+                                }
+                            });
+                            delivery_status
+                                .send(DeliveryStatus { success, failure })
+                                .await
+                                .ok()?;
+                        }
+                    }
+                    Err(err) => {
+                        match err {
+                            EventDeliveryError::RetriableError(_) => {
+                                // Retry later
+                                delivery_status
+                                    .send(DeliveryStatus {
+                                        failure: batch,
+                                        success: Vec::new(),
+                                    })
+                                    .await
+                                    .ok()?;
+                            }
+                            EventDeliveryError::NonRetriableError(_) => {
+                                warn!("Failed to deliver events: {}", err);
+                                // In this case there is no point in retrying delivery since the error is
+                                // non-retriable.
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+}