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

use std::time::SystemTime;

use chrono::Utc;
use prost_types::Timestamp;
use shuttle_common::tracing::JsonVisitor;
use shuttle_proto::runtime::{LogItem, LogLevel};
use tokio::sync::mpsc::UnboundedSender;
use tracing::Subscriber;
use tracing_subscriber::Layer;

pub struct Logger {
    tx: UnboundedSender<LogItem>,
}

impl Logger {
    pub fn new(tx: UnboundedSender<LogItem>) -> Self {
        Self { tx }
    }
}

impl<S> Layer<S> for Logger
where
    S: Subscriber,
{
    fn on_event(
        &self,
        event: &tracing::Event<'_>,
        _ctx: tracing_subscriber::layer::Context<'_, S>,
    ) {
        let datetime = Utc::now();

        let item = {
            let metadata = event.metadata();
            let mut visitor = JsonVisitor::default();

            event.record(&mut visitor);

            LogItem {
                level: LogLevel::from(metadata.level()) as i32,
                timestamp: Some(Timestamp::from(SystemTime::from(datetime))),
                file: visitor.file.or_else(|| metadata.file().map(str::to_string)),
                line: visitor.line.or_else(|| metadata.line()),
                target: visitor
                    .target
                    .unwrap_or_else(|| metadata.target().to_string()),
                fields: serde_json::to_vec(&visitor.fields).unwrap(),
            }
        };

        self.tx.send(item).expect("sending log should succeed");
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use tokio::sync::mpsc;
    use tracing_subscriber::prelude::*;

    #[test]
    fn logging() {
        let (s, mut r) = mpsc::unbounded_channel();

        let logger = Logger::new(s);

        let _guard = tracing_subscriber::registry().with(logger).set_default();

        tracing::debug!("this is");
        tracing::info!("hi");
        tracing::warn!("from");
        tracing::error!("logger");

        assert_eq!(
            r.blocking_recv().map(to_tuple),
            Some(("this is".to_string(), LogLevel::Debug as i32))
        );
        assert_eq!(
            r.blocking_recv().map(to_tuple),
            Some(("hi".to_string(), LogLevel::Info as i32))
        );
        assert_eq!(
            r.blocking_recv().map(to_tuple),
            Some(("from".to_string(), LogLevel::Warn as i32))
        );
        assert_eq!(
            r.blocking_recv().map(to_tuple),
            Some(("logger".to_string(), LogLevel::Error as i32))
        );
    }

    fn to_tuple(log: LogItem) -> (String, i32) {
        let fields: serde_json::Map<String, serde_json::Value> =
            serde_json::from_slice(&log.fields).unwrap();

        let message = fields["message"].as_str().unwrap().to_owned();

        (message, log.level)
    }
}