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