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