1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

use serde::ser::{Serialize, SerializeMap, SerializeSeq};
use serde_json::Value;
use tracing_core::Subscriber;
use tracing_subscriber::{
    fmt::{format::JsonFields, FmtContext, FormattedFields},
    registry::{LookupSpan, SpanRef},
};

/// Serializable tracing span for nesting formatted event fields
pub(crate) struct SerializableSpan<'a, 'b, S>(&'b SpanRef<'a, S>)
where
    S: for<'lookup> LookupSpan<'lookup>;

impl<'a, 'b, S> SerializableSpan<'a, 'b, S>
where
    S: for<'lookup> LookupSpan<'lookup>,
{
    pub(crate) fn new(span: &'b SpanRef<'a, S>) -> Self {
        Self(span)
    }
}

impl<'a, 'b, S> Serialize for SerializableSpan<'a, 'b, S>
where
    S: for<'lookup> LookupSpan<'lookup>,
{
    fn serialize<R>(&self, serializer: R) -> Result<R::Ok, R::Error>
    where
        R: serde::Serializer,
    {
        let name = self.0.name();
        let extensions = self.0.extensions();

        let formatted_fields = extensions
            .get::<FormattedFields<JsonFields>>()
            .expect("No fields!");

        let span_length = formatted_fields.fields.len() + 1;
        let mut map = serializer.serialize_map(Some(span_length))?;

        match serde_json::from_str::<Value>(formatted_fields) {
            // handle string escaping "properly" (this should be fixed upstream)
            // https://github.com/tokio-rs/tracing/issues/391
            Ok(Value::Object(fields)) => {
                for (key, value) in fields {
                    map.serialize_entry(&key, &value)?;
                }
            }
            // these two options should be impossible
            Ok(value) => panic!("Invalid value: {}", value),
            Err(error) => panic!("Error parsing logs: {}", error),
        };

        map.serialize_entry("name", &name)?;
        map.end()
    }
}

/// Serializable tracing context for serializing a collection of spans
pub(crate) struct SerializableContext<'a, 'b, S>(&'b FmtContext<'a, S, JsonFields>)
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>;

impl<'a, 'b, S> SerializableContext<'a, 'b, S>
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    pub(crate) fn new(context: &'b FmtContext<'a, S, JsonFields>) -> Self {
        Self(context)
    }
}

impl<'a, 'b, S> Serialize for SerializableContext<'a, 'b, S>
where
    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
{
    fn serialize<R>(&self, serializer: R) -> Result<R::Ok, R::Error>
    where
        R: serde::Serializer,
    {
        let mut list = serializer.serialize_seq(None)?;

        if let Some(leaf_span) = self.0.lookup_current() {
            for span in leaf_span.scope().from_root() {
                list.serialize_element(&SerializableSpan::new(&span))?;
            }
        }

        list.end()
    }
}