sink.rs

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