foxhunt/trading_engine/src/events/ring_buffer.rs

//! Lock-Free Ring Buffer Implementation for Event Storage
//!
//! This module provides specialized ring buffers optimized for high-frequency trading events
//! with sub-microsecond insertion performance and guaranteed ordering.

use super::event_types::{EventSequence, TradingEvent};
use super::EventProcessingError;
use crate::lockfree::ring_buffer::LockFreeRingBuffer;
use crate::timing::HardwareTimestamp;
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use tokio::sync::RwLock;

/// Specialized ring buffer for trading events with sequence tracking
#[derive(Debug)]
pub struct EventRingBuffer {
    /// Underlying lock-free ring buffer
    buffer: LockFreeRingBuffer<TradingEvent>,
    /// Buffer identifier for load balancing
    buffer_id: usize,
    /// Statistics tracking
    stats: Arc<RwLock<BufferStats>>,
    /// Last sequence number processed
    last_sequence: AtomicU64,
    /// Event counter for this buffer
    event_counter: AtomicU64,
}

impl EventRingBuffer {
    /// Create a new event ring buffer
    pub fn new(buffer_id: usize, capacity: usize) -> Result<Self> {
        let buffer = LockFreeRingBuffer::new(capacity)
            .map_err(|e| anyhow!("Failed to create ring buffer: {}", e))?;

        Ok(Self {
            buffer,
            buffer_id,
            stats: Arc::new(RwLock::new(BufferStats::new(buffer_id, capacity))),
            last_sequence: AtomicU64::new(0),
            event_counter: AtomicU64::new(0),
        })
    }

    /// Try to push an event into the buffer (lock-free)
    #[inline(always)]
    pub async fn try_push(
        &self,
        event: TradingEvent,
    ) -> Result<EventSequence, EventProcessingError> {
        let start_time = HardwareTimestamp::now();

        // Attempt lock-free insertion
        if let Ok(()) = self.buffer.try_push(event.clone()) {
            // Update sequence tracking
            let sequence = event.sequence_number().unwrap_or(0);
            self.last_sequence.store(sequence, Ordering::Relaxed);
            self.event_counter.fetch_add(1, Ordering::Relaxed);

            // Update statistics
            let latency_ns = HardwareTimestamp::now().latency_ns(&start_time);
            self.update_stats_push_success(latency_ns).await;

            Ok(EventSequence::new(sequence))
        } else {
            // Buffer is full
            self.update_stats_push_failure().await;
            Err(EventProcessingError::BufferFull(format!(
                "Buffer {} is full",
                self.buffer_id
            )))
        }
    }

    /// Try to pop events from the buffer (lock-free)
    #[inline(always)]
    pub async fn try_pop_batch(&self, max_events: usize) -> Option<Vec<TradingEvent>> {
        let mut events = Vec::with_capacity(max_events);
        let mut popped = 0;

        while popped < max_events {
            match self.buffer.try_pop() {
                Some(event) => {
                    events.push(event);
                    popped += 1;
                }
                None => break,
            }
        }

        if !events.is_empty() {
            self.update_stats_pop_success(events.len()).await;
    // Some variant
            Some(events)
        } else {
    // None variant
            None
        }
    }

    /// Get current buffer utilization
    pub fn utilization(&self) -> f64 {
        self.buffer.utilization()
    }

    /// Check if buffer is full
    pub fn is_full(&self) -> bool {
        self.buffer.is_full()
    }

    /// Check if buffer is empty
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }

    /// Get buffer capacity
    pub const fn capacity(&self) -> usize {
        self.buffer.capacity()
    }

    /// Get current length
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    /// Get buffer statistics
    pub async fn get_stats(&self) -> BufferStats {
        self.stats.read().await.clone()
    }

    /// Get buffer ID
    pub const fn buffer_id(&self) -> usize {
        self.buffer_id
    }

    /// Get last processed sequence number
    pub fn last_sequence(&self) -> u64 {
        self.last_sequence.load(Ordering::Relaxed)
    }

    /// Update statistics after successful push
    async fn update_stats_push_success(&self, latency_ns: u64) {
        let mut stats = self.stats.write().await;
        stats.push_success_count += 1;
        stats.total_push_latency_ns += latency_ns;
        stats.current_utilization = self.utilization();
        stats.last_updated = std::time::Instant::now();
    }

    /// Update statistics after failed push
    async fn update_stats_push_failure(&self) {
        let mut stats = self.stats.write().await;
        stats.push_failure_count += 1;
        stats.current_utilization = self.utilization();
        stats.last_updated = std::time::Instant::now();
    }

    /// Update statistics after successful pop
    async fn update_stats_pop_success(&self, count: usize) {
        let mut stats = self.stats.write().await;
        stats.pop_success_count += 1;
        stats.total_events_popped += count as u64;
        stats.current_utilization = self.utilization();
        stats.last_updated = std::time::Instant::now();
    }
}

