phytrace_sdk/reliability/

batch.rs

//! Batch processing for efficient event transmission.

use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;

use crate::core::config::BufferConfig;
use crate::error::PhyTraceResult;
use crate::models::event::UdmEvent;
use crate::transport::traits::{BatchResult, Transport};

/// Batch processor for aggregating events before transmission.
#[derive(Debug)]
pub struct BatchProcessor {
    events: Arc<Mutex<Vec<UdmEvent>>>,
    max_size: usize,
    max_age: Duration,
    created_at: Arc<Mutex<Instant>>,
}

impl BatchProcessor {
    /// Create a new batch processor.
    pub fn new(max_size: usize, max_age: Duration) -> Self {
        Self {
            events: Arc::new(Mutex::new(Vec::new())),
            max_size,
            max_age,
            created_at: Arc::new(Mutex::new(Instant::now())),
        }
    }

    /// Create from buffer configuration.
    pub fn from_config(config: &BufferConfig) -> Self {
        Self::new(
            config.batch_size,
            Duration::from_secs(config.flush_interval_secs),
        )
    }

    /// Add an event to the batch.
    ///
    /// Returns `Some(events)` if the batch is ready to be sent.
    pub async fn add(&self, event: UdmEvent) -> Option<Vec<UdmEvent>> {
        let mut events = self.events.lock().await;
        events.push(event);

        if self.should_flush_inner(&events).await {
            let batch = std::mem::take(&mut *events);
            *self.created_at.lock().await = Instant::now();
            Some(batch)
        } else {
            None
        }
    }

    /// Add multiple events.
    ///
    /// Returns batches ready to be sent (may return multiple if events exceed batch size).
    pub async fn add_many(&self, new_events: Vec<UdmEvent>) -> Vec<Vec<UdmEvent>> {
        let mut events = self.events.lock().await;
        let mut batches = Vec::new();

        for event in new_events {
            events.push(event);

            if events.len() >= self.max_size {
                let batch = std::mem::take(&mut *events);
                *self.created_at.lock().await = Instant::now();
                batches.push(batch);
            }
        }

        batches
    }

    /// Check if the batch should be flushed.
    pub async fn should_flush(&self) -> bool {
        let events = self.events.lock().await;
        self.should_flush_inner(&events).await
    }

    /// Flush the current batch regardless of conditions.
    pub async fn flush(&self) -> Vec<UdmEvent> {
        let mut events = self.events.lock().await;
        let batch = std::mem::take(&mut *events);
        *self.created_at.lock().await = Instant::now();
        batch
    }

    /// Get current batch size.
    pub async fn len(&self) -> usize {
        self.events.lock().await.len()
    }

    /// Check if batch is empty.
    pub async fn is_empty(&self) -> bool {
        self.events.lock().await.is_empty()
    }

    /// Get age of current batch.
    pub async fn age(&self) -> Duration {
        self.created_at.lock().await.elapsed()
    }

    /// Send the batch using a transport.
    pub async fn send_batch<T: Transport>(
        &self,
        transport: &T,
    ) -> PhyTraceResult<Option<BatchResult>> {
        let batch = self.flush().await;
        if batch.is_empty() {
            return Ok(None);
        }

        let result = transport.send_batch(&batch).await?;
        Ok(Some(result))
    }

    // Private helper
    async fn should_flush_inner(&self, events: &[UdmEvent]) -> bool {
        // Check size
        if events.len() >= self.max_size {
            return true;
        }

        // Check age
        let age = self.created_at.lock().await.elapsed();
        if !events.is_empty() && age >= self.max_age {
            return true;
        }

        false
    }
}

/// Background batch flusher that periodically checks and sends batches.
pub struct BatchFlusher<T: Transport> {
    processor: Arc<BatchProcessor>,
    transport: Arc<T>,
    interval: Duration,
}

impl<T: Transport + 'static> BatchFlusher<T> {
    /// Create a new batch flusher.
    pub fn new(processor: Arc<BatchProcessor>, transport: Arc<T>, interval: Duration) -> Self {
        Self {
            processor,
            transport,
            interval,
        }
    }

    /// Start the background flush loop.
    ///
    /// Returns a handle that can be used to stop the flusher.
    pub fn start(self) -> FlushHandle {
        let (tx, mut rx) = tokio::sync::oneshot::channel::<()>();

        let processor = self.processor;
        let transport = self.transport;
        let interval = self.interval;

        tokio::spawn(async move {
            let mut interval_timer = tokio::time::interval(interval);

            loop {
                tokio::select! {
                    _ = interval_timer.tick() => {
                        if processor.should_flush().await {
                            let _ = processor.send_batch(transport.as_ref()).await;
                        }
                    }
                    _ = &mut rx => {
                        // Final flush before stopping
                        let _ = processor.send_batch(transport.as_ref()).await;
                        break;
                    }
                }
            }
        });

        FlushHandle { stop_tx: Some(tx) }
    }
}

/// Handle to stop a background flusher.
pub struct FlushHandle {
    stop_tx: Option<tokio::sync::oneshot::Sender<()>>,
}

