1use std::{
2 fs,
3 io::{self, Write},
4 path::PathBuf,
5 sync::{
6 Mutex, OnceLock,
7 atomic::{AtomicU64, Ordering},
8 },
9};
10
11use crate::{SCOPE_STRING_SEP_CHAR, Scope};
12
13// ANSI color escape codes for log levels
14const ANSI_RESET: &str = "\x1b[0m";
15const ANSI_BOLD: &str = "\x1b[1m";
16const ANSI_RED: &str = "\x1b[31m";
17const ANSI_YELLOW: &str = "\x1b[33m";
18const ANSI_GREEN: &str = "\x1b[32m";
19const ANSI_BLUE: &str = "\x1b[34m";
20const ANSI_MAGENTA: &str = "\x1b[35m";
21
22/// Whether stdout output is enabled.
23static mut ENABLED_SINKS_STDOUT: bool = false;
24
25/// Is Some(file) if file output is enabled.
26static ENABLED_SINKS_FILE: Mutex<Option<std::fs::File>> = Mutex::new(None);
27static SINK_FILE_PATH: OnceLock<&'static PathBuf> = OnceLock::new();
28static SINK_FILE_PATH_ROTATE: OnceLock<&'static PathBuf> = OnceLock::new();
29/// Atomic counter for the size of the log file in bytes.
30// TODO: make non-atomic if writing single threaded
31static SINK_FILE_SIZE_BYTES: AtomicU64 = AtomicU64::new(0);
32/// Maximum size of the log file before it will be rotated, in bytes.
33const SINK_FILE_SIZE_BYTES_MAX: u64 = 1024 * 1024; // 1 MB
34
35pub fn init_output_stdout() {
36 unsafe {
37 ENABLED_SINKS_STDOUT = true;
38 }
39}
40
41pub fn init_output_file(
42 path: &'static PathBuf,
43 path_rotate: Option<&'static PathBuf>,
44) -> io::Result<()> {
45 let mut file = std::fs::OpenOptions::new()
46 .create(true)
47 .append(true)
48 .open(path)?;
49
50 SINK_FILE_PATH
51 .set(path)
52 .expect("Init file output should only be called once");
53 if let Some(path_rotate) = path_rotate {
54 SINK_FILE_PATH_ROTATE
55 .set(path_rotate)
56 .expect("Init file output should only be called once");
57 }
58
59 let mut enabled_sinks_file = ENABLED_SINKS_FILE
60 .try_lock()
61 .expect("Log file lock is available during init");
62
63 let size_bytes = file.metadata().map_or(0, |metadata| metadata.len());
64 if size_bytes >= SINK_FILE_SIZE_BYTES_MAX {
65 rotate_log_file(&mut file, Some(path), path_rotate, &SINK_FILE_SIZE_BYTES);
66 } else {
67 SINK_FILE_SIZE_BYTES.store(size_bytes, Ordering::Relaxed);
68 }
69
70 *enabled_sinks_file = Some(file);
71
72 Ok(())
73}
74
75const LEVEL_OUTPUT_STRINGS: [&str; 6] = [
76 " ", // nop: ERROR = 1
77 "ERROR", //
78 "WARN ", //
79 "INFO ", //
80 "DEBUG", //
81 "TRACE", //
82];
83
84// Colors for different log levels
85static LEVEL_ANSI_COLORS: [&str; 6] = [
86 "", // nop
87 ANSI_RED, // Error: Red
88 ANSI_YELLOW, // Warn: Yellow
89 ANSI_GREEN, // Info: Green
90 ANSI_BLUE, // Debug: Blue
91 ANSI_MAGENTA, // Trace: Magenta
92];
93
94pub fn submit(record: Record) {
95 if unsafe { ENABLED_SINKS_STDOUT } {
96 let mut stdout = std::io::stdout().lock();
97 _ = writeln!(
98 &mut stdout,
99 "{} {}{}{}{} {}[{}]{} {}",
100 chrono::Local::now().format("%Y-%m-%dT%H:%M:%S%:z"),
101 ANSI_BOLD,
102 LEVEL_ANSI_COLORS[record.level as usize],
103 LEVEL_OUTPUT_STRINGS[record.level as usize],
104 ANSI_RESET,
105 ANSI_BOLD,
106 ScopeFmt(record.scope),
107 ANSI_RESET,
108 record.message
109 );
110 }
111 let mut file = ENABLED_SINKS_FILE.lock().unwrap_or_else(|handle| {
112 ENABLED_SINKS_FILE.clear_poison();
113 handle.into_inner()
114 });
115 if let Some(file) = file.as_mut() {
116 struct SizedWriter<'a> {
117 file: &'a mut std::fs::File,
118 written: u64,
119 }
120 impl io::Write for SizedWriter<'_> {
121 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
122 self.file.write(buf)?;
123 self.written += buf.len() as u64;
124 Ok(buf.len())
125 }
126
127 fn flush(&mut self) -> io::Result<()> {
128 self.file.flush()
129 }
130 }
131 let file_size_bytes = {
132 let mut writer = SizedWriter { file, written: 0 };
133 _ = writeln!(
134 &mut writer,
135 "{} {} [{}] {}",
136 chrono::Local::now().format("%Y-%m-%dT%H:%M:%S%:z"),
137 LEVEL_OUTPUT_STRINGS[record.level as usize],
138 ScopeFmt(record.scope),
139 record.message
140 );
141 SINK_FILE_SIZE_BYTES.fetch_add(writer.written, Ordering::Relaxed) + writer.written
142 };
143 if file_size_bytes > SINK_FILE_SIZE_BYTES_MAX {
144 rotate_log_file(
145 file,
146 SINK_FILE_PATH.get(),
147 SINK_FILE_PATH_ROTATE.get(),
148 &SINK_FILE_SIZE_BYTES,
149 );
150 }
151 }
152}
153
154pub fn flush() {
155 _ = std::io::stdout().lock().flush();
156}
157
158struct ScopeFmt(Scope);
159
160impl std::fmt::Display for ScopeFmt {
161 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162 use std::fmt::Write;
163 f.write_str(self.0[0])?;
164 for scope in &self.0[1..] {
165 if !scope.is_empty() {
166 f.write_char(SCOPE_STRING_SEP_CHAR)?;
167 }
168 f.write_str(scope)?;
169 }
170 Ok(())
171 }
172}
173
174pub struct Record<'a> {
175 pub scope: Scope,
176 pub level: log::Level,
177 pub message: &'a std::fmt::Arguments<'a>,
178}
179
180fn rotate_log_file<PathRef>(
181 file: &mut fs::File,
182 path: Option<PathRef>,
183 path_rotate: Option<PathRef>,
184 atomic_size: &AtomicU64,
185) where
186 PathRef: AsRef<std::path::Path>,
187{
188 if let Err(err) = file.flush() {
189 eprintln!(
190 "Failed to flush log file before rotating, some logs may be lost: {}",
191 err
192 );
193 }
194 let rotation_error = match (path, path_rotate) {
195 (Some(_), None) => Some(anyhow::anyhow!("No rotation log file path configured")),
196 (None, _) => Some(anyhow::anyhow!("No log file path configured")),
197 (Some(path), Some(path_rotate)) => fs::copy(path, path_rotate)
198 .err()
199 .map(|err| anyhow::anyhow!(err)),
200 };
201 if let Some(err) = rotation_error {
202 eprintln!(
203 "Log file rotation failed. Truncating log file anyways: {}",
204 err,
205 );
206 }
207 _ = file.set_len(0);
208
209 // SAFETY: It is safe to set size to 0 even if set_len fails as
210 // according to the documentation, it only fails if:
211 // - the file is not writeable: should never happen,
212 // - the size would cause an overflow (implementation specific): 0 should never cause an overflow
213 atomic_size.store(0, Ordering::Relaxed);
214}
215
216#[cfg(test)]
217mod tests {
218 use super::*;
219
220 #[test]
221 fn test_rotate_log_file() {
222 let temp_dir = tempfile::tempdir().unwrap();
223 let log_file_path = temp_dir.path().join("log.txt");
224 let rotation_log_file_path = temp_dir.path().join("log_rotated.txt");
225
226 let mut file = fs::File::create(&log_file_path).unwrap();
227 let contents = String::from("Hello, world!");
228 file.write_all(contents.as_bytes()).unwrap();
229
230 let size = AtomicU64::new(contents.len() as u64);
231
232 rotate_log_file(
233 &mut file,
234 Some(&log_file_path),
235 Some(&rotation_log_file_path),
236 &size,
237 );
238
239 assert!(log_file_path.exists());
240 assert_eq!(log_file_path.metadata().unwrap().len(), 0);
241 assert!(rotation_log_file_path.exists());
242 assert_eq!(
243 std::fs::read_to_string(&rotation_log_file_path).unwrap(),
244 contents,
245 );
246 assert_eq!(size.load(Ordering::Relaxed), 0);
247 }
248
249 #[test]
250 fn test_log_level_names() {
251 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Error as usize], "ERROR");
252 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Warn as usize], "WARN ");
253 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Info as usize], "INFO ");
254 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Debug as usize], "DEBUG");
255 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Trace as usize], "TRACE");
256 }
257}