feat[tracing-subscriber]: extra fields

tokio-rs · Oct 2, 2023 · 1157692 · 1157692
1 parent 62d57a6
commit 1157692
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 0 deletions.
diff --git a/tracing-subscriber/Cargo.toml b/tracing-subscriber/Cargo.toml
@@ -36,6 +36,7 @@ json = ["tracing-serde", "serde", "serde_json"]
 # Enables support for local time when using the `time` crate timestamp
 # formatters.
 local-time = ["time/local-offset"]
+extra-fields = []
 
 [dependencies]
 tracing-core = { path = "../tracing-core", version = "0.2", default-features = false }

diff --git a/tracing-subscriber/src/fmt/fmt_subscriber.rs b/tracing-subscriber/src/fmt/fmt_subscriber.rs
@@ -590,6 +590,15 @@ where
             _inner: self._inner,
         }
     }
+
+    /// Set the function to make extra fields.
+    #[cfg(feature = "extra-fields")]
+    pub fn with_extra_fields(self, make_extra_fields: Box<dyn format::MakeExtraFields>) -> Self {
+        Subscriber {
+            fmt_event: self.fmt_event.with_make_extra_fields(make_extra_fields),
+            ..self
+        }
+    }
 }
 
 #[cfg(feature = "json")]

diff --git a/tracing-subscriber/src/fmt/format/json.rs b/tracing-subscriber/src/fmt/format/json.rs
@@ -304,6 +304,13 @@ where
                     .serialize_entry("threadId", &format!("{:?}", std::thread::current().id()))?;
             }
 
+            #[cfg(feature = "extra-fields")]
+            if let Some((make_extra_fields, current_span)) = self.make_extra_fields.as_ref().zip(current_span) {
+                for (k, v) in make_extra_fields(current_span.extensions()) {
+                    serializer.serialize_entry(&k, &v)?;
+                }
+            }
+
             serializer.end()
         };
 

diff --git a/tracing-subscriber/src/fmt/format/mod.rs b/tracing-subscriber/src/fmt/format/mod.rs
@@ -60,6 +60,11 @@ mod pretty;
 #[cfg_attr(docsrs, doc(cfg(feature = "ansi")))]
 pub use pretty::*;
 
+#[cfg(feature = "extra-fields")]
+use crate::registry::Extensions;
+#[cfg(feature = "extra-fields")]
+use std::collections::HashMap;
+
 use fmt::{Debug, Display};
 
 /// A type that can format a tracing [`Event`] to a [`Writer`].
@@ -413,6 +418,43 @@ pub struct Format<F = Full, T = SystemTime> {
     pub(crate) display_thread_name: bool,
     pub(crate) display_filename: bool,
     pub(crate) display_line_number: bool,
+    #[cfg(feature = "extra-fields")]
+    pub(crate) make_extra_fields: Option<Box<dyn MakeExtraFields>>,
+}
+
+/// Make extra fields.
+#[cfg(feature = "extra-fields")]
+pub trait MakeExtraFields: Fn(Extensions<'_>) -> HashMap<String, String> + Send + Sync {
+    /// Clone self into a boxed [MakeTraceId].
+    fn clone_box<'a>(&self) -> Box<dyn 'a + MakeExtraFields>
+    where
+        Self: 'a;
+}
+
+#[cfg(feature = "extra-fields")]
+impl Debug for Box<dyn MakeExtraFields> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "MakeExtraFields")
+    }
+}
+
+#[cfg(feature = "extra-fields")]
+impl<F> MakeExtraFields for F
+    where
+        F: Fn(Extensions<'_>) -> HashMap<String, String> + Clone + Send + Sync {
+    fn clone_box<'a>(&self) -> Box<dyn 'a + MakeExtraFields>
+    where
+        Self: 'a,
+    {
+        Box::new(self.clone())
+    }
+}
+
+#[cfg(feature = "extra-fields")]
+impl<'a> Clone for Box<dyn 'a + MakeExtraFields> {
+    fn clone(&self) -> Self {
+        self.clone_box()
+    }
 }
 
 // === impl Writer ===
@@ -591,6 +633,8 @@ impl Default for Format<Full, SystemTime> {
             display_thread_name: false,
             display_filename: false,
             display_line_number: false,
+            #[cfg(feature = "extra-fields")]
+            make_extra_fields: None,
         }
     }
 }
@@ -611,6 +655,8 @@ impl<F, T> Format<F, T> {
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
             display_line_number: self.display_line_number,
+            #[cfg(feature = "extra-fields")]
+            make_extra_fields: None,
         }
     }
 
@@ -650,6 +696,8 @@ impl<F, T> Format<F, T> {
             display_thread_name: self.display_thread_name,
             display_filename: true,
             display_line_number: true,
+            #[cfg(feature = "extra-fields")]
+            make_extra_fields: None,
         }
     }
 
@@ -682,6 +730,8 @@ impl<F, T> Format<F, T> {
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
             display_line_number: self.display_line_number,
+            #[cfg(feature = "extra-fields")]
+            make_extra_fields: None,
         }
     }
 
@@ -711,6 +761,8 @@ impl<F, T> Format<F, T> {
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
             display_line_number: self.display_line_number,
+            #[cfg(feature = "extra-fields")]
+            make_extra_fields: None,
         }
     }
 
@@ -727,6 +779,8 @@ impl<F, T> Format<F, T> {
             display_thread_name: self.display_thread_name,
             display_filename: self.display_filename,
             display_line_number: self.display_line_number,
+            #[cfg(feature = "extra-fields")]
+            make_extra_fields: None,
         }
     }
 
@@ -829,6 +883,15 @@ impl<F, T> Format<F, T> {
         Ok(())
     }
 
+    /// Set the function to make extra fields.
+    #[cfg(feature = "extra-fields")]
+    pub fn with_make_extra_fields(self, make_trace_id: Box<dyn MakeExtraFields>) -> Self {
+        Format {
+            make_extra_fields: Some(make_trace_id),
+            ..self
+        }
+    }
+
     #[inline]
     fn format_timestamp(&self, writer: &mut Writer<'_>) -> fmt::Result
     where