rusty_bin/monitor/collector/

pipeline.rs

//! Data pipeline for processing market data events and sending them to storage
//!
//! This module handles the flow of market data from the collection system
//! to the storage system with FlatBuffers serialization and compression.
//!
//! Why don't we archittect it like original pipeline?
//! Because we want to more performant. not for scalability.

use crate::monitor::collector::{CollectionError, MarketDataEvent, Result};
use crate::monitor::schema::{OrderBookRecord, TradeRecord};
use crate::monitor::storage::{CompressionMode, manager::StorageManager, roller::FileRoller};
use flume::{Receiver, Sender};
use parking_lot::RwLock;
use quanta::Instant as QuantaInstant;
use rusty_common::collections::FxHashMap;
use smartstring::alias::String as SmartString;
use std::sync::Arc;
use tokio::task::JoinHandle;
use tokio::time::{Duration, Instant};

/// Configuration for the data pipeline
#[derive(Debug, Clone)]
pub struct PipelineConfig {
    /// Buffer size for batching events before processing
    pub batch_size: usize,
    /// Maximum time to wait before flushing a batch
    pub batch_timeout_ms: u64,
    /// Number of worker tasks for processing events
    pub worker_count: usize,
    /// Compression mode for storage
    pub compression_mode: CompressionMode,
    /// File rotation check interval in seconds
    pub rotation_check_interval_seconds: u64,
}

impl Default for PipelineConfig {
    fn default() -> Self {
        Self {
            batch_size: 1000,
            batch_timeout_ms: 100,
            worker_count: 4,
            compression_mode: CompressionMode::None,
            rotation_check_interval_seconds: 300, // Check every 5 minutes
        }
    }
}

/// Statistics for pipeline processing
#[derive(Debug, Clone, Default)]
pub struct PipelineStats {
    /// Total number of events processed
    pub events_processed: u64,
    /// Number of trades processed
    pub trades_processed: u64,
    /// Number of order books processed
    pub orderbooks_processed: u64,
    /// Number of errors encountered
    pub errors_count: u64,
    /// Total bytes written to storage
    pub bytes_written: u64,
    /// Average latency in nanoseconds
    pub avg_latency_nanos: u64,
    /// Timestamp of last processed event in nanoseconds
    pub last_processed_time: Option<u64>,
    /// Number of files rotated
    pub files_rotated: u64,
    /// Number of rotation errors
    pub rotation_errors: u64,
}

/// Data pipeline that processes market data events and stores them
#[derive(Debug)]
pub struct DataPipeline {
    /// Pipeline configuration
    config: PipelineConfig,
    /// Storage manager for persisting data
    storage_manager: Arc<StorageManager>,
    /// File roller for managing file rotation
    file_roller: Option<FileRoller>,
    /// Worker task handles
    worker_handles: Vec<JoinHandle<()>>,
    /// File rotation task handle
    rotation_handle: Option<JoinHandle<()>>,
    /// Pipeline statistics
    stats: Arc<RwLock<PipelineStats>>,
    /// Shutdown signal sender
    shutdown_sender: Option<Sender<()>>,
}

impl DataPipeline {
    /// Create a new data pipeline
    #[must_use]
    pub fn new(config: PipelineConfig, storage_manager: Arc<StorageManager>) -> Self {
        // Create FileRoller with storage manager's config
        let file_roller = FileRoller::new(
            storage_manager.config.market_data_path.clone(),
            storage_manager.config.max_file_size_mb,
            storage_manager.config.max_records_per_file,
            storage_manager.config.retention_days,
        );

        Self {
            config,
            storage_manager,
            file_roller: Some(file_roller),
            worker_handles: Vec::new(),
            rotation_handle: None,
            stats: Arc::new(RwLock::new(PipelineStats::default())),
            shutdown_sender: None,
        }
    }

