rusty_engine/monitoring/

writer.rs

//! Memory-mapped file writer for efficient metric persistence
//!
//! This module provides a memory-mapped file writer that allows for
//! low-latency metric persistence without system call overhead.

use crate::monitoring::metrics::{AggregatedMetric, Metric, MetricsSnapshot};
use crate::monitoring::zerocopy_writer::{MonitoringRecordHeader, ZeroCopyMonitoringWriter};
use anyhow::{Result, anyhow};
use bincode;
use lz4;
use memmap2::{MmapMut, MmapOptions};
use parking_lot::Mutex;
use rusty_common::collections::FxHashMap;
use smartstring::alias::String;
use std::fs::OpenOptions;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};

/// Memory-mapped file writer for metrics
pub struct MmapWriter {
    /// Path to the memory-mapped file
    path: PathBuf,
    /// Memory-mapped region
    mmap: Mutex<MmapMut>,
    /// Current write position
    write_pos: AtomicUsize,
    /// File size
    file_size: usize,
    /// Enable compression
    enable_compression: bool,
}

impl MmapWriter {
    /// Create a new memory-mapped file writer
    pub fn new(path: impl AsRef<Path>, file_size: usize, enable_compression: bool) -> Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Create directory if it doesn't exist
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        // Open or create the file
        let file = OpenOptions::new()
            .read(true)
            .write(true)
            .create(true)
            .truncate(true)
            .open(&path)?;

        // Set file size
        file.set_len(file_size as u64)?;

        // Create memory map
        let mmap = unsafe { MmapOptions::new().map_mut(&file)? };

        Ok(Self {
            path,
            mmap: Mutex::new(mmap),
            write_pos: AtomicUsize::new(0),
            file_size,
            enable_compression,
        })
    }

    /// Write a metrics snapshot to the memory-mapped file
    pub fn write_snapshot(&self, snapshot: &MetricsSnapshot) -> Result<usize> {
        // Serialize the snapshot
        let data = if self.enable_compression {
            // Compress with LZ4
            let serialized = bincode::encode_to_vec(snapshot, bincode::config::standard())?;
            lz4::block::compress(&serialized, None, false)?
        } else {
            bincode::encode_to_vec(snapshot, bincode::config::standard())?
        };

        // Calculate header size (4 bytes for length + 1 byte for compression flag)
        let header_size = 5;
        let total_size = header_size + data.len();

        // Check if we have enough space
        let current_pos = self.write_pos.load(Ordering::Acquire);
        if current_pos + total_size > self.file_size {
            // Wrap around to the beginning
            self.write_pos.store(0, Ordering::Release);
            return Err(anyhow!("Memory-mapped file is full, wrapping around"));
        }

        // Write to memory-mapped file
        let mut mmap = self.mmap.lock();

        // Write header using zerocopy
        let header = MonitoringRecordHeader::new(data.len(), self.enable_compression);
        ZeroCopyMonitoringWriter::write_header(&mut mmap, current_pos, &header)
            .ok_or_else(|| anyhow!("Failed to write header"))?;

        // Write data
        let data_start = current_pos + header_size;
        let data_end = data_start + data.len();
        mmap[data_start..data_end].copy_from_slice(&data);

        // Update write position
        self.write_pos.store(data_end, Ordering::Release);

        // Flush to disk (async in background)
        mmap.flush_async()?;

        Ok(data.len())
    }

    /// Write raw metrics (will aggregate them first)
    pub fn write_metrics(&self, metrics: &[Metric], timestamp_ns: u64) -> Result<usize> {
        // Aggregate metrics by name and type
        let mut aggregated: FxHashMap<
            (String, crate::monitoring::metrics::MetricType),
            AggregatedMetric,
        > = FxHashMap::default();

        for metric in metrics {
            let key = (metric.name().clone(), metric.metric_type());

            aggregated
                .entry(key)
                .and_modify(|agg| agg.update(metric))
                .or_insert_with(|| AggregatedMetric::from_metric(metric));
        }

        // Create snapshot
        let snapshot = MetricsSnapshot {
            timestamp_ns,
            metrics: aggregated.into_values().collect(),
        };

        self.write_snapshot(&snapshot)
    }

    /// Get the current write position
    pub fn write_position(&self) -> usize {
        self.write_pos.load(Ordering::Acquire)
    }

    /// Reset the writer (clear the file)
    pub fn reset(&self) -> Result<()> {
        self.write_pos.store(0, Ordering::Release);

        // Clear the memory-mapped region
        let mut mmap = self.mmap.lock();
        for byte in mmap.iter_mut() {
            *byte = 0;
        }
        mmap.flush()?;

        Ok(())
    }

    /// Read snapshots from the memory-mapped file
    pub fn read_snapshots(&self) -> Result<Vec<MetricsSnapshot>> {
        let mut snapshots = Vec::new();
        let mmap = self.mmap.lock();
        let mut pos = 0;
        let end_pos = self.write_pos.load(Ordering::Acquire);

        while pos + 5 <= end_pos {
            // Read header
            let length =
                u32::from_le_bytes([mmap[pos], mmap[pos + 1], mmap[pos + 2], mmap[pos + 3]])
                    as usize;
            let compressed = mmap[pos + 4] != 0;

            pos += 5;

            if pos + length > end_pos {
                break; // Incomplete record
            }

            // Read data
            let data = &mmap[pos..pos + length];

            // Deserialize
            let snapshot = if compressed {
                let decompressed = lz4::block::decompress(data, None)?;
                bincode::decode_from_slice(&decompressed, bincode::config::standard())?.0
            } else {
                bincode::decode_from_slice(data, bincode::config::standard())?.0
            };

            snapshots.push(snapshot);
            pos += length;
        }

        Ok(snapshots)
    }
}

