Tests

metrics-tracing-context/Cargo.toml

@@ -37,10 +37,12 @@ once_cell = { version = "1", default-features = false, features = ["std"] }
 tracing = { version = "0.1.29", default-features = false }
 tracing-core = { version = "0.1.21", default-features = false }
 tracing-subscriber = { version = "0.3.1", default-features = false, features = ["std"] }
+tracing-test = "0.2.4"
 
 [dev-dependencies]
 criterion = { version = "=0.3.3", default-features = false }
 parking_lot = { version = "0.12.1", default-features = false }
 tracing = { version = "0.1.29", default-features = false, features = ["std"] }
-tracing-subscriber = { version = "0.3.1", default-features = false, features = ["registry"] }
-pretty_assertions = "*"
+tracing-subscriber = { version = "0.3.1", default-features = false, features = ["registry", "fmt"] }
+pretty_assertions = "*"
+itertools = "*"

metrics-tracing-context/src/lib.rs

@@ -110,8 +110,8 @@ pub mod label_filter;
 mod tracing_integration;
 
 pub use label_filter::LabelFilter;
+use tracing_integration::Map;
 pub use tracing_integration::{Labels, MetricsLayer};
-use tracing_integration::{Map, WithContext};
 
 /// [`TracingContextLayer`] provides an implementation of a [`Layer`] for [`TracingContext`].
 pub struct TracingContextLayer<F> {
@@ -176,7 +176,7 @@ where
         tracing::dispatcher::get_default(|dispatch| {
             let current = dispatch.current_span();
             let id = current.id()?;
-            let ctx = dispatch.downcast_ref::<WithContext>()?;
+            let ctx = dispatch.downcast_ref::<MetricsLayer>()?;
 
             let mut f = |mut span_labels: Map| {
                 (!span_labels.is_empty()).then(|| {

metrics-tracing-context/src/tracing_integration.rs

@@ -4,8 +4,9 @@ use indexmap::IndexMap;
 use lockfree_object_pool::{LinearObjectPool, LinearOwnedReusable};
 use metrics::{Key, SharedString};
 use once_cell::sync::OnceCell;
+use std::cmp;
 use std::sync::Arc;
-use std::{any::TypeId, cmp, marker::PhantomData};
+use tracing::span;
 use tracing_core::span::{Attributes, Id, Record};
 use tracing_core::{field::Visit, Dispatch, Field, Subscriber};
 use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
@@ -24,13 +25,23 @@ fn get_pool() -> &'static Arc<LinearObjectPool<Map>> {
 #[doc(hidden)]
 pub struct Labels(pub LinearOwnedReusable<Map>);
 
+impl std::fmt::Debug for Labels {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Labels").field(self.as_ref()).finish()
+    }
+}
+
 impl Labels {
-    pub(crate) fn extend_from_labels(&mut self, other: &Labels) {
+    pub(crate) fn extend_from_labels(&mut self, other: &Labels, overwrite: bool) {
         let new_len = cmp::max(self.as_ref().len(), other.as_ref().len());
         let additional = new_len - self.as_ref().len();
         self.0.reserve(additional);
         for (k, v) in other.as_ref() {
-            self.0.insert(k.clone(), v.clone());
+            if overwrite {
+                self.0.insert(k.clone(), v.clone());
+            } else {
+                self.0.entry(k.clone()).or_insert_with(|| v.clone());
+            }
         }
     }
 }
@@ -81,66 +92,57 @@ impl AsRef<Map> for Labels {
     }
 }
 
-pub struct WithContext {
-    with_labels: fn(&Dispatch, &Id, f: &mut dyn FnMut(&Labels) -> Option<Key>) -> Option<Key>,
-}
-
-impl WithContext {
-    pub fn with_labels(
-        &self,
-        dispatch: &Dispatch,
-        id: &Id,
-        f: &mut dyn FnMut(Map) -> Option<Key>,
-    ) -> Option<Key> {
-        let mut ff = |labels: &Labels| f(labels.0.clone());
-        (self.with_labels)(dispatch, id, &mut ff)
-    }
-}
-
 /// [`MetricsLayer`] is a [`tracing_subscriber::Layer`] that captures the span
 /// fields and allows them to be later on used as metrics labels.
-pub struct MetricsLayer<S> {
-    ctx: WithContext,
-    _subscriber: PhantomData<fn(S)>,
+#[derive(Default)]
+pub struct MetricsLayer {
+    with_labels:
+        Option<fn(&Dispatch, &Id, f: &mut dyn FnMut(&Labels) -> Option<Key>) -> Option<Key>>,
 }
 
-impl<S> MetricsLayer<S>
-where
-    S: Subscriber + for<'span> LookupSpan<'span>,
-{
-    /// Create a new `MetricsLayer`.
+impl MetricsLayer {
+    /// Create a new [`MetricsLayer`].
     pub fn new() -> Self {
-        let ctx = WithContext { with_labels: Self::with_labels };
-
-        Self { ctx, _subscriber: PhantomData }
+        Self::default()
     }
 
-    fn with_labels(
+    pub(crate) fn with_labels(
+        &self,
         dispatch: &Dispatch,
         id: &Id,
-        f: &mut dyn FnMut(&Labels) -> Option<Key>,
+        f: &mut dyn FnMut(Map) -> Option<Key>,
     ) -> Option<Key> {
-        let subscriber = dispatch
-            .downcast_ref::<S>()
-            .expect("subscriber should downcast to expected type; this is a bug!");
-        let span = subscriber.span(id).expect("registry should have a span for the current ID");
-
-        let ext = span.extensions();
-        f(ext.get::<Labels>()?)
+        let mut ff = |labels: &Labels| f(labels.0.clone());
+        (self.with_labels?)(dispatch, id, &mut ff)
     }
 }
 
-impl<S> Layer<S> for MetricsLayer<S>
+impl<S> Layer<S> for MetricsLayer
 where
     S: Subscriber + for<'a> LookupSpan<'a>,
 {
+    fn on_layer(&mut self, _: &mut S) {
+        self.with_labels = Some(
+            |dispatch: &Dispatch, id: &span::Id, f: &mut dyn FnMut(&Labels) -> Option<Key>| {
+                let subscriber = dispatch
+                    .downcast_ref::<S>()
+                    .expect("subscriber should downcast to expected type; this is a bug!");
+                let span =
+                    subscriber.span(id).expect("registry should have a span for the current ID");
+
+                let ext = span.extensions();
+                f(ext.get::<Labels>()?)
+            },
+        );
+    }
+
     fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, cx: Context<'_, S>) {
         let span = cx.span(id).expect("span must already exist!");
         let mut labels = Labels::from_record(&Record::new(attrs.values()));
 
         if let Some(parent) = span.parent() {
             if let Some(parent_labels) = parent.extensions().get::<Labels>() {
-                labels.extend_from_labels(parent_labels);
+                labels.extend_from_labels(parent_labels, false);
             }
         }
 
@@ -153,26 +155,9 @@ where
 
         let ext = &mut span.extensions_mut();
         if let Some(existing) = ext.get_mut::<Labels>() {
-            existing.extend_from_labels(&labels);
+            existing.extend_from_labels(&labels, true);
         } else {
             ext.insert(labels);
         }
     }
-
-    unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
-        match id {
-            id if id == TypeId::of::<Self>() => Some(self as *const _ as *const ()),
-            id if id == TypeId::of::<WithContext>() => Some(&self.ctx as *const _ as *const ()),
-            _ => None,
-        }
-    }
-}
-
-impl<S> Default for MetricsLayer<S>
-where
-    S: Subscriber + for<'span> LookupSpan<'span>,
-{
-    fn default() -> Self {
-        MetricsLayer::new()
-    }
 }