/// Statistics for monitoring buffer performance
#[derive(Debug, Clone)]
/// BufferStats
/// 
/// TODO: Add detailed documentation for this struct
pub struct BufferStats {
    /// Buffer Id
    pub buffer_id: usize,
    /// Capacity
    pub capacity: usize,
    /// Current Utilization
    pub current_utilization: f64,
    /// Push Success Count
    pub push_success_count: u64,
    /// Push Failure Count
    pub push_failure_count: u64,
    /// Pop Success Count
    pub pop_success_count: u64,
    /// Total Events Popped
    pub total_events_popped: u64,
    /// Total Push Latency Ns
    pub total_push_latency_ns: u64,
    /// Avg Push Latency Ns
    pub avg_push_latency_ns: f64,
    /// Last Updated
    pub last_updated: std::time::Instant,
}

impl BufferStats {
    pub fn new(buffer_id: usize, capacity: usize) -> Self {
        Self {
            buffer_id,
            capacity,
            current_utilization: 0.0,
            push_success_count: 0,
            push_failure_count: 0,
            pop_success_count: 0,
            total_events_popped: 0,
            total_push_latency_ns: 0,
            avg_push_latency_ns: 0.0,
            last_updated: std::time::Instant::now(),
        }
    }

    /// Calculate average push latency
    pub fn calculate_avg_latency(&mut self) {
        if self.push_success_count > 0 {
            self.avg_push_latency_ns =
                self.total_push_latency_ns as f64 / self.push_success_count as f64;
        }
    }
}

/// Manager for multiple ring buffers with load balancing
#[derive(Debug)]
pub struct BufferManager {
    /// Array of event ring buffers
    buffers: Vec<Arc<EventRingBuffer>>,
    /// Current buffer index for round-robin selection
    current_buffer: AtomicUsize,
    /// Load balancing strategy
    strategy: LoadBalancingStrategy,
}

impl BufferManager {
    /// Create a new buffer manager
    pub fn new(buffer_count: usize, buffer_size: usize) -> Result<Self> {
        let mut buffers = Vec::with_capacity(buffer_count);

        for i in 0..buffer_count {
            let buffer = Arc::new(EventRingBuffer::new(i, buffer_size)?);
            buffers.push(buffer);
        }

        Ok(Self {
            buffers,
            current_buffer: AtomicUsize::new(0),
            strategy: LoadBalancingStrategy::RoundRobin,
        })
    }

    /// Select optimal buffer for event insertion
    pub fn select_buffer(&self) -> usize {
        match self.strategy {
            LoadBalancingStrategy::RoundRobin => {
                self.current_buffer.fetch_add(1, Ordering::Relaxed) % self.buffers.len()
            }
            LoadBalancingStrategy::LeastUtilized => self.select_least_utilized_buffer(),
            LoadBalancingStrategy::Hash => {
                // Use thread ID for hashing to improve CPU cache locality
                let thread_id = std::thread::current().id();
                let hash = self.hash_thread_id(thread_id);
                hash % self.buffers.len()
            }
        }
    }

    /// Try to push event to specified buffer
    pub async fn try_push(
        &self,
        buffer_index: usize,
        event: TradingEvent,
    ) -> Result<EventSequence, EventProcessingError> {
        if buffer_index >= self.buffers.len() {
            return Err(EventProcessingError::Configuration(format!(
                "Buffer index {} out of range",
                buffer_index
            )));
        }

        self.buffers[buffer_index].try_push(event).await
    }

    /// Drain events from specified buffer
    pub async fn drain_buffer(
        &self,
        buffer_index: usize,
        max_events: usize,
    ) -> Option<Vec<TradingEvent>> {
        if buffer_index >= self.buffers.len() {
            return None;
        }

        self.buffers[buffer_index].try_pop_batch(max_events).await
    }

