Custom format fails to find span and trace ids only for new/close events.

[ShovelJockey]

I am using a custom formatter built from examples I have seen it works well except for new/close events which never seem to get the trace or span id recorded, not sure if the issue lies in how I derive context or perhaps set the parent on the span in the first place or possibly something else.

Any advice would be welcome.

creating layer

pub fn init_subscriber() -> Result<(), axum::BoxError> {
    let registery = tracing_subscriber::registry()
        .with(build_otlp_layer()?)
        .with(build_loglevel_filter_layer())
        .with(build_logger_text());
    info!("started otlp logging & tracing");
    tracing::subscriber::set_global_default(registery)?;
    Ok(())
}

pub fn build_logger_text<S>() -> Box<dyn Layer<S> + Send + Sync + 'static>
where
    S: Subscriber + for<'a> LookupSpan<'a>,
{
    use tracing_subscriber::fmt::format::FmtSpan;
    Box::new(
        tracing_subscriber::fmt::layer()
            .json()
            .with_current_span(true)
            .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
            .event_format(TraceIdFormat),
    )
}

format

use chrono::Utc;
use opentelemetry::trace::{SpanId, TraceId};
use serde::ser::{Serialize, SerializeMap, Serializer};
use std::{io, marker::PhantomData};
use tracing::{Event, Subscriber};
use tracing_serde::{fields::AsMap, AsSerde};
use tracing_subscriber::{
    fmt::{format::Writer, FmtContext, FormatEvent, FormatFields, FormattedFields},
    registry::{LookupSpan, SpanRef},
};

use crate::tracing_funcs::utils::trace_info_from_ref;

pub struct TraceIdFormat;

impl<S, N> FormatEvent<S, N> for TraceIdFormat
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
    N: for<'writer> FormatFields<'writer> + 'static,
{
    fn format_event(
        &self,
        ctx: &FmtContext<'_, S, N>,
        mut writer: Writer<'_>,
        event: &Event<'_>,
    ) -> std::fmt::Result
    where
        S: Subscriber + for<'a> LookupSpan<'a>,
    {
        use opentelemetry::trace::TraceContextExt;
        use tracing_opentelemetry::OpenTelemetrySpanExt;
        let meta = event.metadata();

        let mut visit = || {
            let mut serializer = serde_json::Serializer::new(WriteAdaptor::new(&mut writer));
            let mut serializer = serializer.serialize_map(None)?;
            serializer.serialize_entry("timestamp", &Utc::now().to_rfc3339())?;
            serializer.serialize_entry("level", &meta.level().as_serde())?;
            serializer.serialize_entry("fields", &event.field_map())?;
            serializer.serialize_entry("target", "otel::tracing")?;

            let context = tracing::Span::current().context();
            let trace_info = trace_info_from_ref(context.span());

            if trace_info.trace_id != TraceId::INVALID {
                serializer.serialize_entry("trace_id", &trace_info.trace_id.to_string())?;
            }

            if trace_info.span_id != SpanId::INVALID {
                serializer.serialize_entry("span_id", &trace_info.span_id.to_string())?;
            }

            let format_field_marker: PhantomData<N> = PhantomData;
            if let Some(span_ref) = event
                .parent()
                .and_then(|id| ctx.span(id))
                .or_else(|| ctx.lookup_current()) {
                    serializer
                        .serialize_entry("span", &SerializableSpan(&span_ref, format_field_marker))
                        .unwrap_or(());
                }

            serializer.end()
        };

        visit().map_err(|_| std::fmt::Error)?;
        writeln!(writer)
    }
}

pub struct WriteAdaptor<'a> {
    fmt_write: &'a mut dyn std::fmt::Write,
}

impl<'a> WriteAdaptor<'a> {
    pub fn new(fmt_write: &'a mut dyn std::fmt::Write) -> Self {
        Self { fmt_write }
    }
}

