metrics_tracing_context/
tracing_integration.rs
1use 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#[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#[derive(Default)]
99pub struct MetricsLayer {
100 with_labels:
101 Option<fn(&Dispatch, &Id, f: &mut dyn FnMut(&Labels) -> Option<Key>) -> Option<Key>>,
102}
103
104impl MetricsLayer {
105 pub fn new() -> Self {
107 Self::default()
108 }
109
110 pub(crate) fn with_labels(
111 &self,
112 dispatch: &Dispatch,
113 id: &Id,
114 f: &mut dyn FnMut(Map) -> Option<Key>,
115 ) -> Option<Key> {
116 let mut ff = |labels: &Labels| f(labels.0.clone());
117 (self.with_labels?)(dispatch, id, &mut ff)
118 }
119}
120
121impl<S> Layer<S> for MetricsLayer
122where
123 S: Subscriber + for<'a> LookupSpan<'a>,
124{
125 fn on_layer(&mut self, _: &mut S) {
126 self.with_labels = Some(|dispatch, id, f| {
127 let subscriber = dispatch.downcast_ref::<S>()?;
128 let span = subscriber.span(id)?;
129
130 let ext = span.extensions();
131 f(ext.get::<Labels>()?)
132 });
133 }
134
135 fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, cx: Context<'_, S>) {
136 let span = cx.span(id).expect("span must already exist!");
137 let mut labels = Labels::from_record(&Record::new(attrs.values()));
138
139 if let Some(parent) = span.parent() {
140 if let Some(parent_labels) = parent.extensions().get::<Labels>() {
141 labels.extend_from_labels(parent_labels);
142 }
143 }
144
145 span.extensions_mut().insert(labels);
146 }
147
148 fn on_record(&self, id: &Id, values: &Record<'_>, cx: Context<'_, S>) {
149 let span = cx.span(id).expect("span must already exist!");
150 let labels = Labels::from_record(values);
151
152 let ext = &mut span.extensions_mut();
153 if let Some(existing) = ext.get_mut::<Labels>() {
154 existing.extend_from_labels_overwrite(&labels);
155 } else {
156 ext.insert(labels);
157 }
158 }
159}