rusty_bin/monitor/storage/

roller.rs

//! File rolling functionality for managing file rotation
//!
//! This module handles automatic file rotation based on time, size, and record count criteria.

use crate::monitor::storage::naming::SimpleDate;
use crate::monitor::storage::{FileNaming, Result, StorageError};
use crate::monitor::utils::time::TimingUtils;
use parking_lot::RwLock;
use rusty_common::collections::FxHashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::SystemTime;
use tokio::fs;
use walkdir::WalkDir;

/// File roller for managing file rotation based on configurable policies
#[derive(Debug)]
pub struct FileRoller {
    base_path: PathBuf,
    max_file_size_bytes: u64,
    max_records_per_file: u64,
    retention_days: u32,
    active_files: Arc<RwLock<FxHashMap<String, ActiveFileInfo>>>,
}

/// Information about a currently active file being written to
#[derive(Debug, Clone)]
pub struct ActiveFileInfo {
    path: PathBuf,
    created_at: SystemTime,
    last_written: SystemTime,
    size_bytes: u64,
    record_count: u64,
}

impl ActiveFileInfo {
    /// Get the current file size in bytes
    #[must_use]
    pub const fn size_bytes(&self) -> u64 {
        self.size_bytes
    }

    /// Get the current record count
    #[must_use]
    pub const fn record_count(&self) -> u64 {
        self.record_count
    }

    /// Get when the file was created
    #[must_use]
    pub const fn created_at(&self) -> SystemTime {
        self.created_at
    }

    /// Get when the file was last written to
    #[must_use]
    pub const fn last_written(&self) -> SystemTime {
        self.last_written
    }

    /// Update file statistics
    pub fn update_stats(&mut self, size_delta: u64, record_delta: u64) {
        self.size_bytes += size_delta;
        self.record_count += record_delta;
        self.last_written = SystemTime::now();
    }
}

impl FileRoller {
    /// Create a new file roller with specified policies
    #[must_use]
    pub fn new(
        base_path: PathBuf,
        max_file_size_mb: u64,
        max_records_per_file: u64,
        retention_days: u32,
    ) -> Self {
        Self {
            base_path,
            max_file_size_bytes: max_file_size_mb * 1024 * 1024,
            max_records_per_file,
            retention_days,
            active_files: Arc::new(RwLock::new(FxHashMap::default())),
        }
    }

    /// Register a new active file for rotation management
    pub fn register_active_file(&self, writer_id: String, file_path: PathBuf) {
        let file_info = ActiveFileInfo {
            path: file_path,
            created_at: SystemTime::now(),
            last_written: SystemTime::now(),
            size_bytes: 0,
            record_count: 0,
        };

        let mut active_files = self.active_files.write();
        active_files.insert(writer_id, file_info);
    }

    /// Update statistics for an active file
    pub fn update_file_statistics(&self, writer_id: &str, size_delta: u64, record_delta: u64) {
        let mut active_files = self.active_files.write();
        if let Some(file_info) = active_files.get_mut(writer_id) {
            file_info.update_stats(size_delta, record_delta);
        }
    }

    /// Check if a file needs rotation based on configured policies
    pub fn check_rotation_policy(&self, writer_id: &str) -> bool {
        let active_files = self.active_files.read();
        if let Some(file_info) = active_files.get(writer_id) {
            // Check size limit
            if file_info.size_bytes >= self.max_file_size_bytes {
                return true;
            }

            // Check record count limit
            if file_info.record_count >= self.max_records_per_file {
                return true;
            }

            // Check time-based rotation (daily)
            let current_date = SimpleDate::from_system_time(SystemTime::now());
            let file_date = SimpleDate::from_system_time(file_info.created_at);
            if current_date != file_date {
                return true;
            }
        }

        false
    }

