sink.rs

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