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