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