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