    /// Start the pipeline processing
    pub async fn start(&mut self, event_receiver: Receiver<MarketDataEvent>) -> Result<()> {
        log::info!(
            "Starting data pipeline with {} workers",
            self.config.worker_count
        );

        let (shutdown_sender, shutdown_receiver) = flume::unbounded();
        self.shutdown_sender = Some(shutdown_sender);

        // Start rotation watcher
        if let Err(e) = self.start_file_rotation_watcher(&shutdown_receiver) {
            log::error!("Failed to start file rotation watcher: {e}");
        }

        // Start worker tasks
        for worker_id in 0..self.config.worker_count {
            let worker_handle = self
                .spawn_worker(worker_id, event_receiver.clone(), shutdown_receiver.clone())
                .await;
            self.worker_handles.push(worker_handle);
        }

        log::info!("Data pipeline started successfully");
        Ok(())
    }

    /// Stop the pipeline processing
    pub async fn stop(&mut self) -> Result<()> {
        log::info!("Stopping data pipeline");

        // Send shutdown signal
        if let Some(shutdown_sender) = self.shutdown_sender.take() {
            let _ = shutdown_sender.send(());
        }

        // Wait for rotation task to finish
        if let Some(handle) = self.rotation_handle.take() {
            handle.abort();
            let _ = handle.await;
        }

        // Wait for all workers to finish
        for handle in self.worker_handles.drain(..) {
            handle.abort();
            let _ = handle.await;
        }

        log::info!("Data pipeline stopped");
        Ok(())
    }

    /// Start file rotation watcher that runs in background
    fn start_file_rotation_watcher(&mut self, shutdown_receiver: &Receiver<()>) -> Result<()> {
        let file_roller = self
            .file_roller
            .take()
            .ok_or(CollectionError::DataProcessing(
                "File roller not found".to_string(),
            ))?;

        let stats = self.stats.clone();
        let shutdown_receiver = shutdown_receiver.clone();
        let check_interval = Duration::from_secs(self.config.rotation_check_interval_seconds);

        let handle = tokio::spawn(async move {
            log::info!(
                "Starting file rotation watcher (check interval: {}s)",
                check_interval.as_secs()
            );

            let mut interval = tokio::time::interval(check_interval);
            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

            loop {
                tokio::select! {
                    // Check for shutdown signal
                    _ = shutdown_receiver.recv_async() => {
                        log::info!("File rotation watcher received shutdown signal");
                        break;
                    }

                    // Periodic rotation check
                    _ = interval.tick() => {
                        let rotated_files = match file_roller.execute_rotation_check().await {
                            Ok(files) => files,
                            Err(e) => {
                                log::error!("File rotation check failed: {e}");

                                // Update error stats
                                let mut stats_guard = stats.write();
                                stats_guard.rotation_errors += 1;
                                continue;
                            }
                        };

                        if !rotated_files.is_empty() {
                            let count = rotated_files.len() as u64;
                            log::info!("Rotated {count} files: {rotated_files:?}");

                            // Update stats
                            let mut stats_guard = stats.write();
                            stats_guard.files_rotated += count;
                        }

                        // Run cleanup periodically (every 10th check)
                        if stats.read().files_rotated.is_multiple_of(10) {
                            match file_roller.cleanup_old_files().await {
                                Ok(deleted_count) => {
                                    if deleted_count > 0 {
                                        log::info!("Cleaned up {deleted_count} old files");
                                    }
                                }
                                Err(e) => {
                                    log::error!("File cleanup failed: {e}");
                                }
                            }
                        }

                    }
                }
            }

            log::info!("File rotation watcher stopped");
        });

        self.rotation_handle = Some(handle);

        Ok(())
    }

    /// Get current pipeline statistics
    pub fn get_stats(&self) -> PipelineStats {
        self.stats.read().clone()
    }

    /// Reset pipeline statistics
    pub fn reset_stats(&self) {
        *self.stats.write() = PipelineStats::default();
    }

