metrics_tracing_context/
tracing_integration.rs

1//! The code that integrates with the `tracing` crate.
2
3use indexmap::IndexMap;
4use lockfree_object_pool::{LinearObjectPool, LinearOwnedReusable};
5use metrics::{Key, SharedString};
6use once_cell::sync::OnceCell;
7use std::cmp;
8use std::sync::Arc;
9use tracing_core::span::{Attributes, Id, Record};
10use tracing_core::{field::Visit, Dispatch, Field, Subscriber};
11use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
12
13pub(crate) type Map = IndexMap<SharedString, SharedString>;
14
15fn get_pool() -> &'static Arc<LinearObjectPool<Map>> {
16    static POOL: OnceCell<Arc<LinearObjectPool<Map>>> = OnceCell::new();
17    POOL.get_or_init(|| Arc::new(LinearObjectPool::new(Map::new, Map::clear)))
18}
19
20/// Span fields mapped as metrics labels.
21///
22/// Hidden from documentation as there is no need for end users to ever touch this type, but it must
23/// be public in order to be pulled in by external benchmark code.
24#[doc(hidden)]
25pub struct Labels(pub LinearOwnedReusable<Map>);
26
27impl Labels {
28    fn extend(&mut self, other: &Labels, f: impl Fn(&mut Map, &SharedString, &SharedString)) {
29        let new_len = cmp::max(self.as_ref().len(), other.as_ref().len());
30        let additional = new_len - self.as_ref().len();
31        self.0.reserve(additional);
32        for (k, v) in other.as_ref() {
33            f(&mut self.0, k, v);
34        }
35    }
36
37    fn extend_from_labels(&mut self, other: &Labels) {
38        self.extend(other, |map, k, v| {
39            map.entry(k.clone()).or_insert_with(|| v.clone());
40        });
41    }
42
43    fn extend_from_labels_overwrite(&mut self, other: &Labels) {
44        self.extend(other, |map, k, v| {
45            map.insert(k.clone(), v.clone());
46        });
47    }
48}
49
50impl Default for Labels {
51    fn default() -> Self {
52        Labels(get_pool().pull_owned())
53    }
54}
55
56impl Visit for Labels {
57    fn record_str(&mut self, field: &Field, value: &str) {
58        self.0.insert(field.name().into(), value.to_owned().into());
59    }
60
61    fn record_bool(&mut self, field: &Field, value: bool) {
62        self.0.insert(field.name().into(), if value { "true" } else { "false" }.into());
63    }
64
65    fn record_i64(&mut self, field: &Field, value: i64) {
66        let mut buf = itoa::Buffer::new();
67        let s = buf.format(value);
68        self.0.insert(field.name().into(), s.to_owned().into());
69    }
70
71    fn record_u64(&mut self, field: &Field, value: u64) {
72        let mut buf = itoa::Buffer::new();
73        let s = buf.format(value);
74        self.0.insert(field.name().into(), s.to_owned().into());
75    }
76
77    fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {
78        self.0.insert(field.name().into(), format!("{value:?}").into());
79    }
80}
81
82impl Labels {
83    fn from_record(record: &Record) -> Labels {
84        let mut labels = Labels::default();
85        record.record(&mut labels);
86        labels
87    }
88}
89
90impl AsRef<Map> for Labels {
91    fn as_ref(&self) -> &Map {
92        &self.0
93    }
94}
95
96/// [`MetricsLayer`] is a [`tracing_subscriber::Layer`] that captures the span
97/// fields and allows them to be later on used as metrics labels.
98#[derive(Default)]
99pub struct MetricsLayer {
100    #[allow(clippy::type_complexity)]
101    with_labels:
102        Option<fn(&Dispatch, &Id, f: &mut dyn FnMut(&Labels) -> Option<Key>) -> Option<Key>>,
103}
104
105impl MetricsLayer {
106    /// Create a new [`MetricsLayer`].
107    pub fn new() -> Self {
108        Self::default()
109    }
110
111    pub(crate) fn with_labels(
112        &self,
113        dispatch: &Dispatch,
114        id: &Id,
115        f: &mut dyn FnMut(Map) -> Option<Key>,
116    ) -> Option<Key> {
117        let mut ff = |labels: &Labels| f(labels.0.clone());
118        (self.with_labels?)(dispatch, id, &mut ff)
119    }
120}
121
122impl<S> Layer<S> for MetricsLayer
123where
124    S: Subscriber + for<'a> LookupSpan<'a>,
125{
126    fn on_layer(&mut self, _: &mut S) {
127        self.with_labels = Some(|dispatch, id, f| {
128            let subscriber = dispatch.downcast_ref::<S>()?;
129            let span = subscriber.span(id)?;
130
131            let ext = span.extensions();
132            f(ext.get::<Labels>()?)
133        });
134    }
135
136    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, cx: Context<'_, S>) {
137        let span = cx.span(id).expect("span must already exist!");
138        let mut labels = Labels::from_record(&Record::new(attrs.values()));
139
140        if let Some(parent) = span.parent() {
141            if let Some(parent_labels) = parent.extensions().get::<Labels>() {
142                labels.extend_from_labels(parent_labels);
143            }
144        }
145
146        span.extensions_mut().insert(labels);
147    }
148
149    fn on_record(&self, id: &Id, values: &Record<'_>, cx: Context<'_, S>) {
150        let span = cx.span(id).expect("span must already exist!");
151        let labels = Labels::from_record(values);
152
153        let ext = &mut span.extensions_mut();
154        if let Some(existing) = ext.get_mut::<Labels>() {
155            existing.extend_from_labels_overwrite(&labels);
156        } else {
157            ext.insert(labels);
158        }
159    }
160}
metrics_tracing_context/tracing_integration.rs

metrics_tracing_context/
tracing_integration.rs
