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
123
124
125
126
127
128
129

// #![doc(include = "../README.md")]

extern crate log;
extern crate chrono;
extern crate pretty_env_logger;

#[macro_use]
extern crate serde_json;

use std::fmt;
use std::env;

use log::{Level, Log, Metadata, Record, SetLoggerError, STATIC_MAX_LEVEL};
use serde_json::Value;
use chrono::Utc;

struct StackdriverLogger {
    service_name: String,
    service_version: String,
}

impl StackdriverLogger {
    fn format_record(&self, record: &Record) -> Value {
        json!({
            "eventTime": Utc::now().to_rfc3339(),
            "serviceContext": {
                "service": self.service_name,
                "version": self.service_version
            },
            "message": format!("{}", record.args()),
            "severity": map_level(&record.level()).to_string(),
            "reportLocation": {
                "filePath": record.file(),
                "modulePath": record.module_path(),
                // We need this or errors won't show in Error Reporting
                // There's currently no way to get the parent function name
                "functionName": record.module_path(),
                "lineNumber": record.line(),
            }
        })
    }
}

impl Log for StackdriverLogger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        metadata.level() <= STATIC_MAX_LEVEL
    }

    fn log(&self, record: &Record) {
        if self.enabled(record.metadata()) {
            let formatted = self.format_record(record);

            if record.metadata().level() == Level::Error {
                eprintln!("{}", formatted);
            } else {
                println!("{}", formatted);
            }
        }
    }

    fn flush(&self) {}
}

/// Log levels available in Stackdriver
#[derive(Debug)]
pub enum StackdriverLogLevel {
    Debug,
    Info,
    Notice,
    Warning,
    Error,
    Critical,
    Alert,
    Emergency,
}

impl fmt::Display for StackdriverLogLevel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            StackdriverLogLevel::Debug => write!(f, "DEBUG"),
            StackdriverLogLevel::Info => write!(f, "INFO"),
            StackdriverLogLevel::Notice => write!(f, "NOTICE"),
            StackdriverLogLevel::Warning => write!(f, "WARNING"),
            StackdriverLogLevel::Error => write!(f, "ERROR"),
            StackdriverLogLevel::Critical => write!(f, "CRITICAL"),
            StackdriverLogLevel::Alert => write!(f, "ALERT"),
            StackdriverLogLevel::Emergency => write!(f, "EMERGENCY"),
        }
    }
}

fn try_init() -> Result<(), SetLoggerError> {
    if cfg!(debug_assertions) {
        pretty_env_logger::init();
        Ok(())
    } else {
        let service_name = env::var("SERVICE_NAME")
            .or(env::var("CARGO_PKG_NAME"))
            .unwrap_or("".to_owned());

        let service_version = env::var("SERVICE_VERSION")
            .or(env::var("CARGO_PKG_VERSION"))
            .unwrap_or("".to_owned());

        let logger = StackdriverLogger {
            service_name,
            service_version,
        };

        log::set_max_level(STATIC_MAX_LEVEL);
        log::set_boxed_logger(Box::new(logger))
    }
}

/// Initialize the logger.
/// For debug build, this falls back to pretty_env_logger.
/// For release build, we're using the json structure expected by Stackdriver.
pub fn init() {
    try_init().expect("Could not initialize stackdriver_logger");
}

fn map_level(input: &Level) -> StackdriverLogLevel {
    match input {
        Level::Error => StackdriverLogLevel::Error,
        Level::Warn => StackdriverLogLevel::Warning,
        Level::Info => StackdriverLogLevel::Info,
        Level::Debug | Level::Trace => StackdriverLogLevel::Debug,
    }
}