Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Cole/capture span fields #5

Merged
merged 2 commits into from
Oct 9, 2023
Merged

Cole/capture span fields #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
996ef03
fix: try to capture span fields
cmackenzie1 Oct 8, 2023
719405d
refactor: add visitor pattern for tracing fields
cmackenzie1 Oct 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
110 changes: 72 additions & 38 deletions src/formatter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
}
}
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -118,47 +123,74 @@ 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)?;

let span = event
.parent()
.and_then(|id| ctx.span(id))
.or_else(|| ctx.lookup_current());
// Add the new map to the top-level map
let obj: Option<serde_json::Value> = 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) = span {
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::<FormattedFields<N>>()
.expect("Unable to find FormattedFields in extensions; this is a bug");

if !data.is_empty() {
let obj: Option<serde_json::Value> = 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::Value> = 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)?;
}
}
}
}
}
}

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(())
}
Expand All @@ -183,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(())
}
Expand Down
95 changes: 70 additions & 25 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
//!
//! Licensed under MIT license [LICENSE](./LICENSE)
mod formatter;
mod visitor;

use tracing_core::Subscriber;
use tracing_subscriber::fmt::{Layer, SubscriberBuilder};
Expand Down Expand Up @@ -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<Error> for std::fmt::Error {
Expand Down Expand Up @@ -208,8 +207,8 @@ impl Builder {
S: Subscriber + for<'a> LookupSpan<'a>,
{
tracing_subscriber::fmt::layer()
.event_format(self.events)
.fmt_fields(self.fields)
.event_format(self.events)
}

pub fn subscriber_builder(
Expand All @@ -234,42 +233,57 @@ where
mod tests {

use super::*;
use tracing::{debug, error, info, info_span, trace, warn};
use tracing_core::Level;

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 formatter = formatter::JsonEventFormatter::new();
let subscriber = tracing_subscriber::fmt()
.fmt_fields(formatter::FieldsFormatter::new())
.event_format(formatter)
.with_max_level(Level::TRACE)
.finish();
let subscriber = tracing_subscriber::registry().with(builder().layer());

tracing::subscriber::with_default(subscriber, || {
trace!(a = "b", "hello world from trace");
debug!("hello world from debug");
info!("hello world from info");
warn!("hello world from warn");
error!("hello world from error");

let span = info_span!("test_span", b = "b", d = "d", later = tracing::field::Empty,);
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",
);
});
});
});

let formatter = formatter::JsonEventFormatter::new()
.with_level_name("severity")
.with_message_name("msg")
.with_timestamp_name("ts")
.with_timestamp_format(TimestampFormat::Unix);
let subscriber = tracing_subscriber::registry().with(
builder()
.with_level_name("severity")
.with_message_name("msg")
.with_timestamp_name("ts")
.with_timestamp_format(TimestampFormat::Unix)
.with_flatten_fields(false)
.layer(),
);

let subscriber = tracing_subscriber::fmt()
.fmt_fields(formatter::FieldsFormatter::new())
.event_format(formatter)
.with_max_level(Level::TRACE)
.finish();
tracing::subscriber::with_default(subscriber, || {
trace!(a = "b", "hello world from trace");
debug!("hello world from debug");
Expand All @@ -284,11 +298,42 @@ mod tests {
);
span.in_scope(|| {
info!("some message from inside a info_span");
let inner = info_span!("inner_span", a = "b", c = "d");
let inner = info_span!("inner_span", a = "b", c = "d", inner_span = true);
inner.in_scope(|| {
info!("some message from inside a info_span");
info!(
inner_span_field = true,
later = "populated from inside a span",
"some message from inside a info_span",
);
});
});
});
}

#[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);
});
}
}
Loading