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 eprintln!("Failed to flush log file: {}", err);
200 }
201}
202
203struct SourceFmt<'a> {
204 scope: Scope,
205 module_path: Option<&'a str>,
206 ansi: bool,
207}
208
209impl std::fmt::Display for SourceFmt<'_> {
210 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
211 use std::fmt::Write;
212 f.write_char('[')?;
213 if self.ansi {
214 f.write_str(ANSI_BOLD)?;
215 }
216 // NOTE: if no longer prefixing scopes with their crate name, check if scope[0] is empty
217 if (self.scope[1].is_empty() && self.module_path.is_some()) || self.scope[0].is_empty() {
218 f.write_str(self.module_path.unwrap_or("?"))?;
219 } else {
220 f.write_str(self.scope[0])?;
221 for subscope in &self.scope[1..] {
222 if subscope.is_empty() {
223 break;
224 }
225 f.write_char(SCOPE_STRING_SEP_CHAR)?;
226 f.write_str(subscope)?;
227 }
228 }
229 if self.ansi {
230 f.write_str(ANSI_RESET)?;
231 }
232 f.write_char(']')?;
233 Ok(())
234 }
235}
236
237fn rotate_log_file<PathRef>(
238 file: &mut fs::File,
239 path: Option<PathRef>,
240 path_rotate: Option<PathRef>,
241 atomic_size: &AtomicU64,
242) where
243 PathRef: AsRef<std::path::Path>,
244{
245 if let Err(err) = file.flush() {
246 eprintln!(
247 "Failed to flush log file before rotating, some logs may be lost: {}",
248 err
249 );
250 }
251 let rotation_error = match (path, path_rotate) {
252 (Some(_), None) => Some(anyhow::anyhow!("No rotation log file path configured")),
253 (None, _) => Some(anyhow::anyhow!("No log file path configured")),
254 (Some(path), Some(path_rotate)) => fs::copy(path, path_rotate)
255 .err()
256 .map(|err| anyhow::anyhow!(err)),
257 };
258 if let Some(err) = rotation_error {
259 eprintln!("Log file rotation failed. Truncating log file anyways: {err}",);
260 }
261 _ = file.set_len(0);
262
263 // SAFETY: It is safe to set size to 0 even if set_len fails as
264 // according to the documentation, it only fails if:
265 // - the file is not writeable: should never happen,
266 // - the size would cause an overflow (implementation specific): 0 should never cause an overflow
267 atomic_size.store(0, Ordering::Relaxed);
268}
269
270#[cfg(test)]
271mod tests {
272 use super::*;
273
274 #[test]
275 fn test_rotate_log_file() {
276 let temp_dir = tempfile::tempdir().unwrap();
277 let log_file_path = temp_dir.path().join("log.txt");
278 let rotation_log_file_path = temp_dir.path().join("log_rotated.txt");
279
280 let mut file = fs::File::create(&log_file_path).unwrap();
281 let contents = String::from("Hello, world!");
282 file.write_all(contents.as_bytes()).unwrap();
283
284 let size = AtomicU64::new(contents.len() as u64);
285
286 rotate_log_file(
287 &mut file,
288 Some(&log_file_path),
289 Some(&rotation_log_file_path),
290 &size,
291 );
292
293 assert!(log_file_path.exists());
294 assert_eq!(log_file_path.metadata().unwrap().len(), 0);
295 assert!(rotation_log_file_path.exists());
296 assert_eq!(
297 std::fs::read_to_string(&rotation_log_file_path).unwrap(),
298 contents,
299 );
300 assert_eq!(size.load(Ordering::Relaxed), 0);
301 }
302
303 /// Regression test, ensuring that if log level values change we are made aware
304 #[test]
305 fn test_log_level_names() {
306 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Error as usize], "ERROR");
307 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Warn as usize], "WARN ");
308 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Info as usize], "INFO ");
309 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Debug as usize], "DEBUG");
310 assert_eq!(LEVEL_OUTPUT_STRINGS[log::Level::Trace as usize], "TRACE");
311 }
312}