file_per_thread_logger/
lib.rs

1#[macro_use]
2extern crate log;
3extern crate env_logger;
4
5use std::cell::RefCell;
6use std::env;
7use std::fs::File;
8use std::io::{self, Write};
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::thread;
11
12use env_logger::filter::{Builder, Filter};
13use log::{LevelFilter, Metadata, Record};
14
15thread_local! {
16    static WRITER: RefCell<Option<io::BufWriter<File>>> = RefCell::new(None);
17}
18
19static ALLOW_UNINITIALIZED: AtomicBool = AtomicBool::new(false);
20
21/// Format function to print logs in a custom format.
22pub type FormatFn = fn(&mut io::BufWriter<File>, &Record) -> io::Result<()>;
23
24/// Initializes the current process/thread with a logger, parsing the RUST_LOG environment
25/// variables to set the logging level filter and/or directives to set a filter by module name,
26/// following the usual env_logger conventions.
27///
28/// Must be called on every running thread, or else logging will panic the first time it's used.
29/// ```
30/// use file_per_thread_logger::initialize;
31///
32/// initialize("log-file-prefix");
33/// ```
34pub fn initialize(filename_prefix: &str) {
35    init_logging(filename_prefix, None)
36}
37
38/// Initializes the current process/thread with a logger, parsing the RUST_LOG environment
39/// variables to set the logging level filter and/or directives to set a filter by module name,
40/// following the usual env_logger conventions. The format function specifies the format in which
41/// the logs will be printed.
42///
43/// Must be called on every running thread, or else logging will panic the first time it's used.
44/// ```
45/// use file_per_thread_logger::{initialize_with_formatter, FormatFn};
46/// use std::io::Write;
47///
48/// let formatter: FormatFn = |writer, record| {
49///     writeln!(
50///         writer,
51///         "{} [{}:{}] {}",
52///         record.level(),
53///         record.file().unwrap_or_default(),
54///         record.line().unwrap_or_default(),
55///         record.args()
56///     )
57/// };
58/// initialize_with_formatter("log-file-prefix", formatter);
59/// ```
60pub fn initialize_with_formatter(filename_prefix: &str, formatter: FormatFn) {
61    init_logging(filename_prefix, Some(formatter))
62}
63
64/// Allow logs files to be created from threads in which the logger is specifically uninitialized.
65/// It can be useful when you don't have control on threads spawned by a dependency, for instance.
66///
67/// Should be called before calling code that spawns the new threads.
68pub fn allow_uninitialized() {
69    ALLOW_UNINITIALIZED.store(true, Ordering::Relaxed);
70}
71
72fn init_logging(filename_prefix: &str, formatter: Option<FormatFn>) {
73    let env_var = env::var_os("RUST_LOG");
74    if env_var.is_none() {
75        return;
76    }
77
78    let level_filter = {
79        let mut builder = Builder::new();
80        builder.parse(env_var.unwrap().to_str().unwrap());
81        builder.build()
82    };
83
84    // Ensure the thread local state is always properly initialized.
85    WRITER.with(|rc| {
86        if rc.borrow().is_none() {
87            rc.replace(Some(open_file(filename_prefix)));
88        }
89    });
90
91    let logger = FilePerThreadLogger::new(level_filter, formatter);
92    let _ =
93        log::set_boxed_logger(Box::new(logger)).map(|()| log::set_max_level(LevelFilter::max()));
94
95    info!("Set up logging; filename prefix is {}", filename_prefix);
96}
97
98struct FilePerThreadLogger {
99    filter: Filter,
100    formatter: Option<FormatFn>,
101}
102
103impl FilePerThreadLogger {
104    pub fn new(filter: Filter, formatter: Option<FormatFn>) -> Self {
105        FilePerThreadLogger { filter, formatter }
106    }
107}
108
109impl log::Log for FilePerThreadLogger {
110    fn enabled(&self, metadata: &Metadata) -> bool {
111        self.filter.enabled(metadata)
112    }
113
114    fn log(&self, record: &Record) {
115        if self.enabled(record.metadata()) {
116            WRITER.with(|rc| {
117                if rc.borrow().is_none() && ALLOW_UNINITIALIZED.load(Ordering::Relaxed) {
118                    rc.replace(Some(open_file("")));
119                }
120                let mut opt_writer = rc.borrow_mut();
121                let writer = opt_writer
122                    .as_mut()
123                    .expect("call the logger's initialize() function first");
124                if let Some(format_fn) = &self.formatter {
125                    let _ = format_fn(&mut *writer, record);
126                } else {
127                    let _ = writeln!(*writer, "{} - {}", record.level(), record.args());
128                }
129            })
130        }
131    }
132
133    fn flush(&self) {
134        WRITER.with(|rc| {
135            let mut opt_writer = rc.borrow_mut();
136            let writer = opt_writer
137                .as_mut()
138                .expect("call the logger's initialize() function first");
139            let _ = writer.flush();
140        });
141    }
142}
143
144/// Open the tracing file for the current thread.
145fn open_file(filename_prefix: &str) -> io::BufWriter<File> {
146    let curthread = thread::current();
147    let tmpstr;
148    let mut path = filename_prefix.to_owned();
149    path.extend(
150        match curthread.name() {
151            Some(name) => name.chars(),
152            // The thread is unnamed, so use the thread ID instead.
153            None => {
154                tmpstr = format!("{:?}", curthread.id());
155                tmpstr.chars()
156            }
157        }
158        .filter(|ch| ch.is_alphanumeric() || *ch == '-' || *ch == '_'),
159    );
160    let file = File::create(path).expect("Can't open tracing file");
161    io::BufWriter::new(file)
162}
file_per_thread_logger/lib.rs

file_per_thread_logger/
lib.rs