/// A rotating file writer that manages multiple memory-mapped files
pub struct RotatingMmapWriter {
    /// Base path for files
    base_path: PathBuf,
    /// Current file index
    current_index: AtomicUsize,
    /// Maximum number of files
    max_files: usize,
    /// File size
    file_size: usize,
    /// Enable compression
    enable_compression: bool,
    /// Current writer
    current_writer: Mutex<MmapWriter>,
}

impl RotatingMmapWriter {
    /// Create a new rotating memory-mapped file writer
    pub fn new(
        base_path: impl AsRef<Path>,
        max_files: usize,
        file_size: usize,
        enable_compression: bool,
    ) -> Result<Self> {
        let base_path = base_path.as_ref().to_path_buf();
        let current_path = Self::get_file_path(&base_path, 0);
        let current_writer = MmapWriter::new(current_path, file_size, enable_compression)?;

        Ok(Self {
            base_path,
            current_index: AtomicUsize::new(0),
            max_files,
            file_size,
            enable_compression,
            current_writer: Mutex::new(current_writer),
        })
    }

    /// Get the file path for a given index
    fn get_file_path(base_path: &Path, index: usize) -> PathBuf {
        let file_name = format!("metrics_{index:04}.bin");
        base_path.join(file_name)
    }

    /// Rotate to the next file
    fn rotate(&self) -> Result<()> {
        let next_index = (self.current_index.load(Ordering::Acquire) + 1) % self.max_files;
        let next_path = Self::get_file_path(&self.base_path, next_index);

        let new_writer = MmapWriter::new(next_path, self.file_size, self.enable_compression)?;

        let mut current_writer = self.current_writer.lock();
        *current_writer = new_writer;

        self.current_index.store(next_index, Ordering::Release);

        Ok(())
    }

    /// Write a metrics snapshot
    pub fn write_snapshot(&self, snapshot: &MetricsSnapshot) -> Result<usize> {
        // Try to write to current file
        let result = {
            let writer = self.current_writer.lock();
            writer.write_snapshot(snapshot)
        };

        match result {
            Ok(size) => Ok(size),
            Err(_) => {
                // File is full, rotate and retry
                self.rotate()?;
                let writer = self.current_writer.lock();
                writer.write_snapshot(snapshot)
            }
        }
    }