    /// Execute rotation for all files that meet rotation criteria
    pub async fn execute_rotation_check(&self) -> Result<Vec<String>> {
        let mut rotated_files = Vec::new();

        // Get list of files that need rotation
        let files_to_rotate: Vec<(String, PathBuf)> = {
            let active_files = self.active_files.read();
            active_files
                .iter()
                .filter(|(writer_id, _)| self.check_rotation_policy(writer_id))
                .map(|(writer_id, info)| (writer_id.clone(), info.path.clone()))
                .collect()
        };

        log::info!("Rotation needed for {} files", files_to_rotate.len());

        // Rotate each file
        for (writer_id, file_path) in files_to_rotate {
            match self.execute_single_rotation(&writer_id, &file_path).await {
                Ok(new_path) => {
                    rotated_files.push(new_path);
                    log::info!("Rotated file: {} -> {}", file_path.display(), writer_id);
                }
                Err(e) => {
                    log::error!("Failed to rotate file {}: {}", file_path.display(), e);
                }
            }
        }

        Ok(rotated_files)
    }

    /// Execute rotation for a single file
    async fn execute_single_rotation(&self, writer_id: &str, file_path: &Path) -> Result<String> {
        // Parse the current filename to get components
        let filename = file_path
            .file_name()
            .and_then(|n| n.to_str())
            .ok_or_else(|| StorageError::Path("Invalid filename".to_string()))?;

        let naming = FileNaming::parse_filename(filename)
            .ok_or_else(|| StorageError::Path("Cannot parse filename".to_string()))?;

        // Create new filename with current date
        let current_date = SimpleDate::from_system_time(SystemTime::now());
        let new_naming = FileNaming::new(
            naming.exchange.clone(),
            naming.symbol.clone(),
            naming.data_type.clone(),
            current_date,
        );

        let new_filename = new_naming.filename();
        let new_path = file_path.parent().unwrap().join(&new_filename);

        // Handle filename conflicts by adding sequence numbers
        let final_path = self.resolve_filename_conflict(&new_path).await?;

        // Move the file
        fs::rename(file_path, &final_path).await?;

        // Remove from active files registry
        {
            let mut active_files = self.active_files.write();
            active_files.remove(writer_id);
        }

        // Log the rotation
        self.log_file_rotation(&final_path, &naming).await?;

        Ok(final_path.to_string_lossy().to_string())
    }

    /// Resolve filename conflicts by adding sequence numbers
    async fn resolve_filename_conflict(&self, base_path: &Path) -> Result<PathBuf> {
        if !base_path.exists() {
            return Ok(base_path.to_path_buf());
        }

        let base_name = base_path
            .file_stem()
            .and_then(|n| n.to_str())
            .ok_or_else(|| StorageError::Path("Invalid base filename".to_string()))?;

        let extension = base_path.extension().and_then(|e| e.to_str()).unwrap_or("");

        let parent = base_path
            .parent()
            .ok_or_else(|| StorageError::Path("No parent directory".to_string()))?;

        for sequence in 1..1000 {
            let new_name = if extension.is_empty() {
                format!("{base_name}_{sequence:03}")
            } else {
                format!("{base_name}_{sequence:03}.{extension}")
            };

            let new_path = parent.join(new_name);
            if !new_path.exists() {
                return Ok(new_path);
            }
        }

        Err(StorageError::FileRolling(
            "Cannot find unique filename after 999 attempts".to_string(),
        ))
    }

    /// Log information about a rotated file
    async fn log_file_rotation(&self, file_path: &Path, naming: &FileNaming) -> Result<()> {
        let metadata = fs::metadata(file_path).await?;
        let file_size_bytes = metadata.len();

        log::info!(
            "File rotation completed: {} (size: {} bytes)",
            naming.filename(),
            file_size_bytes
        );

        Ok(())
    }

    /// Clean up old files based on retention policy
    pub async fn cleanup_old_files(&self) -> Result<u64> {
        let cutoff_date = TimingUtils::days_ago(self.retention_days);

        let deleted_count = self.cleanup_directory(&self.base_path, cutoff_date).await?;

        if deleted_count > 0 {
            log::info!("Cleaned up {deleted_count} old files");
        }

        Ok(deleted_count)
    }

