tracing_json/layers/storage/
json_storage.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::time::Instant;
4use tracing::field::{Field, Visit};
5use tracing::span::{Attributes, Record};
6use tracing::{Id, Subscriber};
7use tracing_subscriber::layer::Context;
8use tracing_subscriber::Layer;
9
10#[derive(Clone, Debug)]
11pub struct JsonStorageLayer;
12
13#[derive(Clone, Debug)]
14pub struct JsonStorage<'a> {
15    values: HashMap<&'a str, serde_json::Value>,
16}
17
18impl<'a> JsonStorage<'a> {
19    /// Get the set of stored values, as a set of keys and JSON values.
20    pub fn values(&self) -> &HashMap<&'a str, serde_json::Value> {
21        &self.values
22    }
23}
24
25/// Get a new visitor, with an empty bag of key-value pairs.
26impl Default for JsonStorage<'_> {
27    fn default() -> Self {
28        Self {
29            values: HashMap::new(),
30        }
31    }
32}
33
34/// Taken verbatim from tracing-subscriber
35impl Visit for JsonStorage<'_> {
36    /// Visit a signed 64-bit integer value.
37    fn record_i64(&mut self, field: &Field, value: i64) {
38        self.values
39            .insert(&field.name(), serde_json::Value::from(value));
40    }
41
42    /// Visit an unsigned 64-bit integer value.
43    fn record_u64(&mut self, field: &Field, value: u64) {
44        self.values
45            .insert(&field.name(), serde_json::Value::from(value));
46    }
47
48    /// Visit a boolean value.
49    fn record_bool(&mut self, field: &Field, value: bool) {
50        self.values
51            .insert(&field.name(), serde_json::Value::from(value));
52    }
53
54    /// Visit a string value.
55    fn record_str(&mut self, field: &Field, value: &str) {
56        self.values
57            .insert(&field.name(), serde_json::Value::from(value));
58    }
59
60    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
61        match field.name() {
62            // Skip fields that are actually log metadata that have already been handled
63            name if name.starts_with("log.") => (),
64            name if name.starts_with("r#") => {
65                self.values
66                    .insert(&name[2..], serde_json::Value::from(format!("{:?}", value)));
67            }
68            name => {
69                self.values
70                    .insert(name, serde_json::Value::from(format!("{:?}", value)));
71            }
72        };
73    }
74}
75
76impl<S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>> Layer<S>
77    for JsonStorageLayer
78{
79    /// Span creation.
80    /// This is the only occasion we have to store the fields attached to the span
81    /// given that they might have been borrowed from the surrounding context.
82    fn new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
83        let span = ctx.span(id).expect("Span not found, this is a bug");
84
85        // We want to inherit the fields from the parent span, if there is one.
86        let mut visitor = if let Some(parent_span) = span.parent() {
87            // Extensions can be used to associate arbitrary data to a span.
88            // We'll use it to store our representation of its fields.
89            // We create a copy of the parent visitor!
90            let mut extensions = parent_span.extensions_mut();
91            extensions
92                .get_mut::<JsonStorage>()
93                .map(|v| v.to_owned())
94                .unwrap_or_default()
95        } else {
96            JsonStorage::default()
97        };
98
99        let mut extensions = span.extensions_mut();
100
101        // Register all fields.
102        // Fields on the new span should override fields on the parent span if there is a conflict.
103        attrs.record(&mut visitor);
104        // Associate the visitor with the Span for future usage via the Span's extensions
105        extensions.insert(visitor);
106    }
107
108    fn on_record(&self, span: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
109        let span = ctx.span(span).expect("Span not found, this is a bug");
110
111        // Before you can associate a record to an existing Span, well, that Span has to be created!
112        // We can thus rely on the invariant that we always associate a JsonVisitor with a Span
113        // on creation (`new_span` method), hence it's safe to unwrap the Option.
114        let mut extensions = span.extensions_mut();
115        let visitor = extensions
116            .get_mut::<JsonStorage>()
117            .expect("Visitor not found on 'record', this is a bug");
118        // Register all new fields
119        values.record(visitor);
120    }
121
122    /// When we enter a span **for the first time** save the timestamp in its extensions.
123    fn on_enter(&self, span: &Id, ctx: Context<'_, S>) {
124        let span = ctx.span(span).expect("Span not found, this is a bug");
125
126        let mut extensions = span.extensions_mut();
127        if extensions.get_mut::<Instant>().is_none() {
128            extensions.insert(Instant::now());
129        }
130    }
131
132    /// When we close a span, register how long it took in milliseconds.
133    fn on_close(&self, span: Id, ctx: Context<'_, S>) {
134        let span = ctx.span(&span).expect("Span not found, this is a bug");
135
136        // Using a block to drop the immutable reference to extensions
137        // given that we want to borrow it mutably just below
138        let elapsed = {
139            let extensions = span.extensions();
140            extensions
141                .get::<Instant>()
142                .expect("Timestamp not found on 'record', this is a bug")
143                .elapsed()
144        };
145
146        let mut extensions_mut = span.extensions_mut();
147        let visitor = extensions_mut
148            .get_mut::<JsonStorage>()
149            .expect("Timestamp not found on 'record', this is a bug");
150
151        if let Ok(elapsed) = serde_json::to_value(elapsed.as_millis()) {
152            visitor.values.insert("elapsed_milliseconds", elapsed);
153        }
154    }
155}
tracing_json/layers/storage/json_storage.rs

tracing_json/layers/storage/
json_storage.rs