    /// Drain all buffers sequentially
    pub async fn drain_all_buffers(&self) -> Result<()> {
        for buffer in &self.buffers {
            // Drain each buffer completely
            while !buffer.is_empty() {
                if let Some(events) = buffer.try_pop_batch(1000).await {
                    tracing::warn!("Dropping {} events during shutdown", events.len());
                } else {
                    break;
                }
            }
        }
        Ok(())
    }

    /// Get statistics for all buffers
    pub async fn get_all_stats(&self) -> Vec<BufferStats> {
        let mut all_stats = Vec::with_capacity(self.buffers.len());

        for buffer in &self.buffers {
            let mut stats = buffer.get_stats().await;
            stats.calculate_avg_latency();
            all_stats.push(stats);
        }

        all_stats
    }

    /// Get buffer count
    pub fn buffer_count(&self) -> usize {
        self.buffers.len()
    }

    /// Get total utilization across all buffers
    pub fn total_utilization(&self) -> f64 {
        let total_utilization: f64 = self.buffers.iter().map(|buffer| buffer.utilization()).sum();
        total_utilization / self.buffers.len() as f64
    }

    /// Select the least utilized buffer for load balancing
    fn select_least_utilized_buffer(&self) -> usize {
        let mut min_utilization = f64::MAX;
        let mut selected_buffer = 0;

        for (i, buffer) in self.buffers.iter().enumerate() {
            let utilization = buffer.utilization();
            if utilization < min_utilization {
                min_utilization = utilization;
                selected_buffer = i;
            }
        }

        selected_buffer
    }

    /// Hash thread ID for consistent buffer assignment
    fn hash_thread_id(&self, thread_id: std::thread::ThreadId) -> usize {
        // Simple hash function for thread ID
        let id_bytes = format!("{:?}", thread_id);
        let mut hash = 0_usize;
        for byte in id_bytes.bytes() {
            hash = hash.wrapping_mul(31).wrapping_add(byte as usize);
        }
        hash
    }

    /// Set load balancing strategy
    pub fn set_strategy(&mut self, strategy: LoadBalancingStrategy) {
        self.strategy = strategy;
    }

    /// Get current load balancing strategy
    pub const fn strategy(&self) -> LoadBalancingStrategy {
        self.strategy
    }
}

/// Load balancing strategies for buffer selection
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
/// LoadBalancingStrategy
/// 
/// TODO: Add detailed documentation for this enum
pub enum LoadBalancingStrategy {
    /// Simple round-robin assignment
    RoundRobin,
    /// Select buffer with lowest utilization
    LeastUtilized,
    /// Hash-based assignment for thread locality
    Hash,
}

/// Specialized buffer for sequence-ordered events
#[derive(Debug)]
pub struct SequenceOrderedBuffer {
    /// Events indexed by sequence number
    events: Vec<Option<TradingEvent>>,
    /// Expected next sequence number
    next_expected: AtomicU64,
    /// Highest sequence number seen
    highest_seen: AtomicU64,
    /// Buffer capacity
    capacity: usize,
}

impl SequenceOrderedBuffer {
    /// Create a new sequence-ordered buffer
    pub fn new(capacity: usize, start_sequence: u64) -> Self {
        let mut events = Vec::with_capacity(capacity);
        events.resize_with(capacity, || None);

        Self {
            events,
            next_expected: AtomicU64::new(start_sequence),
            highest_seen: AtomicU64::new(start_sequence.saturating_sub(1)),
            capacity,
        }
    }

    /// Insert event maintaining sequence order
    pub fn insert_ordered(&mut self, event: TradingEvent) -> Result<(), EventProcessingError> {
        let sequence = event.sequence_number().ok_or_else(|| {
            EventProcessingError::Configuration("Event missing sequence number".to_owned())
        })?;

        let index = (sequence % self.capacity as u64) as usize;

        // Check if slot is available
        if self.events[index].is_some() {
            return Err(EventProcessingError::BufferFull(
                "Sequence buffer full".to_owned(),
            ));
        }

        self.events[index] = Some(event);
        self.highest_seen.store(
            sequence.max(self.highest_seen.load(Ordering::Relaxed)),
            Ordering::Relaxed,
        );

        Ok(())
    }

