use std::fs::{File, OpenOptions};
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use anyhow::Result;
use memrecd::config::DaemonConfig;
use tracing::{info, Level};
use tracing_subscriber::fmt::format::FmtSpan;
use tracing_subscriber::fmt::time::FormatTime;
use tracing_subscriber::fmt::writer::BoxMakeWriter;
use tracing_subscriber::FmtSubscriber;
use memrecd::daemon::Daemon;
fn main() -> Result<()> {
let args: Vec<String> = std::env::args().collect();
if args.iter().any(|a| a == "--version" || a == "-V") {
println!("memrecd {}", env!("CARGO_PKG_VERSION"));
return Ok(());
}
let config = DaemonConfig::load()?;
let subscriber = build_subscriber(&config)?;
tracing::subscriber::set_global_default(subscriber)?;
info!("Starting memrecd v{}", env!("CARGO_PKG_VERSION"));
let daemon = Daemon::new()?;
let runtime = tokio::runtime::Builder::new_multi_thread()
.thread_name("tokio-worker")
.enable_all()
.build()?;
runtime.block_on(async { daemon.run().await })?;
info!("memrecd stopped");
Ok(())
}
#[derive(Clone)]
struct LocalDailyWriter {
dir: PathBuf,
prefix: String,
suffix: String,
max_files: usize,
file: Arc<Mutex<Option<(String, File)>>>,
}
impl LocalDailyWriter {
fn new(dir: PathBuf, prefix: String, suffix: String, max_files: usize) -> Self {
Self {
dir,
prefix,
suffix,
max_files,
file: Arc::new(Mutex::new(None)),
}
}
fn current_filename(&self) -> String {
format!(
"{}.{}.{}",
self.prefix,
chrono::Local::now().format("%Y-%m-%d"),
self.suffix
)
}
fn cleanup_old_files(&self) {
let Ok(entries) = std::fs::read_dir(&self.dir) else {
return;
};
let pattern_prefix = format!("{}.", self.prefix);
let pattern_suffix = format!(".{}", self.suffix);
let mut files: Vec<PathBuf> = entries
.flatten()
.map(|e| e.path())
.filter(|p| {
p.file_name()
.and_then(|n| n.to_str())
.map(|n| n.starts_with(&pattern_prefix) && n.ends_with(&pattern_suffix))
.unwrap_or(false)
})
.collect();
if files.len() <= self.max_files {
return;
}
files.sort();
let remove_count = files.len() - self.max_files;
for old in files.into_iter().take(remove_count) {
if let Err(e) = std::fs::remove_file(&old) {
eprintln!("Failed to remove old log file {}: {}", old.display(), e);
}
}
}
}
impl Write for LocalDailyWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
let filename = self.current_filename();
let mut guard = self.file.lock().unwrap_or_else(|e| e.into_inner());
match &mut *guard {
Some((cur, f)) if *cur == filename => f.write(buf),
_ => {
let path = self.dir.join(&filename);
let mut f = OpenOptions::new().create(true).append(true).open(&path)?;
let n = f.write(buf)?;
*guard = Some((filename, f));
self.cleanup_old_files();
Ok(n)
}
}
}
fn flush(&mut self) -> io::Result<()> {
if let Some((_, f)) = &mut *self.file.lock().unwrap_or_else(|e| e.into_inner()) {
f.flush()?;
}
Ok(())
}
}
#[derive(Clone, Copy)]
struct LocalTimestamp;
impl FormatTime for LocalTimestamp {
fn format_time(&self, w: &mut tracing_subscriber::fmt::format::Writer<'_>) -> std::fmt::Result {
write!(
w,
"{}",
chrono::Local::now().format("%Y-%m-%dT%H:%M:%S%.3f%:z")
)
}
}
fn build_subscriber(config: &DaemonConfig) -> Result<impl tracing::Subscriber + Send + Sync> {
let level = parse_level(&config.log.level);
let writer = if config.log.enabled {
let log_path = expand_tilde(&config.log.memrecd_log);
let (log_dir, prefix, suffix) = split_log_path(&log_path);
std::fs::create_dir_all(&log_dir)?;
let file_writer =
LocalDailyWriter::new(log_dir, prefix, suffix, config.log.rotation_days as usize);
BoxMakeWriter::new(move || file_writer.clone())
} else {
BoxMakeWriter::new(std::io::stdout)
};
Ok(FmtSubscriber::builder()
.with_max_level(level)
.with_ansi(false)
.with_writer(writer)
.with_timer(LocalTimestamp)
.with_span_events(FmtSpan::CLOSE)
.finish())
}
fn parse_level(level: &str) -> Level {
match level.to_ascii_lowercase().as_str() {
"trace" => Level::TRACE,
"debug" => Level::DEBUG,
"warn" | "warning" => Level::WARN,
"error" => Level::ERROR,
_ => Level::INFO,
}
}
fn expand_tilde(path: &str) -> PathBuf {
if let Some(rest) = path.strip_prefix('~') {
if let Some(home) = dirs::home_dir() {
return home.join(rest.strip_prefix('/').unwrap_or(rest));
}
}
PathBuf::from(path)
}
fn split_log_path(path: &Path) -> (PathBuf, String, String) {
let dir = path
.parent()
.map(Path::to_path_buf)
.unwrap_or_else(|| PathBuf::from("."));
let file_name = path
.file_name()
.and_then(|n| n.to_str())
.unwrap_or("memrecd.log")
.to_string();
let (prefix, suffix) = match file_name.rsplit_once('.') {
Some((p, s)) => (p.to_string(), s.to_string()),
None => (file_name, "log".to_string()),
};
(dir, prefix, suffix)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_local_daily_writer_creates_file() {
let dir = tempfile::tempdir().unwrap();
let mut writer = LocalDailyWriter::new(
dir.path().to_path_buf(),
"memrecd".to_string(),
"log".to_string(),
7,
);
writer.write_all(b"hello local time log\n").expect("write");
writer.flush().expect("flush");
let today = chrono::Local::now().format("%Y-%m-%d").to_string();
let expected = dir.path().join(format!("memrecd.{}.log", today));
assert!(expected.exists(), "expected {}", expected.display());
let content = std::fs::read_to_string(&expected).unwrap();
assert!(content.contains("hello local time log"));
}
#[test]
fn test_local_daily_writer_switches_date() {
let dir = tempfile::tempdir().unwrap();
let mut writer = LocalDailyWriter::new(
dir.path().to_path_buf(),
"memrecd".to_string(),
"log".to_string(),
7,
);
writer.write_all(b"line1\n").unwrap();
let yesterday = (chrono::Local::now() - chrono::Duration::days(1))
.format("%Y-%m-%d")
.to_string();
let yesterday_path = dir.path().join(format!("memrecd.{}.log", yesterday));
std::fs::write(&yesterday_path, "old\n").unwrap();
writer.write_all(b"line2\n").unwrap();
writer.flush().unwrap();
assert!(yesterday_path.exists(), "yesterday log kept");
let today = chrono::Local::now().format("%Y-%m-%d").to_string();
assert!(dir.path().join(format!("memrecd.{}.log", today)).exists());
}
#[test]
fn test_local_daily_writer_cleanup_excess_files() {
let dir = tempfile::tempdir().unwrap();
let mut writer = LocalDailyWriter::new(
dir.path().to_path_buf(),
"memrecd".to_string(),
"log".to_string(),
2,
);
for i in 0..4 {
let day = (chrono::Local::now() - chrono::Duration::days(i as i64 + 1))
.format("%Y-%m-%d")
.to_string();
let p = dir.path().join(format!("memrecd.{}.log", day));
std::fs::write(&p, format!("old {}\n", i)).unwrap();
}
writer.write_all(b"new\n").unwrap();
writer.flush().unwrap();
let today = chrono::Local::now().format("%Y-%m-%d").to_string();
let mut remaining: Vec<String> = std::fs::read_dir(dir.path())
.unwrap()
.flatten()
.map(|e| e.file_name().to_string_lossy().to_string())
.filter(|n| n.starts_with("memrecd."))
.collect();
remaining.sort();
assert!(remaining.iter().any(|n| n.contains(&today)));
assert!(remaining.len() <= 2, "remaining: {:?}", remaining);
}
}