1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
//! Provides a very simple Log output implementation for testing and CLI.

use std::cell::RefCell;
use std::error::Error as StdError;
use std::io::Write;
use std::rc::Rc;
use std::sync::Once;
use std::thread_local;

/// Conveniently compact type alias for dyn Trait `std::error::Error`.
pub type Flaw = Box<dyn StdError + Send + Sync + 'static>;

struct Piccolog {
    filter: bool
}

impl Piccolog {
    fn thread_name(&self) -> Rc<String> {
        thread_local! {
            pub static TNAME: RefCell<Option<Rc<String>>> = RefCell::new(None);
        }

        let tn = TNAME.with(|c| {
            if let Some(n) = c.borrow().as_ref() {
                Some(n.clone())
            } else {
                None
            }
        });

        if let Some(n) = tn {
            return n;
        }

        TNAME.with(|c| {
            let t = std::thread::current();
            let mut tn = t.name().unwrap_or("-").to_owned();
            if tn == "tokio-runtime-worker" {
                tn = format!("tokio-w-{:?}", t.id());
            }
            *c.borrow_mut() = Some(Rc::new(tn));
            c.borrow().as_ref().unwrap().clone()
        })
    }
}

impl log::Log for Piccolog {
    fn enabled(&self, meta: &log::Metadata<'_>) -> bool {
        // If filtering, not from a known "local" target and Debug or Trace
        // level, then return false (disabled).
        !( self.filter &&
           !meta.target().starts_with("body_image") &&
           !meta.target().starts_with("barc") &&
           !meta.target().starts_with("blocking_permit") &&
           meta.level() > log::Level::Info )
    }

    fn log(&self, record: &log::Record<'_>) {
        if self.enabled(record.metadata()) {
            let tn = self.thread_name();
            writeln!(
                std::io::stderr(),
                "{:5} {} {}: {}",
                record.level(), record.target(), tn, record.args()
            ).ok();
        }
    }

    fn flush(&self) {
        std::io::stderr().flush().ok();
    }
}

/// Setup logger for a test run, if not already setup, based on TEST_LOG
/// environment variable.
///
/// Always returns true.
pub fn test_logger() -> bool {
    static TEST_LOG_INIT: Once = Once::new();

    TEST_LOG_INIT.call_once(|| {
        let level = if let Ok(l) = std::env::var("TEST_LOG") {
            l.parse().expect("TEST_LOG parse integer")
        } else {
            0
        };
        if level > 0 {
            setup_logger(level - 1).expect("setup logger");
        }
    });
    true
}

/// Setup logger based on specified level.
///
/// Will fail if already setup.
pub fn setup_logger(level: u32) -> Result<(), Flaw> {
    if level == 0 {
        log::set_max_level(log::LevelFilter::Info)
    } else if level < 3 {
        log::set_max_level(log::LevelFilter::Debug)
    } else {
        log::set_max_level(log::LevelFilter::Trace)
    }

    let filter = level == 1;
    log::set_boxed_logger(Box::new(Piccolog { filter }))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::test_logger;
    use log::debug;

    #[test]
    fn log_setup() {
        assert!(test_logger());
        debug!("log message");
        debug!("log message 2");
    }
}