use std::{
fs::{self, File},
io::{BufWriter, Write},
path::{Path, PathBuf},
sync::Arc,
time::SystemTime,
};
use chrono::{DateTime, Datelike, Local};
use parking_lot::Mutex;
use crate::{
config::{LogLevel, SizeRollingConfig, TimeRollingConfig},
error::Result,
};
use super::LogAppender;
#[derive(Debug, Clone)]
pub enum RollingStrategy {
Time(TimeRollingConfig),
Size(SizeRollingConfig),
}
#[derive(Debug)]
pub struct RollingFileAppender {
base_path: PathBuf,
current_path: Arc<Mutex<PathBuf>>,
writer: Arc<Mutex<BufWriter<File>>>,
strategy: RollingStrategy,
last_check: Arc<Mutex<SystemTime>>,
current_size: Arc<Mutex<u64>>,
}
impl RollingFileAppender {
pub fn new(base_path: PathBuf, strategy: RollingStrategy) -> Result<Self> {
if let Some(parent) = base_path.parent() {
fs::create_dir_all(parent)?;
}
let current_path = Self::get_current_path(&base_path, &strategy)?;
if let Some(parent) = current_path.parent() {
fs::create_dir_all(parent)?;
}
let file = Self::open_file(¤t_path)?;
let current_size = file.metadata()?.len();
Ok(Self {
base_path,
current_path: Arc::new(Mutex::new(current_path)),
writer: Arc::new(Mutex::new(BufWriter::new(file))),
strategy,
last_check: Arc::new(Mutex::new(SystemTime::now())),
current_size: Arc::new(Mutex::new(current_size)),
})
}
fn get_current_path(base_path: &Path, strategy: &RollingStrategy) -> Result<PathBuf> {
match strategy {
RollingStrategy::Time(config) => {
let now = Local::now();
let suffix = match config.period {
crate::config::TimePeriod::Hourly => now.format("%Y%m%d%H"),
crate::config::TimePeriod::Daily => now.format("%Y%m%d"),
crate::config::TimePeriod::Weekly => now.format("%Y%V"),
crate::config::TimePeriod::Monthly => now.format("%Y%m"),
};
let mut path = base_path.to_path_buf();
if let Some(ext) = path.extension() {
let mut new_ext = ext.to_string_lossy().to_string();
new_ext.push('.');
new_ext.push_str(&suffix.to_string());
path.set_extension(new_ext);
} else {
path.set_extension(suffix.to_string());
}
Ok(path)
}
RollingStrategy::Size(_) => Ok(base_path.to_path_buf()),
}
}
fn open_file(path: &Path) -> Result<File> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
Ok(std::fs::OpenOptions::new()
.create(true)
.write(true)
.append(true)
.open(path)?)
}
fn check_roll(&self) -> Result<bool> {
match &self.strategy {
RollingStrategy::Time(config) => {
let now = SystemTime::now();
let last_check = *self.last_check.lock();
let duration = now.duration_since(last_check).unwrap();
let should_roll = match config.period {
crate::config::TimePeriod::Hourly => duration.as_secs() >= 3600,
crate::config::TimePeriod::Daily => duration.as_secs() >= 86400,
crate::config::TimePeriod::Weekly => duration.as_secs() >= 604800,
crate::config::TimePeriod::Monthly => {
let now: DateTime<Local> = now.into();
let last: DateTime<Local> = last_check.into();
now.month() != last.month()
}
};
if should_roll {
*self.last_check.lock() = now;
}
Ok(should_roll)
}
RollingStrategy::Size(config) => {
let current_size = *self.current_size.lock();
Ok(current_size >= config.max_size)
}
}
}
fn roll(&self) -> Result<()> {
let new_path = Self::get_current_path(&self.base_path, &self.strategy)?;
let new_file = Self::open_file(&new_path)?;
{
let mut current_path = self.current_path.lock();
let mut writer = self.writer.lock();
let mut current_size = self.current_size.lock();
writer.flush()?;
*current_path = new_path;
*writer = BufWriter::new(new_file);
*current_size = 0;
}
self.cleanup()?;
Ok(())
}
fn cleanup(&self) -> Result<()> {
match &self.strategy {
RollingStrategy::Time(config) => {
let now = Local::now();
let keep_days = config.keep_days;
let dir = self.base_path.parent().unwrap_or_else(|| Path::new("."));
let prefix = self
.base_path
.file_name()
.unwrap()
.to_string_lossy()
.to_string();
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
if let Some(name) = path.file_name() {
let name = name.to_string_lossy();
if !name.starts_with(&prefix) {
continue;
}
}
let metadata = fs::metadata(&path)?;
let modified: DateTime<Local> = metadata.modified()?.into();
let days_old = now.signed_duration_since(modified).num_days();
if days_old > keep_days as i64 {
fs::remove_file(path)?;
}
}
}
RollingStrategy::Size(config) => {
let dir = self.base_path.parent().unwrap_or_else(|| Path::new("."));
let prefix = self
.base_path
.file_name()
.unwrap()
.to_string_lossy()
.to_string();
let mut files: Vec<_> = fs::read_dir(dir)?
.filter_map(|e| e.ok())
.filter(|e| {
if let Some(name) = e.path().file_name() {
name.to_string_lossy().starts_with(&prefix)
} else {
false
}
})
.collect();
if files.len() > config.max_files as usize {
files.sort_by_key(|f| f.metadata().unwrap().modified().unwrap());
for file in files.iter().take(files.len() - config.max_files as usize) {
fs::remove_file(file.path())?;
}
}
}
}
Ok(())
}
}
impl LogAppender for RollingFileAppender {
fn write(&self, _level: LogLevel, message: &str) -> Result<()> {
if self.check_roll()? {
self.roll()?;
}
let mut writer = self.writer.lock();
writer.write_all(message.as_bytes())?;
writer.write_all(b"\n")?;
*self.current_size.lock() += message.len() as u64 + 1;
Ok(())
}
fn flush(&self) -> Result<()> {
let mut writer = self.writer.lock();
writer.flush()?;
Ok(())
}
}
impl Clone for RollingFileAppender {
fn clone(&self) -> Self {
Self {
base_path: self.base_path.clone(),
current_path: Arc::clone(&self.current_path),
writer: Arc::clone(&self.writer),
strategy: self.strategy.clone(),
last_check: Arc::clone(&self.last_check),
current_size: Arc::clone(&self.current_size),
}
}
}