    /// Clean up files in a specific directory older than cutoff date
    async fn cleanup_directory(&self, directory: &Path, cutoff_date: SimpleDate) -> Result<u64> {
        if !directory.exists() {
            return Ok(0);
        }

        let mut deleted_count = 0;

        for entry in WalkDir::new(directory)
            .into_iter()
            .filter_map(|e| e.ok())
            .filter(|e| e.file_type().is_file())
        {
            if let Some(filename_str) = entry.file_name().to_str()
                && let Some(naming) = FileNaming::parse_filename(filename_str)
                && naming.date < cutoff_date
            {
                match fs::remove_file(entry.path()).await {
                    Ok(()) => {
                        deleted_count += 1;
                        log::debug!("Deleted old file: {}", entry.path().display());
                    }
                    Err(e) => {
                        log::warn!(
                            "Failed to delete old file {}: {}",
                            entry.path().display(),
                            e
                        );
                    }
                }
            }
        }

        Ok(deleted_count)
    }

    /// Get statistics for all currently active files
    pub fn get_active_files_statistics(&self) -> FxHashMap<String, ActiveFileInfo> {
        let active_files = self.active_files.read();
        active_files.clone()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    use tokio::fs;

    #[tokio::test]
    async fn test_file_roller_creation() {
        let roller = FileRoller::new(PathBuf::from("/tmp"), 1, 100, 30);

        let stats = roller.get_active_files_statistics();
        assert!(stats.is_empty());
    }

    #[tokio::test]
    async fn test_register_and_update_file() {
        let roller = FileRoller::new(PathBuf::from("/tmp"), 1, 100, 30);

        let temp_dir = tempdir().unwrap();
        let file_path = temp_dir.path().join("test_file.dat");
        fs::write(&file_path, "test data").await.unwrap();

        // Register the file
        let writer_id = "test:BTCUSDT:trades".to_string();
        roller.register_active_file(writer_id.clone(), file_path);

        // Update stats
        roller.update_file_statistics(&writer_id, 100, 1);

        let stats = roller.get_active_files_statistics();
        assert_eq!(stats.len(), 1);
        assert!(stats.contains_key(&writer_id));

        let file_info = &stats[&writer_id];
        assert_eq!(file_info.size_bytes(), 100);
        assert_eq!(file_info.record_count(), 1);
    }

    #[tokio::test]
    async fn test_rotation_policy_checks() {
        let temp_dir = tempdir().unwrap();
        let roller = FileRoller::new(temp_dir.path().to_path_buf(), 1, 100, 30);
        let writer_id = "test_writer".to_string();
        let file_path = temp_dir.path().join("test_file.dat");
        fs::write(&file_path, "initial data").await.unwrap();

        roller.register_active_file(writer_id.clone(), file_path);

        // Should not rotate initially
        assert!(!roller.check_rotation_policy(&writer_id));

        // Test size-based rotation
        roller.update_file_statistics(&writer_id, 2 * 1024 * 1024, 10); // 2MB, 10 records
        assert!(roller.check_rotation_policy(&writer_id));

        // Test record-based rotation (create new roller for clean state)
        let roller2 = FileRoller::new(temp_dir.path().to_path_buf(), 10, 50, 30);
        let file_path2 = temp_dir.path().join("test_file2.dat");
        fs::write(&file_path2, "initial data").await.unwrap();
        roller2.register_active_file(writer_id.clone(), file_path2);
        roller2.update_file_statistics(&writer_id, 1024, 100); // 1KB, 100 records
        assert!(roller2.check_rotation_policy(&writer_id));
    }

    #[tokio::test]
    async fn test_filename_conflict_resolution() {
        let temp_dir = tempdir().unwrap();
        let roller = FileRoller::new(temp_dir.path().to_path_buf(), 1, 100, 30);

        let base_path = temp_dir.path().join("test.parquet");

        // Create the base file to trigger conflict
        fs::write(&base_path, "existing file").await.unwrap();

        // Resolve conflict
        let unique_path = roller.resolve_filename_conflict(&base_path).await.unwrap();
        assert!(unique_path.to_string_lossy().contains("_001"));
        assert!(!unique_path.exists());
    }
}