impl<'a> io::Write for WriteAdaptor<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let s =
            std::str::from_utf8(buf).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        self.fmt_write
            .write_str(s)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;

        Ok(s.as_bytes().len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

struct SerializableSpan<'a, 'b, Span, N>(&'b SpanRef<'a, Span>, PhantomData<N>)
where
    Span: for<'lookup> LookupSpan<'lookup>,
    N: for<'writer> FormatFields<'writer> + 'static;

impl<'a, 'b, Span, N> Serialize for SerializableSpan<'a, 'b, Span, N>
where
    Span: for<'lookup> LookupSpan<'lookup>,
    N: for<'writer> FormatFields<'writer> + 'static,
{
    fn serialize<Ser>(&self, serializer: Ser) -> Result<Ser::Ok, Ser::Error>
    where
        Ser: serde::ser::Serializer,
    {
        let mut serializer = serializer.serialize_map(None)?;

        let ext = self.0.extensions();
        let data = ext
            .get::<FormattedFields<N>>()
            .expect("Unable to find FormattedFields in extensions; this is a bug");

        match serde_json::from_str::<serde_json::Value>(data) {
            Ok(serde_json::Value::Object(fields)) => {
                for field in fields {
                    serializer.serialize_entry(&field.0, &field.1)?;
                }
            }
            Ok(_) if cfg!(debug_assertions) => panic!(
                "span '{}' had malformed fields! this is a bug.\n  error: invalid JSON object\n  fields: {:?}",
                self.0.metadata().name(),
                data
            ),
            Ok(value) => {
                serializer.serialize_entry("field", &value)?;
                serializer.serialize_entry("field_error", "field was no a valid object")?
            }
            Err(e) if cfg!(debug_assertions) => panic!(
                "span '{}' had malformed fields! this is a bug.\n  error: {}\n  fields: {:?}",
                self.0.metadata().name(),
                e,
                data
            ),
            Err(e) => serializer.serialize_entry("field_error", &format!("{e}"))?,
        };
        serializer.serialize_entry("name", self.0.metadata().name())?;
        serializer.end()
    }
}

log examples

{"timestamp":"2024-04-08T17:20:03.336313220+00:00","level":"INFO","fields":{"message":"new"},"target":"otel::tracing","span"{"http.request.method":"GET","network.protocol.version":"1.1","otel.kind":"Server","otel.name":"GET","request_id":"f05a0c70-a233-416e-b67c-806b28a8c551","server.address":"localhost:8000","url.path":"/status"}}

{"timestamp":"2024-04-08T17:20:03.336355288+00:00","level":"INFO","fields":{"message":"Getting status"},"target":"otel::tracing","trace_id":"6f412abc3343d52d7f5f5279d78a947b","span_id":"d726fe0c3ed40295","span":{"http.request.method":"GET","http.route":"/status","network.protocol.version":"1.1","otel.kind":"Server","otel.name":"GET /prevent/v1.0/status","request_id":"f05a0c70-a233-416e-b67c-806b28a8c551","server.address":"localhost:8000","url.path":"/prevent/v1.0/status","url.scheme":"","user_agent.original":"curl/7.68.0","name":"HTTP request"}}

{"timestamp":"2024-04-08T17:20:03.361217853+00:00","level":"INFO","fields":{"message":"Unable to process"},"target":"otel::tracing","trace_id":"6f412abc3343d52d7f5f5279d78a947b","span_id":"d726fe0c3ed40295","span"{"http.request.method":"GET","http.route":"/status","network.protocol.version":"1.1","otel.kind":"Server","otel.name":"GET /prevent/v1.0/status","request_id":"f05a0c70-a233-416e-b67c-806b28a8c551","server.address":"localhost:8000","url.path":"/status","url.scheme":"","user_agent.original":"curl/7.68.0","name":"HTTP request"}}

{"timestamp":"2024-04-08T17:20:03.438278707+00:00","level":"INFO","fields":{"message":"close","time.busy":"83.6ms","time.idle":"18.4ms"},"target":"otel::tracing","span":{"http.request.method":"GET","http.response.status_code":200,"http.route":"/status","network.protocol.version":"1.1","otel.kind":"Server","otel.name":"GET /status","request_id":"f05a0c70-a233-416e-b67c-806b28a8c551","server.address":"localhost:8000","url.path":"/status","url.scheme":"","user_agent.original":"curl/7.68.0","name":"HTTP request"}}