    /// Extract events in sequence order
    pub fn extract_ordered(&mut self) -> Vec<TradingEvent> {
        let mut ordered_events = Vec::new();
        let mut current_seq = self.next_expected.load(Ordering::Relaxed);

        loop {
            let index = (current_seq % self.capacity as u64) as usize;

            if let Some(event) = self.events[index].take() {
                if event.sequence_number() == Some(current_seq) {
                    ordered_events.push(event);
                    current_seq += 1;
                } else {
                    // Put it back and break
                    self.events[index] = Some(event);
                    break;
                }
            } else {
                break;
            }
        }

        self.next_expected.store(current_seq, Ordering::Relaxed);
        ordered_events
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::events::event_types::TradingEvent;
    use crate::timing::HardwareTimestamp;
    use rust_decimal::Decimal;

    #[tokio::test]
    async fn test_event_ring_buffer_creation() {
        let buffer = EventRingBuffer::new(0, 1024).expect("Event ring buffer should be created successfully");
        assert_eq!(buffer.buffer_id(), 0);
        assert_eq!(buffer.capacity(), 1024);
        assert!(buffer.is_empty());
        assert!(!buffer.is_full());
    }

    #[tokio::test]
    async fn test_event_ring_buffer_push_pop() {
        let buffer = EventRingBuffer::new(0, 64).expect("Small event ring buffer should be created successfully");

        // Create test event
        let event = TradingEvent::OrderSubmitted {
            order_id: "TEST-001".to_string(),
            symbol: "EURUSD".to_string(),
            quantity: Decimal::new(100000, 0),
            price: Decimal::new(10850, 4),
            timestamp: HardwareTimestamp::now(),
            sequence_number: Some(1),
            metadata: None,
        };

        // Test push
        let result = buffer.try_push(event.clone()).await;
        assert!(result.is_ok());
        assert_eq!(buffer.len(), 1);

        // Test pop
        let popped = buffer.try_pop_batch(10).await;
        assert!(popped.is_some());
        let events = popped.expect("Popped events should contain the pushed event");
        assert_eq!(events.len(), 1);
        assert!(buffer.is_empty());
    }

    #[tokio::test]
    async fn test_buffer_manager_creation() {
        let manager = BufferManager::new(4, 1024).expect("Buffer manager should be created successfully");
        assert_eq!(manager.buffer_count(), 4);
    }

    #[tokio::test]
    async fn test_buffer_manager_selection() {
        let manager = BufferManager::new(4, 1024).expect("Buffer manager should be created successfully");

        // Test round-robin selection
        for i in 0..8 {
            let selected = manager.select_buffer();
            assert_eq!(selected, i % 4);
        }
    }

    #[tokio::test]
    async fn test_buffer_stats() {
        let buffer = EventRingBuffer::new(0, 64).expect("Small event ring buffer should be created successfully");

        let event = TradingEvent::OrderSubmitted {
            order_id: "TEST-001".to_string(),
            symbol: "EURUSD".to_string(),
            quantity: Decimal::new(100000, 0),
            price: Decimal::new(10850, 4),
            timestamp: HardwareTimestamp::now(),
            sequence_number: Some(1),
            metadata: None,
        };

        buffer.try_push(event).await.expect("Event push should succeed in performance test");

        let stats = buffer.get_stats().await;
        assert_eq!(stats.push_success_count, 1);
        assert_eq!(stats.buffer_id, 0);
        assert!(stats.current_utilization > 0.0);
    }

    #[tokio::test]
    async fn test_sequence_ordered_buffer() {
        let mut buffer = SequenceOrderedBuffer::new(10, 1);

        let event1 = TradingEvent::OrderSubmitted {
            order_id: "TEST-001".to_string(),
            symbol: "EURUSD".to_string(),
            quantity: Decimal::new(100000, 0),
            price: Decimal::new(10850, 4),
            timestamp: HardwareTimestamp::now(),
            sequence_number: Some(1),
            metadata: None,
        };

        let event2 = TradingEvent::OrderSubmitted {
            order_id: "TEST-002".to_string(),
            symbol: "EURUSD".to_string(),
            quantity: Decimal::new(100000, 0),
            price: Decimal::new(10851, 4),
            timestamp: HardwareTimestamp::now(),
            sequence_number: Some(2),
            metadata: None,
        };

        // Insert in order
        buffer.insert_ordered(event1).expect("First ordered event should insert successfully");
        buffer.insert_ordered(event2).expect("Second ordered event should insert successfully");

        // Extract in order
        let ordered_events = buffer.extract_ordered();
        assert_eq!(ordered_events.len(), 2);
        assert_eq!(ordered_events[0].sequence_number(), Some(1));
        assert_eq!(ordered_events[1].sequence_number(), Some(2));
    }
}