    /// Write raw metrics
    pub fn write_metrics(&self, metrics: &[Metric], timestamp_ns: u64) -> Result<usize> {
        // Try to write to current file
        let result = {
            let writer = self.current_writer.lock();
            writer.write_metrics(metrics, timestamp_ns)
        };

        match result {
            Ok(size) => Ok(size),
            Err(_) => {
                // File is full, rotate and retry
                self.rotate()?;
                let writer = self.current_writer.lock();
                writer.write_metrics(metrics, timestamp_ns)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::monitoring::metrics::{Counter, Gauge};
    use tempfile::TempDir;

    #[test]
    fn test_mmap_writer_basic() {
        let temp_dir = TempDir::new().unwrap();
        let file_path = temp_dir.path().join("test_metrics.bin");

        let writer = MmapWriter::new(&file_path, 1024 * 1024, false).unwrap();

        // Create a snapshot
        let snapshot = MetricsSnapshot {
            timestamp_ns: 1234567890,
            metrics: vec![AggregatedMetric {
                name: "test_counter".into(),
                metric_type: crate::monitoring::metrics::MetricType::Counter,
                count: 10,
                sum: 100.0,
                min: 5.0,
                max: 15.0,
                last_value: 12.0,
                start_timestamp_ns: 1000,
                end_timestamp_ns: 2000,
            }],
        };

        // Write snapshot
        let size = writer.write_snapshot(&snapshot).unwrap();
        assert!(size > 0);
        assert_eq!(writer.write_position(), size + 5); // data + header

        // Read snapshots
        let snapshots = writer.read_snapshots().unwrap();
        assert_eq!(snapshots.len(), 1);
        assert_eq!(snapshots[0].timestamp_ns, snapshot.timestamp_ns);
        assert_eq!(snapshots[0].metrics.len(), 1);
    }

    #[test]
    fn test_mmap_writer_compression() {
        let temp_dir = TempDir::new().unwrap();
        let file_path = temp_dir.path().join("test_metrics_compressed.bin");

        let writer = MmapWriter::new(&file_path, 1024 * 1024, true).unwrap();

        // Create metrics
        let metrics = vec![
            Metric::Counter(Counter {
                name: "orders".into(),
                value: 100,
                timestamp_ns: 1000,
            }),
            Metric::Gauge(Gauge {
                name: "pnl".into(),
                value: 1234.56,
                timestamp_ns: 2000,
            }),
        ];

        // Write metrics
        let size = writer.write_metrics(&metrics, 3000).unwrap();
        assert!(size > 0);

        // Read back
        let snapshots = writer.read_snapshots().unwrap();
        assert_eq!(snapshots.len(), 1);
        assert_eq!(snapshots[0].metrics.len(), 2);
    }

    #[test]
    fn test_rotating_writer() {
        let temp_dir = TempDir::new().unwrap();
        let base_path = temp_dir.path();

        let writer = RotatingMmapWriter::new(
            base_path, 3,     // max 3 files
            10240, // 10KB per file
            false,
        )
        .unwrap();

        // Write multiple snapshots
        for i in 0..5 {
            let snapshot = MetricsSnapshot {
                timestamp_ns: i as u64,
                metrics: vec![AggregatedMetric {
                    name: format!("metric_{i}"),
                    metric_type: crate::monitoring::metrics::MetricType::Counter,
                    count: 1,
                    sum: i as f64,
                    min: i as f64,
                    max: i as f64,
                    last_value: i as f64,
                    start_timestamp_ns: i as u64,
                    end_timestamp_ns: i as u64,
                }],
            };

            writer.write_snapshot(&snapshot).unwrap();
        }

        // Check that files were created
        let entries: Vec<_> = std::fs::read_dir(base_path)
            .unwrap()
            .filter_map(|e| e.ok())
            .filter(|e| e.path().extension().is_some_and(|ext| ext == "bin"))
            .collect();

        assert!(!entries.is_empty());
    }

    #[test]
    fn test_writer_reset() {
        let temp_dir = TempDir::new().unwrap();
        let file_path = temp_dir.path().join("test_reset.bin");

        let writer = MmapWriter::new(&file_path, 1024, false).unwrap();

        // Write some data
        let metrics = vec![Metric::Counter(Counter {
            name: "test".into(),
            value: 1,
            timestamp_ns: 1000,
        })];

        writer.write_metrics(&metrics, 2000).unwrap();
        assert!(writer.write_position() > 0);

        // Reset
        writer.reset().unwrap();
        assert_eq!(writer.write_position(), 0);

        // Should be able to read empty
        let snapshots = writer.read_snapshots().unwrap();
        assert_eq!(snapshots.len(), 0);
    }
}