    /// Spawn a worker task to process events
    async fn spawn_worker(
        &self,
        worker_id: usize,
        event_receiver: Receiver<MarketDataEvent>,
        shutdown_receiver: Receiver<()>,
    ) -> JoinHandle<()> {
        let storage_manager = self.storage_manager.clone();
        let stats = self.stats.clone();
        let config = self.config.clone();

        tokio::spawn(async move {
            log::info!("Starting pipeline worker {worker_id}");

            let mut batch = Vec::with_capacity(config.batch_size);
            let mut last_flush = Instant::now();
            let batch_timeout = Duration::from_millis(config.batch_timeout_ms);

            loop {
                tokio::select! {
                    // Check for shutdown signal
                    _ = shutdown_receiver.recv_async() => {
                        log::info!("Worker {worker_id} received shutdown signal");
                        break;
                    }

                    // Check for timeout to flush batch
                    _ = tokio::time::sleep_until(last_flush + batch_timeout) => {
                        if !batch.is_empty() {
                            Self::process_batch(&batch, &storage_manager, &stats, &config).await;
                            batch.clear();
                            last_flush = Instant::now();
                        }
                    }

                    // Receive events
                    event_result = event_receiver.recv_async() => {
                        match event_result {
                            Ok(event) => {
                                batch.push(event);

                                // Process batch if it's full
                                if batch.len() >= config.batch_size {
                                    Self::process_batch(&batch, &storage_manager, &stats, &config).await;
                                    batch.clear();
                                    last_flush = Instant::now();
                                }
                            }
                            Err(_) => {
                                log::info!("Worker {worker_id} event channel closed");
                                break;
                            }
                        }
                    }
                }
            }

            // Process remaining events in batch
            if !batch.is_empty() {
                Self::process_batch(&batch, &storage_manager, &stats, &config).await;
            }

            log::info!("Pipeline worker {worker_id} stopped");
        })
    }

    /// Process a batch of events
    async fn process_batch(
        batch: &[MarketDataEvent],
        storage_manager: &StorageManager,
        stats: &Arc<RwLock<PipelineStats>>,
        config: &PipelineConfig,
    ) {
        let start_time = QuantaInstant::now();

        // Group events by exchange and symbol for efficient processing
        let mut trade_groups: FxHashMap<(SmartString, SmartString), Vec<&TradeRecord>> =
            FxHashMap::default();
        let mut orderbook_groups: FxHashMap<(SmartString, SmartString), Vec<&OrderBookRecord>> =
            FxHashMap::default();

        for event in batch {
            match event {
                MarketDataEvent::Trade(trade) => {
                    let key = (trade.exchange.clone(), trade.symbol.clone());
                    trade_groups.entry(key).or_default().push(trade);
                }
                MarketDataEvent::OrderBook(orderbook) => {
                    let key = (orderbook.exchange.clone(), orderbook.symbol.clone());
                    orderbook_groups.entry(key).or_default().push(orderbook);
                }
                MarketDataEvent::ConnectionStatus { .. } | MarketDataEvent::Error { .. } => {
                    // Log status/error events but don't store them
                    log::debug!("Received status/error event: {event:?}");
                }
            }
        }

        let mut total_bytes_written = 0u64;
        let mut trades_processed = 0u64;
        let mut orderbooks_processed = 0u64;
        let mut errors_count = 0u64;

        // Process trade groups
        for ((exchange, symbol), trades) in trade_groups {
            let trade_count = trades.len() as u64;
            match Self::process_trade_batch(&exchange, &symbol, trades, storage_manager, config)
                .await
            {
                Ok(bytes_written) => {
                    total_bytes_written += bytes_written;
                    trades_processed += trade_count;
                }
                Err(e) => {
                    log::error!("Failed to process trade batch for {exchange}:{symbol}: {e}");
                    errors_count += 1;
                }
            }
        }

        // Process orderbook groups
        for ((exchange, symbol), orderbooks) in orderbook_groups {
            let orderbook_count = orderbooks.len() as u64;
            match Self::process_orderbook_batch(
                &exchange,
                &symbol,
                orderbooks,
                storage_manager,
                config,
            )
            .await
            {
                Ok(bytes_written) => {
                    total_bytes_written += bytes_written;
                    orderbooks_processed += orderbook_count;
                }
                Err(e) => {
                    log::error!("Failed to process orderbook batch for {exchange}:{symbol}: {e}");
                    errors_count += 1;
                }
            }
        }

        // Update statistics
        let processing_time_nanos = start_time.elapsed().as_nanos() as u64;
        {
            let mut stats_guard = stats.write();
            stats_guard.events_processed += batch.len() as u64;
            stats_guard.trades_processed += trades_processed;
            stats_guard.orderbooks_processed += orderbooks_processed;
            stats_guard.errors_count += errors_count;
            stats_guard.bytes_written += total_bytes_written;
            stats_guard.avg_latency_nanos =
                (stats_guard.avg_latency_nanos * 9 + processing_time_nanos) / 10;
            stats_guard.last_processed_time = Some(crate::monitor::schema::timestamp::now_nanos());
        }

        log::debug!(
            "Processed batch of {} events in {}μs",
            batch.len(),
            processing_time_nanos / 1000
        );
    }

