tracing_ext/sub/
mod.rs

1//! Extensions for subscribers
2//!
3//! This module provides utilities for subscribers
4
5use std::{collections::HashMap, time::Instant};
6
7use tracing_subscriber::registry::SpanRef;
8
9mod pretty;
10
11pub use pretty::*;
12
13#[cfg(test)]
14mod tests;
15
16/// Trait for a span extension
17pub trait SpanExtension {
18    /// Registers an extension with default values
19    fn register_default<S>(span_ref: &SpanRef<S>)
20    where
21        S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
22        Self: Default + Send + Sync + 'static,
23    {
24        let mut extensions = span_ref.extensions_mut();
25        if extensions.get_mut::<Self>().is_none() {
26            let ext = Self::default();
27            extensions.insert(ext);
28        }
29    }
30
31    /// Registers an extension with initial value
32    fn register_value<S>(initial_value: Self, span_ref: &SpanRef<S>)
33    where
34        S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
35        Self: Send + Sync + Sized + 'static,
36    {
37        let mut extensions = span_ref.extensions_mut();
38        if extensions.get_mut::<Self>().is_none() {
39            extensions.insert(initial_value);
40        }
41    }
42
43    /// Records the span attributes for the extension
44    fn record_attrs<S>(span_ref: &SpanRef<S>, attrs: &tracing::span::Attributes<'_>)
45    where
46        S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
47        Self: tracing::field::Visit + Sized + 'static,
48    {
49        let mut extensions = span_ref.extensions_mut();
50        let ext = extensions
51            .get_mut::<Self>()
52            .expect("Extension not initialized");
53        attrs.record(ext);
54    }
55}
56
57/// A span extension to record the span attributes
58#[derive(Debug, Default)]
59pub struct SpanExtAttrs {
60    /// Attributes values
61    attrs: HashMap<&'static str, String>,
62}
63
64impl SpanExtension for SpanExtAttrs {}
65
66impl tracing::field::Visit for SpanExtAttrs {
67    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
68        let value = format!("{value:?}");
69        self.attrs.insert(field.name(), value);
70    }
71}
72
73/// A span extensison to record timing info
74#[derive(Debug)]
75pub struct SpanExtTiming {
76    /// Instant when the span was entered
77    pub entered: Instant,
78}
79
80impl Default for SpanExtTiming {
81    fn default() -> Self {
82        Self {
83            entered: Instant::now(),
84        }
85    }
86}
87
88impl SpanExtension for SpanExtTiming {}
89
90/// A visitor for events
91///
92/// The visitor saves the event data
93#[derive(Debug, Default)]
94pub struct EventVisitor {
95    /// Fields
96    fields: HashMap<&'static str, String>,
97}
98
99impl tracing::field::Visit for EventVisitor {
100    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
101        let value_str = format!("{value:?}");
102        self.fields.insert(field.name(), value_str);
103    }
104}
105
106impl EventVisitor {
107    /// Records an event fields
108    ///
109    /// Returns the event message and the event fields
110    pub fn record_event(event: &tracing::Event) -> Self {
111        let mut f_visitor = EventVisitor::default();
112        event.record(&mut f_visitor);
113        f_visitor
114    }
115
116    /// Returns the event message
117    pub fn message(&self) -> &str {
118        match self.fields.get("message") {
119            Some(s) => s,
120            None => {
121                panic!("Event message not found")
122            }
123        }
124    }
125
126    /// Returns the event fields (exc. message)
127    pub fn meta_fields(&self) -> HashMap<&'static str, &str> {
128        self.fields
129            .iter()
130            .filter_map(|(k, v)| {
131                if *k == "message" {
132                    return None;
133                }
134                Some((*k, v.as_str()))
135            })
136            .collect()
137    }
138}