impl FlushHandle {
    /// Stop the flusher.
    pub fn stop(mut self) {
        if let Some(tx) = self.stop_tx.take() {
            let _ = tx.send(());
        }
    }
}

impl Drop for FlushHandle {
    fn drop(&mut self) {
        if let Some(tx) = self.stop_tx.take() {
            let _ = tx.send(());
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::config::BufferConfig;
    use crate::models::domains::IdentityDomain;
    use crate::models::enums::SourceType;
    use crate::transport::mock::MockTransport;

    fn test_event(id: &str) -> UdmEvent {
        UdmEvent::new(SourceType::Amr).with_identity(IdentityDomain {
            source_id: Some(id.to_string()),
            ..Default::default()
        })
    }

    #[tokio::test]
    async fn test_batch_add_below_threshold() {
        let processor = BatchProcessor::new(5, Duration::from_secs(60));

        let result = processor.add(test_event("1")).await;
        assert!(result.is_none());
        assert_eq!(processor.len().await, 1);
    }

    #[tokio::test]
    async fn test_batch_add_reaches_threshold() {
        let processor = BatchProcessor::new(3, Duration::from_secs(60));

        processor.add(test_event("1")).await;
        processor.add(test_event("2")).await;
        let result = processor.add(test_event("3")).await;

        assert!(result.is_some());
        let batch = result.unwrap();
        assert_eq!(batch.len(), 3);
        assert!(processor.is_empty().await);
    }

    #[tokio::test]
    async fn test_batch_add_many() {
        let processor = BatchProcessor::new(3, Duration::from_secs(60));

        let events = vec![
            test_event("1"),
            test_event("2"),
            test_event("3"),
            test_event("4"),
            test_event("5"),
        ];

        let batches = processor.add_many(events).await;

        // Should produce one full batch of 3
        assert_eq!(batches.len(), 1);
        assert_eq!(batches[0].len(), 3);
        // Remaining 2 events in processor
        assert_eq!(processor.len().await, 2);
    }

    #[tokio::test]
    async fn test_batch_flush() {
        let processor = BatchProcessor::new(10, Duration::from_secs(60));

        processor.add(test_event("1")).await;
        processor.add(test_event("2")).await;

        let batch = processor.flush().await;

        assert_eq!(batch.len(), 2);
        assert!(processor.is_empty().await);
    }

    #[tokio::test]
    async fn test_batch_age_trigger() {
        let processor = BatchProcessor::new(100, Duration::from_millis(10));

        processor.add(test_event("1")).await;

        // Wait for batch to age
        tokio::time::sleep(Duration::from_millis(20)).await;

        assert!(processor.should_flush().await);
    }

    #[tokio::test]
    async fn test_empty_batch_no_flush() {
        let processor = BatchProcessor::new(5, Duration::from_millis(1));

        // Even with zero max age, empty batch shouldn't flush
        tokio::time::sleep(Duration::from_millis(5)).await;

        assert!(!processor.should_flush().await);
    }

    // Additional tests for improved coverage

    #[tokio::test]
    async fn test_from_config() {
        let config = BufferConfig {
            batch_size: 10,
            flush_interval_secs: 30,
            ..Default::default()
        };

        let processor = BatchProcessor::from_config(&config);

        // Add 10 events - should trigger batch
        for i in 0..9 {
            let result = processor.add(test_event(&i.to_string())).await;
            assert!(result.is_none());
        }
        let result = processor.add(test_event("10")).await;
        assert!(result.is_some());
        assert_eq!(result.unwrap().len(), 10);
    }

    #[tokio::test]
    async fn test_batch_age() {
        let processor = BatchProcessor::new(100, Duration::from_secs(60));

        // Age should be very small initially
        let age = processor.age().await;
        assert!(age < Duration::from_secs(1));

        // Wait a bit
        tokio::time::sleep(Duration::from_millis(50)).await;

        let age2 = processor.age().await;
        assert!(age2 >= Duration::from_millis(50));
    }

    #[tokio::test]
    async fn test_batch_is_empty() {
        let processor = BatchProcessor::new(5, Duration::from_secs(60));

        assert!(processor.is_empty().await);

        processor.add(test_event("1")).await;
        assert!(!processor.is_empty().await);

        processor.flush().await;
        assert!(processor.is_empty().await);
    }

    #[tokio::test]
    async fn test_batch_len() {
        let processor = BatchProcessor::new(10, Duration::from_secs(60));

        assert_eq!(processor.len().await, 0);

        processor.add(test_event("1")).await;
        assert_eq!(processor.len().await, 1);

        processor.add(test_event("2")).await;
        processor.add(test_event("3")).await;
        assert_eq!(processor.len().await, 3);
    }

    #[tokio::test]
    async fn test_add_many_multiple_batches() {
        let processor = BatchProcessor::new(3, Duration::from_secs(60));

        // Add 7 events - should produce 2 full batches
        let events: Vec<UdmEvent> = (0..7).map(|i| test_event(&i.to_string())).collect();

        let batches = processor.add_many(events).await;

        assert_eq!(batches.len(), 2);
        assert_eq!(batches[0].len(), 3);
        assert_eq!(batches[1].len(), 3);
        // 1 event remaining
        assert_eq!(processor.len().await, 1);
    }

    #[tokio::test]
    async fn test_add_many_exact_multiple() {
        let processor = BatchProcessor::new(3, Duration::from_secs(60));

        // Add exactly 6 events - should produce 2 full batches with nothing remaining
        let events: Vec<UdmEvent> = (0..6).map(|i| test_event(&i.to_string())).collect();

        let batches = processor.add_many(events).await;

        assert_eq!(batches.len(), 2);
        assert!(processor.is_empty().await);
    }

    #[tokio::test]
    async fn test_add_many_less_than_batch() {
        let processor = BatchProcessor::new(10, Duration::from_secs(60));

        let events: Vec<UdmEvent> = (0..5).map(|i| test_event(&i.to_string())).collect();

        let batches = processor.add_many(events).await;

        // No complete batches
        assert!(batches.is_empty());
        assert_eq!(processor.len().await, 5);
    }

    #[tokio::test]
    async fn test_send_batch_with_events() {
        let processor = BatchProcessor::new(10, Duration::from_secs(60));
        let transport = MockTransport::new();

        processor.add(test_event("1")).await;
        processor.add(test_event("2")).await;

        let result = processor.send_batch(&transport).await.unwrap();

        assert!(result.is_some());
        let batch_result = result.unwrap();
        assert_eq!(batch_result.total, 2);
        assert!(processor.is_empty().await);
    }

    #[tokio::test]
    async fn test_send_batch_empty() {
        let processor = BatchProcessor::new(10, Duration::from_secs(60));
        let transport = MockTransport::new();

        let result = processor.send_batch(&transport).await.unwrap();

        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_flush_resets_timer() {
        let processor = BatchProcessor::new(100, Duration::from_millis(50));

        processor.add(test_event("1")).await;

        // Wait for batch to age
        tokio::time::sleep(Duration::from_millis(60)).await;
        assert!(processor.should_flush().await);

        // Flush resets the timer
        processor.flush().await;

        // Add new event
        processor.add(test_event("2")).await;

        // Should not flush yet (timer was reset)
        assert!(!processor.should_flush().await);
    }

    #[tokio::test]
    async fn test_batch_flusher_start_and_stop() {
        let processor = Arc::new(BatchProcessor::new(100, Duration::from_millis(10)));
        let transport = Arc::new(MockTransport::new());

        processor.add(test_event("1")).await;

        let flusher = BatchFlusher::new(
            processor.clone(),
            transport.clone(),
            Duration::from_millis(20),
        );

        let handle = flusher.start();

        // Wait for at least one flush cycle
        tokio::time::sleep(Duration::from_millis(50)).await;

        handle.stop();

        // Transport should have received the event
        assert!(transport.event_count() >= 1);
    }

    #[tokio::test]
    async fn test_batch_flusher_final_flush_on_stop() {
        let processor = Arc::new(BatchProcessor::new(100, Duration::from_secs(60)));
        let transport = Arc::new(MockTransport::new());

        processor.add(test_event("1")).await;

        let flusher = BatchFlusher::new(
            processor.clone(),
            transport.clone(),
            Duration::from_secs(60), // Long interval - won't trigger naturally
        );

        let handle = flusher.start();

        // Immediately stop - should trigger final flush
        tokio::time::sleep(Duration::from_millis(10)).await;
        handle.stop();

        // Give time for final flush
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Event should have been sent during final flush
        assert_eq!(transport.event_count(), 1);
    }

    #[tokio::test]
    async fn test_flush_handle_drop() {
        let processor = Arc::new(BatchProcessor::new(100, Duration::from_secs(60)));
        let transport = Arc::new(MockTransport::new());

        processor.add(test_event("1")).await;

        let flusher = BatchFlusher::new(
            processor.clone(),
            transport.clone(),
            Duration::from_secs(60),
        );

        {
            let _handle = flusher.start();
            // Handle dropped here
        }

        // Give time for cleanup
        tokio::time::sleep(Duration::from_millis(50)).await;

        // Event should have been sent during drop
        assert_eq!(transport.event_count(), 1);
    }

    #[tokio::test]
    async fn test_should_flush_size_threshold() {
        let processor = BatchProcessor::new(3, Duration::from_secs(60));

        processor.add(test_event("1")).await;
        assert!(!processor.should_flush().await);

        processor.add(test_event("2")).await;
        assert!(!processor.should_flush().await);

        // Note: third add will auto-flush, so we check just before
        assert_eq!(processor.len().await, 2);
    }

    #[tokio::test]
    async fn test_batch_processor_concurrent_access() {
        let processor = Arc::new(BatchProcessor::new(100, Duration::from_secs(60)));

        let mut handles = vec![];

        for i in 0..10 {
            let proc = processor.clone();
            let handle = tokio::spawn(async move {
                proc.add(test_event(&i.to_string())).await;
            });
            handles.push(handle);
        }

        for handle in handles {
            handle.await.unwrap();
        }

        assert_eq!(processor.len().await, 10);
    }
}