From aa174a5ea05eb49d3515c5c05adfdd81ee91c983 Mon Sep 17 00:00:00 2001 From: Quentin Gliech Date: Wed, 16 Apr 2025 16:06:06 +0200 Subject: [PATCH] Roll our own event formatter --- Cargo.lock | 5 ++ Cargo.toml | 4 ++ crates/cli/src/main.rs | 3 +- crates/context/Cargo.toml | 5 ++ crates/context/src/fmt.rs | 141 ++++++++++++++++++++++++++++++++++++++ crates/context/src/lib.rs | 2 + 6 files changed, 158 insertions(+), 2 deletions(-) create mode 100644 crates/context/src/fmt.rs diff --git a/Cargo.lock b/Cargo.lock index 96915d6f8..ee216fcbd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3251,11 +3251,16 @@ dependencies = [ name = "mas-context" version = "0.15.0-rc.0" dependencies = [ + "console", + "opentelemetry", "pin-project-lite", "quanta", "tokio", "tower-layer", "tower-service", + "tracing", + "tracing-opentelemetry", + "tracing-subscriber", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index b6fe86947..8718e70f1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -112,6 +112,10 @@ version = "1.1.9" [workspace.dependencies.compact_str] version = "0.9.0" +# Terminal formatting +[workspace.dependencies.console] +version = "0.15.11" + # Time utilities [workspace.dependencies.chrono] version = "0.4.40" diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 85c5a89f1..1477b634a 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -91,8 +91,7 @@ async fn try_main() -> anyhow::Result { let (log_writer, _guard) = tracing_appender::non_blocking(output); let fmt_layer = tracing_subscriber::fmt::layer() .with_writer(log_writer) - .with_file(true) - .with_line_number(true) + .event_format(mas_context::EventFormatter) .with_ansi(with_ansi); let filter_layer = EnvFilter::try_from_default_env() .or_else(|_| EnvFilter::try_new("info")) diff --git a/crates/context/Cargo.toml b/crates/context/Cargo.toml index b0f422b51..762985080 100644 --- a/crates/context/Cargo.toml +++ b/crates/context/Cargo.toml @@ -12,8 +12,13 @@ publish = false workspace = true [dependencies] +console.workspace = true pin-project-lite.workspace = true quanta.workspace = true tokio.workspace = true tower-service.workspace = true tower-layer.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true +tracing-opentelemetry.workspace = true +opentelemetry.workspace = true diff --git a/crates/context/src/fmt.rs b/crates/context/src/fmt.rs new file mode 100644 index 000000000..345025c10 --- /dev/null +++ b/crates/context/src/fmt.rs @@ -0,0 +1,141 @@ +// Copyright 2025 New Vector Ltd. +// +// SPDX-License-Identifier: AGPL-3.0-only +// Please see LICENSE in the repository root for full details. + +use console::{Color, Style}; +use opentelemetry::{TraceId, trace::TraceContextExt}; +use tracing::{Level, Subscriber}; +use tracing_opentelemetry::OtelData; +use tracing_subscriber::{ + fmt::{ + FormatEvent, FormatFields, + format::{DefaultFields, Writer}, + time::{FormatTime, SystemTime}, + }, + registry::LookupSpan, +}; + +use crate::LogContext; + +/// An event formatter usable by the [`tracing-subscriber`] crate, which +/// includes the log context and the OTEL trace ID. +#[derive(Debug, Default)] +pub struct EventFormatter; + +struct FmtLevel<'a> { + level: &'a Level, + ansi: bool, +} + +impl<'a> FmtLevel<'a> { + pub(crate) fn new(level: &'a Level, ansi: bool) -> Self { + Self { level, ansi } + } +} + +const TRACE_STR: &str = "TRACE"; +const DEBUG_STR: &str = "DEBUG"; +const INFO_STR: &str = " INFO"; +const WARN_STR: &str = " WARN"; +const ERROR_STR: &str = "ERROR"; + +const TRACE_STYLE: Style = Style::new().fg(Color::Magenta); +const DEBUG_STYLE: Style = Style::new().fg(Color::Blue); +const INFO_STYLE: Style = Style::new().fg(Color::Green); +const WARN_STYLE: Style = Style::new().fg(Color::Yellow); +const ERROR_STYLE: Style = Style::new().fg(Color::Red); + +impl std::fmt::Display for FmtLevel<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let msg = match *self.level { + Level::TRACE => TRACE_STYLE.force_styling(self.ansi).apply_to(TRACE_STR), + Level::DEBUG => DEBUG_STYLE.force_styling(self.ansi).apply_to(DEBUG_STR), + Level::INFO => INFO_STYLE.force_styling(self.ansi).apply_to(INFO_STR), + Level::WARN => WARN_STYLE.force_styling(self.ansi).apply_to(WARN_STR), + Level::ERROR => ERROR_STYLE.force_styling(self.ansi).apply_to(ERROR_STR), + }; + write!(f, "{msg}") + } +} + +struct TargetFmt<'a> { + target: &'a str, + line: Option, +} + +impl<'a> TargetFmt<'a> { + pub(crate) fn new(metadata: &tracing::Metadata<'a>) -> Self { + Self { + target: metadata.target(), + line: metadata.line(), + } + } +} + +impl std::fmt::Display for TargetFmt<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.target)?; + if let Some(line) = self.line { + write!(f, ":{line}")?; + } + Ok(()) + } +} + +impl FormatEvent for EventFormatter +where + S: Subscriber + for<'a> LookupSpan<'a>, + N: for<'writer> FormatFields<'writer> + 'static, +{ + fn format_event( + &self, + ctx: &tracing_subscriber::fmt::FmtContext<'_, S, N>, + mut writer: Writer<'_>, + event: &tracing::Event<'_>, + ) -> std::fmt::Result { + let ansi = writer.has_ansi_escapes(); + let metadata = event.metadata(); + + SystemTime.format_time(&mut writer)?; + + let level = FmtLevel::new(metadata.level(), ansi); + let target = Style::new() + .dim() + .force_styling(ansi) + .apply_to(TargetFmt::new(metadata)); + write!(&mut writer, " {level} {target} ")?; + + if let Some(log_context) = LogContext::current() { + let log_context = Style::new() + .bold() + .force_styling(ansi) + .apply_to(log_context); + write!(&mut writer, "{log_context} - ")?; + } + + let field_fromatter = DefaultFields::new(); + field_fromatter.format_fields(writer.by_ref(), event)?; + + // If we have a OTEL span, we can add the trace ID to the end of the log line + if let Some(span) = ctx.lookup_current() { + if let Some(otel) = span.extensions().get::() { + // If it is the root span, the trace ID will be in the span builder. Else, it + // will be in the parent OTEL context + let trace_id = otel + .builder + .trace_id + .unwrap_or_else(|| otel.parent_cx.span().span_context().trace_id()); + if trace_id != TraceId::INVALID { + let label = Style::new() + .italic() + .force_styling(ansi) + .apply_to("trace.id"); + write!(&mut writer, " {label}={trace_id}")?; + } + } + } + + writeln!(&mut writer) + } +} diff --git a/crates/context/src/lib.rs b/crates/context/src/lib.rs index d189cb53e..b289bfcda 100644 --- a/crates/context/src/lib.rs +++ b/crates/context/src/lib.rs @@ -3,6 +3,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Please see LICENSE in the repository root for full details. +mod fmt; mod future; mod layer; mod service; @@ -20,6 +21,7 @@ use quanta::Instant; use tokio::task_local; pub use self::{ + fmt::EventFormatter, future::{LogContextFuture, PollRecordingFuture}, layer::LogContextLayer, service::LogContextService,