From ae8b401d3062acf8a6e12121c2bc704d7cd1f180 Mon Sep 17 00:00:00 2001 From: Cole MacKenzie Date: Sun, 8 Oct 2023 15:54:38 -0700 Subject: [PATCH] refactor: add visitor pattern for tracing fields By adding a visitor pattern, we can further customize the output of each event. For example, the `message` field was never actually being formatted with the provided `message_name` field. Also, this fixes an issue where the number of nested spans exceeds one. Fields from any previous spans were being ignored and not formatted in the final output. --- Cargo.toml | 1 - src/formatter.rs | 105 +++++++++++++++++++++++++++++++--------------- src/lib.rs | 43 +++++++++++++++++-- src/visitor.rs | 107 +++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 218 insertions(+), 38 deletions(-) create mode 100644 src/visitor.rs diff --git a/Cargo.toml b/Cargo.toml index f88a39f..eb775c0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,7 +16,6 @@ serde = "1.0.188" serde_json = "1.0.107" thiserror = "1.0.49" tracing-core = "0.1.31" -tracing-serde = "0.1.3" tracing-subscriber = "0.3.17" [dev-dependencies] diff --git a/src/formatter.rs b/src/formatter.rs index 22877f5..5dbd105 100644 --- a/src/formatter.rs +++ b/src/formatter.rs @@ -3,13 +3,19 @@ use std::fmt; use serde::{ser::SerializeMap, Serializer}; use tracing_core::{Event, Subscriber}; -use tracing_serde::fields::AsMap; use tracing_subscriber::{ fmt::{format, FmtContext, FormatEvent, FormatFields, FormattedFields}, registry::LookupSpan, }; -use crate::Error; +use crate::{visitor, Error}; + +const DEFAULT_TIMESTAMP_FORMAT: crate::TimestampFormat = crate::TimestampFormat::Rfc3339; +const DEFAULT_LEVEL_NAME: &str = "level"; +const DEFAULT_MESSAGE_NAME: &str = "message"; +const DEFAULT_TARGET_NAME: &str = "target"; +const DEFAULT_TIMESTAMP_NAME: &str = "timestamp"; +const DEFAULT_FLATTEN_FIELDS: bool = true; /// A JSON formatter for tracing events. /// This is used to format the event field in the JSON output. @@ -25,12 +31,12 @@ pub struct JsonEventFormatter { impl Default for JsonEventFormatter { fn default() -> Self { Self { - level_name: "level", - message_name: "message", - target_name: "target", - timestamp_name: "timestamp", - timestamp_format: crate::TimestampFormat::Rfc3339, - flatten_fields: true, + level_name: DEFAULT_LEVEL_NAME, + message_name: DEFAULT_MESSAGE_NAME, + target_name: DEFAULT_TARGET_NAME, + timestamp_name: DEFAULT_TIMESTAMP_NAME, + timestamp_format: DEFAULT_TIMESTAMP_FORMAT, + flatten_fields: DEFAULT_FLATTEN_FIELDS, } } } @@ -84,8 +90,7 @@ where ) -> fmt::Result { let now = chrono::Utc::now(); - let mut buffer = Vec::new(); - let mut binding = serde_json::Serializer::new(&mut buffer); + let mut binding = serde_json::Serializer::new(Vec::new()); let mut serializer = binding.serialize_map(None).map_err(Error::Serde)?; serializer @@ -118,33 +123,61 @@ where .serialize_entry(self.target_name, event.metadata().target()) .map_err(Error::Serde)?; + let msg_name = if self.message_name != DEFAULT_MESSAGE_NAME { + Some(self.message_name) + } else { + None + }; + if self.flatten_fields { - let mut visitor = tracing_serde::SerdeMapVisitor::new(serializer); + // record fields in the top-level map + let mut visitor = visitor::Visitor::new(&mut serializer, msg_name); event.record(&mut visitor); - - serializer = visitor.take_serializer().map_err(|_| Error::Unknown)?; + visitor.finish().map_err(Error::Serde)?; } else { - serializer - .serialize_entry("fields", &event.field_map()) - .map_err(Error::Serde)?; - }; + // record fields in a nested map under the key "fields" + let mut binding = serde_json::Serializer::pretty(Vec::new()); + let mut field_serializer = binding.serialize_map(None).map_err(Error::Serde)?; + let mut visitor = visitor::Visitor::new(&mut field_serializer, msg_name); + event.record(&mut visitor); + visitor.finish().map_err(Error::Serde)?; + field_serializer.end().map_err(Error::Serde)?; + + // Add the new map to the top-level map + let obj: Option = serde_json::from_str( + std::str::from_utf8(&binding.into_inner()).map_err(Error::Utf8)?, + ) + .ok(); + if matches!(obj, Some(serde_json::Value::Object(_))) { + let obj = obj.expect("matched object"); + serializer + .serialize_entry("fields", &obj) + .map_err(Error::Serde)?; + } + } // Write all fields from spans if let Some(leaf_span) = ctx.lookup_current() { for span in leaf_span.scope().from_root() { let ext = span.extensions(); - let data = ext + let formatted_fields = ext .get::>() .expect("Unable to find FormattedFields in extensions; this is a bug"); - if !data.is_empty() { - let obj: Option = serde_json::from_str(data.as_str()).ok(); - if matches!(obj, Some(serde_json::Value::Object(_))) { - let obj = obj.expect("matched object"); - for (key, value) in obj.as_object().unwrap() { - serializer - .serialize_entry(key, value) - .map_err(Error::Serde)?; + // formatted_fields actually contains multiple ndjson objects, one for every time a spans fields are formatted. + // re-parse these into JSON for serialization into the final map. Any fields redefined in a subsequent span + // will overwrite the previous value. + // TODO(cmackenzie1): There has to be a better way to do this. + for data in formatted_fields.split('\n') { + if !data.is_empty() { + let obj: Option = serde_json::from_str(data).ok(); + if matches!(obj, Some(serde_json::Value::Object(_))) { + let obj = obj.expect("matched object"); + for (key, value) in obj.as_object().unwrap() { + serializer + .serialize_entry(key, value) + .map_err(Error::Serde)?; + } } } } @@ -152,8 +185,12 @@ where } serializer.end().map_err(Error::Serde)?; - writer.write_str(std::str::from_utf8(&buffer).map_err(Error::Utf8)?)?; - writer.write_char('\n')?; + + writeln!( + writer, + "{}", + std::str::from_utf8(&binding.into_inner()).map_err(Error::Utf8)? + )?; Ok(()) } @@ -178,16 +215,18 @@ impl<'writer> FormatFields<'writer> for FieldsFormatter { where R: tracing_subscriber::field::RecordFields, { - let mut buffer = Vec::new(); - let mut binding = serde_json::Serializer::new(&mut buffer); + let mut binding = serde_json::Serializer::new(Vec::new()); let mut serializer = binding.serialize_map(None).map_err(Error::Serde)?; - let mut visitor = tracing_serde::SerdeMapVisitor::new(serializer); + let mut visitor = visitor::Visitor::new(&mut serializer, None); fields.record(&mut visitor); - serializer = visitor.take_serializer().map_err(|_| Error::Unknown)?; serializer.end().map_err(Error::Serde)?; - writer.write_str(std::str::from_utf8(&buffer).map_err(Error::Utf8)?)?; + writeln!( + writer, + "{}", + std::str::from_utf8(&binding.into_inner()).map_err(Error::Utf8)? + )?; Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index 2a7916d..4aff8b3 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -43,6 +43,7 @@ //! //! Licensed under MIT license [LICENSE](./LICENSE) mod formatter; +mod visitor; use tracing_core::Subscriber; use tracing_subscriber::fmt::{Layer, SubscriberBuilder}; @@ -97,8 +98,6 @@ enum Error { Serde(#[from] serde_json::Error), #[error("utf8 error: {0}")] Utf8(#[from] std::str::Utf8Error), - #[error("unknown error")] - Unknown, } impl From for std::fmt::Error { @@ -235,9 +234,17 @@ mod tests { use super::*; - use tracing::{debug, error, info, info_span, trace, warn}; + use tracing::{debug, error, info, info_span, instrument, trace, warn}; use tracing_subscriber::prelude::*; + #[instrument] + fn some_function(a: u32, b: u32) { + let span = info_span!("some_span", a = a, b = b); + span.in_scope(|| { + info!("some message from inside a span"); + }); + } + #[test] fn test_json_event_formatter() { let subscriber = tracing_subscriber::registry().with(builder().layer()); @@ -270,9 +277,10 @@ mod tests { let subscriber = tracing_subscriber::registry().with( builder() .with_level_name("severity") - .with_message_name("message") + .with_message_name("msg") .with_timestamp_name("ts") .with_timestamp_format(TimestampFormat::Unix) + .with_flatten_fields(false) .layer(), ); @@ -301,4 +309,31 @@ mod tests { }); }); } + + #[test] + fn test_nested_spans() { + let subscriber = tracing_subscriber::registry().with(builder().layer()); + + tracing::subscriber::with_default(subscriber, || { + let span = info_span!( + "test_span", + person.firstname = "cole", + person.lastname = "mackenzie", + later = tracing::field::Empty, + ); + span.in_scope(|| { + info!("some message from inside a info_span"); + let inner = info_span!("inner_span", a = "b", c = "d", inner_span = true); + inner.in_scope(|| { + info!( + inner_span_field = true, + later = "populated from inside a span", + "some message from inside a info_span", + ); + }); + }); + + some_function(1, 2); + }); + } } diff --git a/src/visitor.rs b/src/visitor.rs new file mode 100644 index 0000000..a10ce13 --- /dev/null +++ b/src/visitor.rs @@ -0,0 +1,107 @@ +use tracing_core::field::Visit; + +pub struct Visitor<'a, W> +where + W: serde::ser::SerializeMap, +{ + serializer: &'a mut W, + state: Result<(), W::Error>, + overwrite_message_name: Option<&'static str>, +} + +impl<'a, W> Visitor<'a, W> +where + W: serde::ser::SerializeMap, +{ + pub fn new(serializer: &'a mut W, overwrite_message_name: Option<&'static str>) -> Self { + Self { + serializer, + state: Ok(()), + overwrite_message_name, + } + } + + pub fn finish(self) -> Result<(), W::Error> { + self.state + } +} + +impl<'a, W> Visit for Visitor<'a, W> +where + W: serde::ser::SerializeMap, +{ + fn record_f64(&mut self, field: &tracing_core::Field, value: f64) { + if self.state.is_ok() { + self.state = self.serializer.serialize_entry(field.name(), &value); + } + } + + fn record_i64(&mut self, field: &tracing_core::Field, value: i64) { + if self.state.is_ok() { + self.state = self.serializer.serialize_entry(field.name(), &value); + } + } + + fn record_u64(&mut self, field: &tracing_core::Field, value: u64) { + if self.state.is_ok() { + self.state = self.serializer.serialize_entry(field.name(), &value); + } + } + + fn record_i128(&mut self, field: &tracing_core::Field, value: i128) { + if self.state.is_ok() { + self.state = self.serializer.serialize_entry(field.name(), &value); + } + } + + fn record_u128(&mut self, field: &tracing_core::Field, value: u128) { + if self.state.is_ok() { + self.state = self.serializer.serialize_entry(field.name(), &value); + } + } + + fn record_bool(&mut self, field: &tracing_core::Field, value: bool) { + if self.state.is_ok() { + self.state = self.serializer.serialize_entry(field.name(), &value); + } + } + + fn record_str(&mut self, field: &tracing_core::Field, value: &str) { + if self.state.is_ok() { + if self.overwrite_message_name.is_some() && field.name() == "message" { + self.state = self + .serializer + .serialize_entry(self.overwrite_message_name.expect("message"), &value); + } else { + self.state = self.serializer.serialize_entry(field.name(), &value); + } + } + } + + fn record_error( + &mut self, + field: &tracing_core::Field, + value: &(dyn std::error::Error + 'static), + ) { + if self.state.is_ok() { + self.state = self + .serializer + .serialize_entry(field.name(), &format_args!("{}", value).to_string()); + } + } + + fn record_debug(&mut self, field: &tracing_core::Field, value: &dyn std::fmt::Debug) { + if self.state.is_ok() { + if self.overwrite_message_name.is_some() && field.name() == "message" { + self.state = self.serializer.serialize_entry( + self.overwrite_message_name.expect("message"), + &format_args!("{:?}", value).to_string(), + ); + } else { + self.state = self + .serializer + .serialize_entry(field.name(), &format_args!("{:?}", value).to_string()); + } + } + } +}