technimal_logger/

logger.rs

use crossbeam_channel::{unbounded, Sender};
use once_cell::sync::Lazy;
use std::{
    io::stderr,
    sync::{
        atomic::{AtomicU64, AtomicUsize, Ordering},
        Mutex,
    },
    thread,
    time::{Duration, Instant},
};

static LOG_MESSAGE_BUFFER_SIZE: Lazy<AtomicUsize> = Lazy::new(|| AtomicUsize::new(1_000_000));
static LOG_MESSAGE_FLUSH_INTERVAL: Lazy<AtomicU64> = Lazy::new(|| AtomicU64::new(500_000_000));
pub static LOGGER_HANDLER: Lazy<Mutex<Option<thread::JoinHandle<()>>>> =
    Lazy::new(|| Mutex::new(None));

pub static LOG_SENDER: Lazy<Sender<LogMessage>> = Lazy::new(|| {
    let (sender, receiver) = unbounded();
    let mut message_queue: Vec<String> =
        Vec::with_capacity(LOG_MESSAGE_BUFFER_SIZE.load(Ordering::SeqCst));
    let msg_buffer_size = LOG_MESSAGE_BUFFER_SIZE.load(Ordering::SeqCst);
    let mut last_flush_time = Instant::now();

    *LOGGER_HANDLER.lock().unwrap() = Some(thread::spawn(move || {
        use humantime::format_rfc3339_micros;
        while let Ok(msg) = receiver.recv() {
            match msg {
                LogMessage::LazyMessage(lazy_message) => {
                    let message = lazy_message.eval();
                    let new_msg_length = message.len();
                    let buffer_size = message_queue.len();
                    let timestamp = Instant::now();
                    message_queue.push(message);

                    if (buffer_size + new_msg_length > msg_buffer_size)
                        || (timestamp - last_flush_time > Duration::from_secs(1))
                    // flush if the buffer is full or the time interval is passed
                    {
                        let output = message_queue.join("");
                        let _ = std::io::Write::write_all(
                            &mut stderr(),
                            output.as_bytes(),
                        );

                        message_queue.clear();
                        last_flush_time = Instant::now();
                    }
                }
                LogMessage::FlushingMessage(lazy_message) => {
                    let message = lazy_message.eval();
                    message_queue.push(message);

                    let output = message_queue.join("");
                    println!("{}", output);

                    message_queue.clear();
                    last_flush_time = Instant::now();
                }
                LogMessage::Flush => {
                    let output = message_queue.join("");
                    println!("{}", output);
                    message_queue.clear();
                    last_flush_time = Instant::now();
                }
                LogMessage::Close => {
                    let output = message_queue.join("");
                    println!("{}", output);
                    break;
                }
            }
        }
    }));
    sender
});

pub enum LogLevel {
    Error = 1,
    Warn = 2,
    Info = 3,
    Debug = 4,
    Trace = 5,
}

impl LogLevel {
    pub fn as_str(&self) -> &str {
        match self {
            LogLevel::Trace => "TRACE",
            LogLevel::Debug => "DEBUG",
            LogLevel::Info => "INFO ",
            LogLevel::Error => "ERROR",
            LogLevel::Warn => "WARN ",
        }
    }
}

pub struct LoggerGuard;

impl Drop for LoggerGuard {
    fn drop(&mut self) {
        Logger::finalize();
    }
}

pub struct Logger {}

impl Logger {
    pub fn finalize() {
        let _ = LOG_SENDER.try_send(LogMessage::Close);
        if let Some(handler) = LOGGER_HANDLER.lock().unwrap().take() {
            let _ = handler.join();
        }
    }

    pub fn initialize() -> Logger {
        LOG_MESSAGE_BUFFER_SIZE.store(1_000_000, Ordering::Relaxed);
        LOG_MESSAGE_FLUSH_INTERVAL.store(1_000_000, Ordering::Relaxed);
        Logger {}
    }

    pub fn launch(self) -> LoggerGuard {
        LoggerGuard {}
    }
}