    /// Process a batch of trades for a specific exchange/symbol
    async fn process_trade_batch(
        _exchange: &str,
        _symbol: &str,
        trades: Vec<&TradeRecord>,
        storage_manager: &StorageManager,
        _config: &PipelineConfig,
    ) -> Result<u64> {
        if trades.is_empty() {
            return Ok(0);
        }

        let mut bytes_written = 0u64;

        // Write each trade to storage
        for trade in trades {
            storage_manager.write_trade(trade).await.map_err(|e| {
                CollectionError::DataProcessing(format!("Storage write failed: {e}"))
            })?;

            // Estimate bytes written (approximate)
            bytes_written += std::mem::size_of::<TradeRecord>() as u64;
        }

        Ok(bytes_written)
    }

    /// Process a batch of orderbooks for a specific exchange/symbol
    async fn process_orderbook_batch(
        _exchange: &str,
        _symbol: &str,
        orderbooks: Vec<&OrderBookRecord>,
        storage_manager: &StorageManager,
        _config: &PipelineConfig,
    ) -> Result<u64> {
        if orderbooks.is_empty() {
            return Ok(0);
        }

        let mut bytes_written = 0u64;

        // Write each orderbook to storage
        for orderbook in orderbooks {
            storage_manager
                .write_orderbook(orderbook)
                .await
                .map_err(|e| {
                    CollectionError::DataProcessing(format!("Storage write failed: {e}"))
                })?;

            // Estimate bytes written (approximate)
            bytes_written += std::mem::size_of::<OrderBookRecord>() as u64;
        }

        Ok(bytes_written)
    }
}

impl Drop for DataPipeline {
    fn drop(&mut self) {
        // Send shutdown signal if still available
        if let Some(shutdown_sender) = self.shutdown_sender.take() {
            let _ = shutdown_sender.try_send(());
        }

        // Abort all worker handles
        for handle in &self.worker_handles {
            handle.abort();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::monitor::config::storage::StorageConfig;
    use tempfile::TempDir;

    #[tokio::test]
    async fn test_pipeline_creation() {
        let temp_dir = TempDir::new().unwrap();
        let storage_config = StorageConfig {
            market_data_path: temp_dir.path().join("market_data"),
            ..Default::default()
        };
        let storage_manager = Arc::new(StorageManager::new(storage_config).await.unwrap());

        let config = PipelineConfig::default();
        let pipeline = DataPipeline::new(config, storage_manager);

        assert_eq!(pipeline.worker_handles.len(), 0);
        assert_eq!(pipeline.get_stats().events_processed, 0);
    }

    #[tokio::test]
    async fn test_pipeline_stats() {
        let temp_dir = TempDir::new().unwrap();
        let storage_config = StorageConfig {
            market_data_path: temp_dir.path().join("market_data"),
            ..Default::default()
        };
        let storage_manager = Arc::new(StorageManager::new(storage_config).await.unwrap());

        let config = PipelineConfig::default();
        let pipeline = DataPipeline::new(config, storage_manager);

        let initial_stats = pipeline.get_stats();
        assert_eq!(initial_stats.events_processed, 0);
        assert_eq!(initial_stats.trades_processed, 0);
        assert_eq!(initial_stats.orderbooks_processed, 0);

        pipeline.reset_stats();
        let reset_stats = pipeline.get_stats();
        assert_eq!(reset_stats.events_processed, 0);
    }
}