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