rs-zero 0.2.4

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},
    time::Duration,
};

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 before startup when configured boundaries are exceeded.
    RollingFile(RollingFileConfig),
}

/// Rolling file configuration.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RollingFileConfig {
    /// Active log file path.
    pub path: PathBuf,
    /// Maximum bytes before the active file is rotated at runtime.
    pub max_bytes: Option<u64>,
    /// Number of rotated files to retain.
    pub max_files: usize,
    /// Maximum age boundary documented for runtime rotation.
    pub max_age: Option<Duration>,
}

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(),
            max_bytes: Some(max_bytes),
            max_files,
            max_age: None,
        }
    }
}

/// 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: rolling.max_bytes.is_some() || rolling.max_age.is_some(),
            })
        }
    }
}

/// 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()))?;
        }
        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 {
                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,
}

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<()> {
        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.file.flush()?;
        self.rotate_files()?;
        self.file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(&self.config.path)?;
        self.bytes_written = 0;
        Ok(())
    }

    fn rotate_files(&self) -> io::Result<()> {
        let max_files = self.config.max_files.max(1);
        let oldest = rotated_path(&self.config.path, max_files);
        if oldest.exists() {
            fs::remove_file(&oldest)?;
        }
        for index in (1..max_files).rev() {
            let source = rotated_path(&self.config.path, index);
            if source.exists() {
                fs::rename(source, rotated_path(&self.config.path, index + 1))?;
            }
        }
        if self.config.path.exists() {
            fs::rename(&self.config.path, rotated_path(&self.config.path, 1))?;
        }
        Ok(())
    }
}

fn rotated_path(path: &Path, index: usize) -> PathBuf {
    path.with_file_name(format!(
        "{}.{index}",
        path.file_name()
            .and_then(|name| name.to_str())
            .unwrap_or("service.log")
    ))
}