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