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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.

use std::fs::{File, OpenOptions};
use std::io::prelude::*;
use std::sync::{Mutex, Once};
use std::{env, fmt, io};

use slog::{self, Drain, OwnedKVList, Record};

struct Serializer<'a>(&'a mut dyn std::io::Write);

impl<'a> slog::Serializer for Serializer<'a> {
    fn emit_arguments(&mut self, key: slog::Key, val: &std::fmt::Arguments<'_>) -> slog::Result {
        write!(self.0, ", {}: {}", key, val)?;
        Ok(())
    }
}

/// A logger that add a test case tag before each line of log.
struct CaseTraceLogger {
    f: Option<Mutex<File>>,
    skip_tags: Vec<&'static str>,
}

// FIXME: Remove this type when slog::Never implements Display.
#[derive(Debug)]
enum Never {}

impl fmt::Display for Never {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(self, f)
    }
}

impl CaseTraceLogger {
    fn write_log(
        w: &mut dyn std::io::Write,
        record: &Record<'_>,
        values: &OwnedKVList,
        skip_tags: &[&str],
    ) -> Result<(), std::io::Error> {
        use slog::KV;
        if skip_tags.contains(&record.tag()) {
            return Ok(());
        }

        let tag = tikv_util::get_tag_from_thread_name().map_or_else(|| "".to_owned(), |s| s + " ");
        let t = time::now();
        let time_str = time::strftime("%Y/%m/%d %H:%M:%S.%f", &t).unwrap();
        write!(
            w,
            "{}{} {}:{}: [{}] {}",
            tag,
            &time_str[..time_str.len() - 6],
            record.file().rsplit('/').next().unwrap(),
            record.line(),
            record.level(),
            record.msg(),
        )?;
        {
            let mut s = Serializer(w);
            record.kv().serialize(record, &mut s)?;
            values.serialize(record, &mut s)?;
        }
        writeln!(w)?;
        w.flush()?;
        Ok(())
    }
}

impl Drain for CaseTraceLogger {
    type Ok = ();
    type Err = Never;
    fn log(&self, record: &Record<'_>, values: &OwnedKVList) -> Result<Self::Ok, Self::Err> {
        if let Some(ref out) = self.f {
            let mut w = out.lock().unwrap();
            let _ = Self::write_log(&mut *w, record, values, &self.skip_tags);
        } else {
            let mut w = io::stderr();
            let _ = Self::write_log(&mut w, record, values, &self.skip_tags);
        }
        Ok(())
    }
}

impl Drop for CaseTraceLogger {
    fn drop(&mut self) {
        if let Some(ref w) = self.f {
            w.lock().unwrap().flush().unwrap();
        }
    }
}

// A help function to initial logger.
pub fn init_log_for_test() {
    static START: Once = Once::new();
    START.call_once(|| {
        let output = env::var("LOG_FILE").ok();
        let level = tikv_util::logger::get_level_by_string(
            &env::var("LOG_LEVEL").unwrap_or_else(|_| "debug".to_owned()),
        )
        .unwrap();
        let append_instead_truncate = env::var("LOG_APPEND").is_ok();
        let writer = output.map(|f| {
            Mutex::new(
                OpenOptions::new()
                    .create(true)
                    .write(!append_instead_truncate)
                    .truncate(!append_instead_truncate)
                    .append(append_instead_truncate)
                    .open(f)
                    .unwrap(),
            )
        });
        // We don't mind set it multiple times.
        // We hardly ever read rocksdb log in tests.
        let drainer = CaseTraceLogger {
            f: writer,
            skip_tags: vec![
                "rocksdb_log",
                "raftdb_log",
                "rocksdb_log_header",
                "raftdb_log_header",
            ],
        };

        // Default disabled log targets for test.
        let disabled_targets = vec!["tokio_core".to_owned(), "tokio_reactor".to_owned()];

        // CaseTraceLogger relies on test's thread name, however slog_async has
        // its own thread, and the name is "".
        // TODO: Enable the slog_async when the [Custom test frameworks][1] is mature,
        //       and hook the slog_async logger to every test cases.
        //
        // [1]: https://github.com/rust-lang/rfcs/blob/master/text/2318-custom-test-frameworks.md
        tikv_util::logger::init_log(
            drainer,
            level,
            false, // disable async drainer
            true,  // init std log
            disabled_targets,
            0,
        )
        .unwrap();
    });
}