use std::io::Write;
use std::path::Path;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::Mutex;
use std::thread;
use tracing::Level;
use crate::base::alias::OnLog;
use crate::env::env_log::EnvLog;
pub struct TeeWriter {
file: Arc<Mutex<std::fs::File>>,
path: PathBuf,
on_log: Option<Arc<OnLog>>,
log_level: Level,
}
impl TeeWriter {
pub fn new(path: PathBuf, log_level: Level, on_log: Option<OnLog>) -> std::io::Result<Self> {
let file = std::fs::File::create(&path)?;
Ok(Self {
file: Arc::new(Mutex::new(file)),
path,
on_log: on_log.map(Arc::new),
log_level,
})
}
pub fn path(&self) -> &Path {
&self.path
}
fn emit(&self, on_log: &OnLog, line: &str) {
if self.should_print(line) {
on_log(EnvLog {
timestamp: chrono::Utc::now().to_rfc3339(),
message: line.trim_end_matches('\n').to_string(),
});
}
}
fn should_print(&self, line: &str) -> bool {
let level = if line.contains("ERROR") && line.contains("[0m") {
Level::ERROR
} else if line.contains("WARN") && line.contains("[0m") {
Level::WARN
} else if line.contains("INFO") && line.contains("[0m") {
Level::INFO
} else if line.contains("DEBUG") && line.contains("[0m") {
Level::DEBUG
} else if line.contains("TRACE") && line.contains("[0m") {
Level::TRACE
} else {
Level::TRACE
};
level <= self.log_level
}
}
impl Write for TeeWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
if let Ok(mut file) = self.file.lock() {
let _ = file.write_all(buf);
let _ = file.flush();
}
if let Some(ref on_log) = self.on_log {
let mut line_buffer = Vec::new();
for &byte in buf {
line_buffer.push(byte);
if byte == b'\n' {
let line = String::from_utf8_lossy(&line_buffer);
self.emit(on_log, &line);
line_buffer.clear();
}
}
if !line_buffer.is_empty() {
let line = String::from_utf8_lossy(&line_buffer);
self.emit(on_log, &line);
}
}
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
if let Ok(mut file) = self.file.lock() {
file.flush()
} else {
Ok(())
}
}
}
impl Clone for TeeWriter {
fn clone(&self) -> Self {
Self {
file: Arc::clone(&self.file),
path: self.path.clone(),
on_log: self.on_log.clone(),
log_level: self.log_level,
}
}
}
impl From<TeeWriter> for Stdio {
fn from(writer: TeeWriter) -> Self {
let (mut reader, writer_pipe) =
os_pipe::pipe().expect("Failed to create pipe for TeeWriter");
thread::spawn(move || {
let mut tee = writer;
let _ = std::io::copy(&mut reader, &mut tee);
});
Stdio::from(writer_pipe)
}
}