opentelemetry_tracing_utils/
trace_output_fmt.rs

1use opentelemetry::trace::TraceContextExt;
2use serde::ser::{SerializeMap, Serializer as _};
3use std::io;
4use tracing::{Event, Subscriber};
5use tracing_serde::AsSerde;
6use tracing_subscriber::fmt::format::Writer;
7use tracing_subscriber::fmt::{FmtContext, FormatEvent, FormatFields};
8use tracing_subscriber::registry::LookupSpan;
9
10pub struct WriteAdaptor<'a> {
11    fmt_write: &'a mut dyn std::fmt::Write,
12}
13
14impl<'a> WriteAdaptor<'a> {
15    pub fn new(fmt_write: &'a mut dyn std::fmt::Write) -> Self {
16        Self { fmt_write }
17    }
18}
19
20impl<'a> io::Write for WriteAdaptor<'a> {
21    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
22        let s =
23            std::str::from_utf8(buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
24
25        self.fmt_write
26            .write_str(s)
27            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
28
29        Ok(s.as_bytes().len())
30    }
31
32    fn flush(&mut self) -> io::Result<()> {
33        Ok(())
34    }
35}
36
37/// Derived from https://github.com/tokio-rs/tracing/issues/1531#issuecomment-1136971089 combined
38/// with default Json formatter
39pub struct JsonWithTraceId;
40
41pub struct TraceInfo {
42    pub trace_id: String,
43    pub span_id: String,
44}
45
46pub fn lookup_trace_info<S>(
47    span_ref: &tracing_subscriber::registry::SpanRef<S>,
48) -> Option<TraceInfo>
49where
50    S: Subscriber + for<'a> LookupSpan<'a>,
51{
52    span_ref
53        .extensions()
54        .get::<tracing_opentelemetry::OtelData>()
55        .map(|o| {
56            TraceInfo {
57                // commented out line was from the original, conversation here:
58                // https://github.com/tokio-rs/tracing/issues/1531#issuecomment-1137296115
59                // trace_id: o.parent_cx.span().span_context().trace_id().to_string(),
60                trace_id: o
61                    .builder
62                    .trace_id
63                    .unwrap_or(o.parent_cx.span().span_context().trace_id())
64                    .to_string(),
65                span_id: o
66                    .builder
67                    .span_id
68                    .unwrap_or(opentelemetry::trace::SpanId::INVALID)
69                    .to_string(),
70            }
71        })
72}
73
74impl<S, N> FormatEvent<S, N> for JsonWithTraceId
75where
76    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
77    N: for<'writer> FormatFields<'writer> + 'static,
78{
79    fn format_event(
80        &self,
81        ctx: &FmtContext<'_, S, N>,
82        mut writer: Writer<'_>,
83        event: &Event<'_>,
84    ) -> std::fmt::Result
85    where
86        S: Subscriber + for<'a> LookupSpan<'a>,
87    {
88        let meta = event.metadata();
89
90        let mut visit = || {
91            let mut serializer = serde_json::Serializer::new(WriteAdaptor::new(&mut writer));
92
93            let mut serializer = serializer.serialize_map(None)?;
94            serializer.serialize_entry("level", &meta.level().as_serde())?;
95
96            let _format_field_marker: std::marker::PhantomData<N> = std::marker::PhantomData;
97
98            use tracing_serde::fields::AsMap;
99            serializer.serialize_entry("fields", &event.field_map())?;
100
101            serializer.serialize_entry("target", meta.target())?;
102
103            if let Some(ref span_ref) = ctx.lookup_current() {
104                if let Some(trace_info) = lookup_trace_info(span_ref) {
105                    serializer.serialize_entry("span_id", &trace_info.span_id)?;
106                    serializer.serialize_entry("trace_id", &trace_info.trace_id)?;
107                }
108            }
109
110            serializer.end()
111        };
112
113        visit().map_err(|_| std::fmt::Error)?;
114        writeln!(writer)
115    }
116}