pub enum LogMessage {
    LazyMessage(LazyMessage),
    FlushingMessage(LazyMessage),
    Flush,
    Close,
}

pub struct LazyMessage {
    data: Box<dyn (FnOnce() -> String) + Send + 'static>,
}

impl LazyMessage {
    pub fn new<F>(data: F) -> LazyMessage
    where
        F: (FnOnce() -> String) + Send + 'static,
    {
        LazyMessage {
            data: Box::new(data),
        }
    }

    pub fn eval(self) -> String {
        (self.data)()
    }
}

#[macro_export]
macro_rules! flush {
    () => {{
        $crate::LOG_SENDER
            .try_send($crate::LogMessage::Flush)
            .unwrap();
    }};
}

#[macro_export]
macro_rules! log_with_level {
    ($level:expr, $fmt:expr, $($arg:expr),* $(,)?) => {{
        $crate::log_fn!($level, $fmt, $($arg),*);
    }};

    ($level:expr, $msg:expr) => {{
        $crate::log_fn!($level, $msg);
    }};
}

#[macro_export]
macro_rules! trace {
    ( $($args:tt)* ) => {
        $crate::log_with_level!($crate::LogLevel::Trace, $($args)* )
    };
}

#[macro_export]
macro_rules! debug {
    ( $($args:tt)* ) => {
        $crate::log_with_level!($crate::LogLevel::Debug, $($args)* )
    };
}

#[macro_export]
macro_rules! info {
    ( $($args:tt)* ) => {
        $crate::log_with_level!($crate::LogLevel::Info, $($args)* )
    };
}

#[macro_export]
macro_rules! warn {
    ( $($args:tt)* ) => {
        $crate::log_with_level!($crate::LogLevel::Warn, $($args)* )
    };
}

#[macro_export]
macro_rules! error {
    ( $($args:tt)* ) => {
        $crate::log_with_level!($crate::LogLevel::Error, $($args)* )
    };
}

#[macro_export]
macro_rules! log_fn {
    ($level:expr, $fmt:expr, $($arg:expr),* $(,)?) => {{
        let timestamp = humantime::format_rfc3339_micros(std::time::SystemTime::now()).to_string();
        let func = move || {
            let cmd = $crate::concat_strs_derive::concat_strs!(
                &timestamp, ' ',
                $level.as_str(),
                " [",
                &format!("{}:{}", file!(), line!()),
                "] ",
                &format!($fmt, $($arg),*),
                "\n"
            );
            cmd
        };
        $crate::LOG_SENDER.try_send($crate::LogMessage::LazyMessage($crate::LazyMessage::new(func))).unwrap();
    }};

    ($level:expr, $msg:expr) => {{
        let timestamp = humantime::format_rfc3339_micros(std::time::SystemTime::now()).to_string();
        let func = move || {
            let cmd = $crate::concat_strs_derive::concat_strs!(
                &timestamp, ' ',
                $level.as_str(),
                " [",
                &format!("{}:{}", file!(), line!()),
                "] ",
                $msg,
                "\n"
            );
            cmd
        };
        $crate::LOG_SENDER.try_send($crate::LogMessage::LazyMessage($crate::LazyMessage::new(func))).unwrap();
    }};
}

#[cfg(test)]
mod tests {
    use crate::Logger;

    #[test]
    fn test_logger() {
        let _guard = Logger::initialize().launch();

        error!("Hello");
        error!("Hello");
        error!("Hello");
        error!("Hello");
        error!("Hello");
        error!("Hello");
        error!("Hello");
        let cmd = "x";
        let v = "y";
        error!("{} -> {}{v}", cmd, v);
        error!("Hello2");
        error!("Hello2");
        error!("Hello2");
        error!("Hello2");
        error!("Hello2");
        warn!("Hello");
        info!("Hello");
        debug!("Hello");
        debug!("message {} {}", 1, 2);
        trace!("Hello");

        // drop(_guard);
    }
}