tracing_gcp_formatter/
lib.rs

1use std::{borrow::Cow, io::Write};
2
3use chrono::{DateTime, Utc};
4use tracing::{Level, Metadata, Subscriber, field::Visit, span};
5use tracing_subscriber::{Layer, fmt::MakeWriter, registry::LookupSpan};
6
7use crate::models::{Severity, SimplifiedLogEntry, SourceLocation};
8
9mod models;
10
11pub struct GCPFormattingLayer<W: for<'a> MakeWriter<'a> + 'static> {
12    make_writer: W,
13    pid: u32,
14    hostname: Option<String>,
15}
16
17impl<W> GCPFormattingLayer<W>
18where
19    W: for<'a> MakeWriter<'a> + 'static,
20{
21    pub fn new(make_writer: W) -> Self {
22        Self {
23            make_writer,
24            pid: Self::pid(),
25            hostname: Self::hostname(),
26        }
27    }
28
29    #[cfg(test)]
30    fn pid() -> u32 {
31        1
32    }
33
34    #[cfg(not(test))]
35    fn pid() -> u32 {
36        std::process::id()
37    }
38
39    fn hostname() -> Option<String> {
40        match (cfg!(test), cfg!(feature = "hostname")) {
41            (true, _) => Some("test-hostname".to_owned()),
42            (_, true) => Some(gethostname::gethostname().to_string_lossy().into_owned()),
43            (_, false) => None,
44        }
45    }
46
47    #[cfg(test)]
48    fn now() -> DateTime<Utc> {
49        chrono::DateTime::<Utc>::UNIX_EPOCH
50    }
51
52    #[cfg(not(test))]
53    fn now() -> DateTime<Utc> {
54        Utc::now()
55    }
56
57    fn emit(&self, entry: &SimplifiedLogEntry, meta: &Metadata<'_>) -> Result<(), std::io::Error> {
58        let buffer = {
59            let mut b = serde_json::to_string(entry).expect("Serializing SimplifiedLogEntry");
60            b.push('\n');
61            b
62        };
63
64        self.make_writer
65            .make_writer_for(meta)
66            .write_all(buffer.as_bytes())
67    }
68}
69
70impl<S, W> Layer<S> for GCPFormattingLayer<W>
71where
72    S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
73    W: for<'a> MakeWriter<'a> + 'static,
74{
75    fn on_event(&self, event: &tracing::Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
76        let mut visitor = EventVisitor::default();
77        event.record(&mut visitor);
78
79        let message = match (visitor.message, ctx.lookup_current()) {
80            (Some(message), Some(span)) => {
81                format!("[{} - EVENT] {}", span.name().to_uppercase(), message)
82            }
83            (_, Some(span)) => format!("[{} - EVENT]", span.name().to_uppercase()),
84            (Some(message), _) => message,
85            _ => "[EVENT]".to_owned(),
86        };
87
88        if let Some(span) = ctx.event_span(event) {
89            let mut current_span = Some(span);
90
91            while let Some(span) = current_span {
92                let extensions = span.extensions();
93                if let Some(span_fields) = extensions.get::<SpanFields>() {
94                    for (name, value) in span_fields.fields.iter() {
95                        visitor.other_fields.push((name.to_owned(), value.clone()));
96                    }
97                }
98                current_span = span.parent().and_then(|i| ctx.span(&i.id()));
99            }
100        }
101
102        let entry = SimplifiedLogEntry {
103            severity: map_severity(*event.metadata().level()),
104            time: Self::now(),
105            message: Cow::Borrowed(&message),
106            labels: map_labels(visitor.other_fields, self.pid, self.hostname.as_deref()),
107            source_location: map_source_location(
108                event.metadata().file(),
109                event.metadata().module_path(),
110                event.metadata().line(),
111            ),
112            ..Default::default()
113        };
114
115        let _ = self.emit(&entry, event.metadata());
116    }
117
118    fn on_new_span(
119        &self,
120        attrs: &span::Attributes<'_>,
121        _id: &span::Id,
122        _ctx: tracing_subscriber::layer::Context<'_, S>,
123    ) {
124        let mut visitor = EventVisitor::default();
125        attrs.record(&mut visitor);
126
127        let entry = SimplifiedLogEntry {
128            severity: map_severity(*attrs.metadata().level()),
129            time: Self::now(),
130            message: Cow::Borrowed(&format!(
131                "[{} - START]",
132                attrs.metadata().name().to_uppercase()
133            )),
134            labels: map_labels(visitor.other_fields, self.pid, self.hostname.as_deref()),
135            source_location: map_source_location(
136                attrs.metadata().file(),
137                attrs.metadata().module_path(),
138                attrs.metadata().line(),
139            ),
140            ..Default::default()
141        };
142
143        let _ = self.emit(&entry, attrs.metadata());
144    }
145
146    fn on_close(&self, id: span::Id, ctx: tracing_subscriber::layer::Context<'_, S>) {
147        let Some(span) = ctx.span(&id) else {
148            return;
149        };
150
151        let labels = {
152            let mut l = Vec::<(String, serde_json::Value)>::new();
153
154            let mut current_span_id = Some(id);
155
156            while let Some(span) = current_span_id.and_then(|i| ctx.span(&i)) {
157                let extensions = span.extensions();
158                if let Some(span_fields) = extensions.get::<SpanFields>() {
159                    for (name, value) in span_fields.fields.iter() {
160                        l.push((name.to_owned(), value.clone()));
161                    }
162                }
163
164                current_span_id = span.parent().map(|i| i.id().clone());
165            }
166
167            l
168        };
169
170        let entry = SimplifiedLogEntry {
171            severity: map_severity(*span.metadata().level()),
172            time: Self::now(),
173            message: Cow::Borrowed(&format!(
174                "[{} - END]",
175                span.metadata().name().to_uppercase()
176            )),
177            labels: map_labels(labels, self.pid, self.hostname.as_deref()),
178            source_location: map_source_location(
179                span.metadata().file(),
180                span.metadata().module_path(),
181                span.metadata().line(),
182            ),
183            ..Default::default()
184        };
185
186        let _ = self.emit(&entry, span.metadata());
187    }
188}
189
190fn map_severity(level: Level) -> Severity {
191    match level {
192        Level::TRACE => Severity::Default,
193        Level::DEBUG => Severity::Debug,
194        Level::INFO => Severity::Info,
195        Level::WARN => Severity::Warning,
196        Level::ERROR => Severity::Error,
197    }
198}
199
200fn map_source_location(
201    file: Option<&str>,
202    function: Option<&str>,
203    line: Option<u32>,
204) -> Option<SourceLocation> {
205    match (file, function, line) {
206        (None, None, None) => None,
207        _ => Some(SourceLocation {
208            file: file.map(|i| i.to_owned()),
209            function: function.map(|i| i.to_owned()),
210            line: line.map(|i| i.to_string()),
211        }),
212    }
213}
214
215fn map_labels(
216    labels: Vec<(String, serde_json::Value)>,
217    pid: u32,
218    hostname: Option<&str>,
219) -> Option<serde_json::Map<String, serde_json::Value>> {
220    let mut output = serde_json::Map::new();
221    output.insert("pid".to_owned(), serde_json::json!(pid));
222    if let Some(hostname) = hostname {
223        output.insert("hostname".to_owned(), serde_json::json!(hostname));
224    }
225    for (key, value) in labels {
226        output.insert(key, value);
227    }
228    Some(output)
229}
230
231#[derive(Default, Debug)]
232struct EventVisitor {
233    message: Option<String>,
234    other_fields: Vec<(String, serde_json::Value)>,
235}
236
237impl Visit for EventVisitor {
238    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
239        self.other_fields
240            .push((field.name().to_owned(), value.into()))
241    }
242
243    fn record_f64(&mut self, field: &tracing::field::Field, value: f64) {
244        self.other_fields
245            .push((field.name().to_owned(), value.into()))
246    }
247
248    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
249        self.other_fields
250            .push((field.name().to_owned(), value.into()))
251    }
252
253    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
254        self.other_fields
255            .push((field.name().to_owned(), value.trim_matches('"').into()))
256    }
257
258    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
259        self.other_fields
260            .push((field.name().to_owned(), value.into()))
261    }
262
263    fn record_bytes(&mut self, field: &tracing::field::Field, value: &[u8]) {
264        self.other_fields
265            .push((field.name().to_owned(), value.into()))
266    }
267
268    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn core::fmt::Debug) {
269        match field.name() {
270            "message" => self.message = Some(format!("{value:?}")),
271            name => self
272                .other_fields
273                .push((name.to_owned(), format!("{value:?}").into())),
274        }
275    }
276}
277
278pub struct SpanDataLayer {}
279
280impl SpanDataLayer {
281    pub fn new() -> Self {
282        Self {}
283    }
284}
285
286impl Default for SpanDataLayer {
287    fn default() -> Self {
288        Self::new()
289    }
290}
291
292struct SpanFields {
293    fields: Vec<(String, serde_json::Value)>,
294}
295
296impl<S> Layer<S> for SpanDataLayer
297where
298    S: Subscriber + for<'a> LookupSpan<'a>,
299{
300    fn on_new_span(
301        &self,
302        attrs: &span::Attributes<'_>,
303        id: &span::Id,
304        ctx: tracing_subscriber::layer::Context<'_, S>,
305    ) {
306        let mut visitor = EventVisitor::default();
307        attrs.record(&mut visitor);
308
309        if let Some(span) = ctx.span(id) {
310            let mut extensions = span.extensions_mut();
311
312            extensions.insert(SpanFields {
313                fields: visitor
314                    .other_fields
315                    .iter()
316                    .map(|(k, v)| (k.to_string(), v.clone()))
317                    .collect(),
318            });
319        }
320    }
321}
322
323#[cfg(test)]
324mod tests;
tracing_gcp_formatter/lib.rs

tracing_gcp_formatter/
lib.rs
