rs-zero 0.2.11

Rust-first microservice framework inspired by go-zero engineering practices
Documentation
use std::{
    fs::{self, File, OpenOptions},
    io::{self, Write},
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
};

use chrono::{Local, NaiveDate};
use serde::Deserialize;

use crate::core::logging::rotation::{cleanup_rotated_files, rotate_files};
use crate::core::{CoreError, CoreResult};

/// Log writer destination.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LogWriterConfig {
    /// Standard output.
    Stdout,
    /// Standard error.
    Stderr,
    /// Append to one file.
    File(PathBuf),
    /// Append to one file and rotate while the process is running.
    RollingFile(RollingFileConfig),
}

/// Rolling policy for file-backed logs.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum RotationPolicy {
    /// Rotate when local date changes.
    #[default]
    Daily,
    /// Rotate when configured size boundary is exceeded.
    Size,
}

/// Rolling file configuration.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RollingFileConfig {
    /// Active log file path.
    pub path: PathBuf,
    /// Rotation policy.
    pub rotation: RotationPolicy,
    /// Maximum bytes before the active file is rotated at runtime. Used by size rotation.
    pub max_bytes: Option<u64>,
    /// Number of rotated files to retain. `0` disables count-based cleanup.
    pub max_files: usize,
    /// Number of days to retain rotated files. `None` disables age-based cleanup.
    pub keep_days: Option<u64>,
    /// Compress rotated files as gzip.
    pub compress: bool,
}

impl RollingFileConfig {
    /// Creates a runtime rolling file config with a size boundary.
    pub fn by_size(path: impl Into<PathBuf>, max_bytes: u64, max_files: usize) -> Self {
        Self {
            path: path.into(),
            rotation: RotationPolicy::Size,
            max_bytes: Some(max_bytes),
            max_files,
            keep_days: None,
            compress: false,
        }
    }

    /// Creates a runtime rolling file config with a daily boundary.
    pub fn daily(path: impl Into<PathBuf>, max_files: usize) -> Self {
        Self {
            path: path.into(),
            rotation: RotationPolicy::Daily,
            max_bytes: None,
            max_files,
            keep_days: None,
            compress: false,
        }
    }
}

/// Prepared writer information.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PreparedLogWriter {
    /// File path for file-backed writers.
    pub path: Option<PathBuf>,
    /// Whether rotation is configured.
    pub rotation_enabled: bool,
}

impl LogWriterConfig {
    /// Creates a file writer.
    pub fn file(path: impl Into<PathBuf>) -> Self {
        Self::File(path.into())
    }
}

/// Validates the writer and prepares file paths when necessary.
pub fn validate_writer(config: &LogWriterConfig) -> CoreResult<PreparedLogWriter> {
    match config {
        LogWriterConfig::Stdout | LogWriterConfig::Stderr => Ok(PreparedLogWriter {
            path: None,
            rotation_enabled: false,
        }),
        LogWriterConfig::File(path) => {
            open_append_file(path)?;
            Ok(PreparedLogWriter {
                path: Some(path.clone()),
                rotation_enabled: false,
            })
        }
        LogWriterConfig::RollingFile(rolling) => {
            open_append_file(&rolling.path)?;
            Ok(PreparedLogWriter {
                path: Some(rolling.path.clone()),
                rotation_enabled: true,
            })
        }
    }
}

/// Shared writer that rotates the active file while the process is running.
#[derive(Clone)]
pub struct RuntimeRollingFileWriter {
    inner: Arc<Mutex<RuntimeRollingFileState>>,
}

impl RuntimeRollingFileWriter {
    /// Opens a rolling writer from config.
    pub fn new(config: RollingFileConfig) -> CoreResult<Self> {
        if let Some(parent) = config.path.parent() {
            fs::create_dir_all(parent).map_err(|error| CoreError::Logging(error.to_string()))?;
        }
        cleanup_rotated_files(&config).map_err(|error| CoreError::Logging(error.to_string()))?;
        let file = open_append_file(&config.path)?;
        let bytes_written = file
            .metadata()
            .map_err(|error| CoreError::Logging(error.to_string()))?
            .len();
        Ok(Self {
            inner: Arc::new(Mutex::new(RuntimeRollingFileState {
                current_day: Local::now().date_naive(),
                config,
                file,
                bytes_written,
            })),
        })
    }
}

impl Write for RuntimeRollingFileWriter {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let mut state = self
            .inner
            .lock()
            .map_err(|_| io::Error::other("rolling log writer mutex poisoned"))?;
        state.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        let mut state = self
            .inner
            .lock()
            .map_err(|_| io::Error::other("rolling log writer mutex poisoned"))?;
        state.file.flush()
    }
}

pub(crate) fn open_append_file(path: &Path) -> CoreResult<File> {
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent).map_err(|error| CoreError::Logging(error.to_string()))?;
    }
    OpenOptions::new()
        .create(true)
        .append(true)
        .open(path)
        .map_err(|error| CoreError::Logging(error.to_string()))
}

struct RuntimeRollingFileState {
    config: RollingFileConfig,
    file: File,
    bytes_written: u64,
    current_day: NaiveDate,
}

impl RuntimeRollingFileState {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.rotate_if_needed(buf.len() as u64)?;
        self.file.write_all(buf)?;
        self.bytes_written = self.bytes_written.saturating_add(buf.len() as u64);
        Ok(buf.len())
    }

    fn rotate_if_needed(&mut self, incoming_bytes: u64) -> io::Result<()> {
        match self.config.rotation {
            RotationPolicy::Size => self.rotate_by_size_if_needed(incoming_bytes),
            RotationPolicy::Daily => self.rotate_by_day_if_needed(),
        }
    }

    fn rotate_by_size_if_needed(&mut self, incoming_bytes: u64) -> io::Result<()> {
        let Some(max_bytes) = self.config.max_bytes else {
            return Ok(());
        };
        if max_bytes == 0 || self.bytes_written.saturating_add(incoming_bytes) <= max_bytes {
            return Ok(());
        }
        self.rotate_active_file()
    }

    fn rotate_by_day_if_needed(&mut self) -> io::Result<()> {
        let today = Local::now().date_naive();
        if today <= self.current_day {
            return Ok(());
        }
        self.current_day = today;
        self.rotate_active_file()
    }

    fn rotate_active_file(&mut self) -> io::Result<()> {
        self.file.flush()?;
        rotate_files(&self.config)?;
        self.file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.config.path)?;
        self.bytes_written = 0;
        cleanup_rotated_files(&self.config)
    }
}