solink_tracing_flat_json/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use serde::ser::SerializeMap;
4use serde::Serializer;
5use tracing::{Event, Subscriber};
6use tracing_serde::AsSerde;
7use tracing_subscriber::{
8    fmt::{format::Writer, FmtContext, FormatEvent, FormatFields, FormattedFields},
9    registry::LookupSpan,
10};
11
12/// `FormatEvent` for serializing data as JSON.
13///
14/// Adapted from the example in https://github.com/tokio-rs/tracing/issues/2670.
15///
16pub struct SolinkJsonFormat {
17    add_timestamp: bool,
18    add_target: bool,
19}
20
21impl SolinkJsonFormat {
22    pub fn new() -> Self {
23        Self {
24            add_timestamp: true,
25            add_target: true,
26        }
27    }
28
29    /// Set whether to add a timestamp to the log.
30    pub fn with_timestamp(mut self, add_timestamp: bool) -> Self {
31        self.add_timestamp = add_timestamp;
32        self
33    }
34
35    /// Set whether to add the target to the log.
36    pub fn with_target(mut self, add_target: bool) -> Self {
37        self.add_target = add_target;
38        self
39    }
40}
41
42impl Default for SolinkJsonFormat {
43    fn default() -> Self {
44        Self::new()
45    }
46}
47
48impl<S, N> FormatEvent<S, N> for SolinkJsonFormat
49where
50    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
51    N: for<'writer> FormatFields<'writer> + 'static,
52{
53    fn format_event(
54        &self,
55        ctx: &FmtContext<'_, S, N>,
56        mut writer: Writer<'_>,
57        event: &Event<'_>,
58    ) -> std::fmt::Result
59    where
60        S: Subscriber + for<'a> LookupSpan<'a>,
61    {
62        let meta = event.metadata();
63
64        let mut s = Vec::<u8>::new();
65        let mut serializer = serde_json::Serializer::new(&mut s);
66        let mut serializer_map = serializer.serialize_map(None).unwrap();
67
68        if self.add_timestamp {
69            let timestamp = chrono::Utc::now().to_rfc3339_opts(chrono::SecondsFormat::Nanos, true);
70            serializer_map
71                .serialize_entry("timestamp", &timestamp)
72                .unwrap();
73        }
74
75        serializer_map
76            .serialize_entry("level", &meta.level().as_serde())
77            .unwrap();
78
79        if self.add_target {
80            serializer_map
81                .serialize_entry("target", meta.target())
82                .unwrap();
83        }
84
85        let mut visitor = tracing_serde::SerdeMapVisitor::new(serializer_map);
86        event.record(&mut visitor);
87        let mut serializer_map = visitor.take_serializer().unwrap();
88
89        if let Some(scope) = ctx.event_scope() {
90            for (index, span) in scope.enumerate() {
91                if index == 0 {
92                    serializer_map.serialize_entry("span", span.name()).unwrap();
93                }
94
95                let ext = span.extensions();
96                if let Some(data) = ext.get::<FormattedFields<N>>() {
97                    if let serde_json::Value::Object(fields) =
98                        serde_json::from_str::<serde_json::Value>(data).unwrap()
99                    {
100                        for field in fields {
101                            serializer_map.serialize_entry(&field.0, &field.1).unwrap();
102                        }
103                    }
104                }
105            }
106        }
107
108        serializer_map.end().unwrap();
109
110        writer.write_str(std::str::from_utf8(&s).unwrap()).unwrap();
111        writeln!(writer)
112    }
113}
114
115#[cfg(test)]
116mod tests {
117
118    use std::{
119        io,
120        sync::{Arc, Mutex},
121    };
122
123    use tracing::{dispatcher, info};
124    use tracing_subscriber::{fmt::format::JsonFields, Layer, Registry};
125
126    use super::*;
127
128    #[derive(Debug, Clone)]
129    struct TestWriter {
130        data: Arc<Mutex<Vec<u8>>>,
131    }
132
133    impl TestWriter {
134        fn new() -> Self {
135            Self {
136                data: Arc::new(Mutex::new(Vec::new())),
137            }
138        }
139    }
140
141    impl io::Write for TestWriter {
142        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
143            self.data.lock().unwrap().write(buf)
144        }
145
146        fn flush(&mut self) -> io::Result<()> {
147            Ok(())
148        }
149    }
150
151    #[tokio::test]
152    async fn should_write_a_log() {
153        let writer = TestWriter::new();
154
155        let log_to_file = {
156            let writer = writer.clone();
157            tracing_subscriber::fmt::layer()
158                .event_format(SolinkJsonFormat::new().with_timestamp(false))
159                .fmt_fields(JsonFields::default())
160                .with_writer(move || writer.clone())
161        };
162
163        let subscriber = log_to_file.with_subscriber(Registry::default());
164        let dispatch = dispatcher::Dispatch::new(subscriber);
165        dispatcher::with_default(&dispatch, || {
166            let span1 = tracing::info_span!("parent", x = 7);
167            let span2 = tracing::info_span!(parent: &span1, "child", y = 9);
168
169            let _s1 = span1.enter();
170            let _s2 = span2.enter();
171
172            info!(z = 10, "Test")
173        });
174
175        let data = writer.data.lock().unwrap();
176        let data = std::str::from_utf8(&data).unwrap();
177        assert_eq!(
178            data.trim(),
179            r#"{"level":"INFO","target":"solink_tracing_flat_json::tests","message":"Test","z":10,"span":"child","y":9,"x":7}"#,
180        );
181    }
182}