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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

//! Extensions for subscribers
//!
//! This module provides utilities for subscribers

use std::{collections::HashMap, time::Instant};

use tracing_subscriber::registry::SpanRef;

mod pretty;

pub use pretty::*;

#[cfg(test)]
mod tests;

/// Trait for a span extension
pub trait SpanExtension {
    /// Registers an extension with default values
    fn register_default<S>(span_ref: &SpanRef<S>)
    where
        S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
        Self: Default + Send + Sync + 'static,
    {
        let mut extensions = span_ref.extensions_mut();
        if extensions.get_mut::<Self>().is_none() {
            let ext = Self::default();
            extensions.insert(ext);
        }
    }

    /// Registers an extension with initial value
    fn register_value<S>(initial_value: Self, span_ref: &SpanRef<S>)
    where
        S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
        Self: Send + Sync + Sized + 'static,
    {
        let mut extensions = span_ref.extensions_mut();
        if extensions.get_mut::<Self>().is_none() {
            extensions.insert(initial_value);
        }
    }

    /// Records the span attributes for the extension
    fn record_attrs<S>(span_ref: &SpanRef<S>, attrs: &tracing::span::Attributes<'_>)
    where
        S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
        Self: tracing::field::Visit + Sized + 'static,
    {
        let mut extensions = span_ref.extensions_mut();
        let ext = extensions
            .get_mut::<Self>()
            .expect("Extension not initialized");
        attrs.record(ext);
    }
}

/// A span extension to record the span attributes
#[derive(Debug, Default)]
pub struct SpanExtAttrs {
    /// Attributes values
    attrs: HashMap<&'static str, String>,
}

impl SpanExtension for SpanExtAttrs {}

impl tracing::field::Visit for SpanExtAttrs {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        let value = format!("{value:?}");
        self.attrs.insert(field.name(), value);
    }
}

/// A span extensison to record timing info
#[derive(Debug)]
pub struct SpanExtTiming {
    /// Instant when the span was entered
    pub entered: Instant,
}

impl Default for SpanExtTiming {
    fn default() -> Self {
        Self {
            entered: Instant::now(),
        }
    }
}

impl SpanExtension for SpanExtTiming {}

/// A visitor for events
///
/// The visitor saves the event data
#[derive(Debug, Default)]
pub struct EventVisitor {
    /// Fields
    fields: HashMap<&'static str, String>,
}

impl tracing::field::Visit for EventVisitor {
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        let value_str = format!("{value:?}");
        self.fields.insert(field.name(), value_str);
    }
}

impl EventVisitor {
    /// Records an event fields
    ///
    /// Returns the event message and the event fields
    pub fn record_event(event: &tracing::Event) -> Self {
        let mut f_visitor = EventVisitor::default();
        event.record(&mut f_visitor);
        f_visitor
    }

    /// Returns the event message
    pub fn message(&self) -> &str {
        match self.fields.get("message") {
            Some(s) => s,
            None => {
                panic!("Event message not found")
            }
        }
    }

    /// Returns the event fields (exc. message)
    pub fn meta_fields(&self) -> HashMap<&'static str, &str> {
        self.fields
            .iter()
            .filter_map(|(k, v)| {
                if *k == "message" {
                    return None;
                }
                Some((*k, v.as_str()))
            })
            .collect()
    }
}