test_span/
layer.rs

1use ::daggy::Dag;
2use ::tracing::span;
3use ::tracing::{Event, Metadata};
4use indexmap::IndexMap;
5use once_cell::sync::Lazy;
6use tracing_subscriber::layer::Context;
7use tracing_subscriber::registry::LookupSpan;
8
9use crate::log::LogsRecorder;
10use crate::record::Recorder;
11use crate::report::ALL_DAGS;
12use crate::LazyMutex;
13
14pub(crate) static ALL_SPANS: LazyMutex<IndexMap<u64, Recorder>> = Lazy::new(Default::default);
15pub(crate) static ALL_LOGS: LazyMutex<LogsRecorder> = Lazy::new(Default::default);
16pub(crate) static SPAN_ID_TO_ROOT_AND_NODE_INDEX: LazyMutex<
17    IndexMap<u64, (u64, daggy::NodeIndex)>,
18> = Lazy::new(Default::default);
19
20#[derive(Debug, Default)]
21pub struct Layer {}
22
23impl Layer {
24    fn record(&self, id: span::Id, record: &span::Record<'_>) {
25        ALL_SPANS
26            .lock()
27            .unwrap()
28            .get_mut(&id.into_u64())
29            .unwrap_or_else(|| panic!("no record for id {}", id.into_u64()))
30            .record(record);
31    }
32
33    fn event(&self, event: &Event<'_>, ctx: Context<'_, impl tracing::Subscriber>) {
34        let current_span = ctx.current_span();
35        let current_span = current_span.id().cloned();
36        ALL_LOGS.lock().unwrap().event(current_span, event);
37    }
38
39    fn attributes(
40        &self,
41        span_id: span::Id,
42        attributes: &span::Attributes<'_>,
43        parent_id: Option<span::Id>,
44    ) {
45        let raw_span_id = span_id.into_u64();
46
47        if let Some(id) = parent_id {
48            // We have a parent, we can store the span in the right DAG
49            let raw_parent_id = id.into_u64();
50
51            let mut id_to_node_index = SPAN_ID_TO_ROOT_AND_NODE_INDEX.lock().unwrap();
52
53            let (root_span_id, parent_node_index) = id_to_node_index
54                .get(&raw_parent_id)
55                .copied()
56                .unwrap_or_else(|| panic!("missing parent attributes for {}.", raw_parent_id));
57
58            let (_, node_index) =
59                if let Some(span_dag) = ALL_DAGS.lock().unwrap().get_mut(&root_span_id) {
60                    span_dag.add_child(parent_node_index, (), raw_span_id)
61                } else {
62                    panic!("missing dag for root {}", root_span_id);
63                };
64
65            id_to_node_index.insert(raw_span_id, (root_span_id, node_index));
66        } else {
67            // We're dealing with a root, let's create a new DAG
68            let mut new_dag: Dag<u64, ()> = Default::default();
69            let root_index = new_dag.add_node(raw_span_id);
70
71            // The span is the root here
72            SPAN_ID_TO_ROOT_AND_NODE_INDEX
73                .lock()
74                .unwrap()
75                .insert(raw_span_id, (raw_span_id, root_index));
76
77            let mut all_dags = ALL_DAGS.lock().unwrap();
78            all_dags.insert(raw_span_id, new_dag);
79        }
80
81        ALL_SPANS
82            .lock()
83            .unwrap()
84            .entry(raw_span_id)
85            .or_default()
86            .attributes(span_id, attributes);
87    }
88}
89
90impl<S> tracing_subscriber::Layer<S> for Layer
91where
92    S: tracing::Subscriber + for<'lookup> LookupSpan<'lookup>,
93{
94    fn register_callsite(
95        &self,
96        _metadata: &'static Metadata<'static>,
97    ) -> tracing::subscriber::Interest {
98        tracing::subscriber::Interest::always()
99    }
100
101    fn on_new_span(
102        &self,
103        attrs: &span::Attributes<'_>,
104        id: &span::Id,
105        ctx: tracing_subscriber::layer::Context<'_, S>,
106    ) {
107        let maybe_parent_id = attrs
108            .parent()
109            .cloned()
110            .or_else(|| ctx.current_span().id().cloned());
111
112        self.attributes(id.clone(), attrs, maybe_parent_id)
113    }
114
115    fn on_record(
116        &self,
117        span: &span::Id,
118        values: &span::Record<'_>,
119        _ctx: tracing_subscriber::layer::Context<'_, S>,
120    ) {
121        self.record(span.clone(), values)
122    }
123
124    fn on_event(&self, event: &Event<'_>, ctx: tracing_subscriber::layer::Context<'_, S>) {
125        self.event(event, ctx)
126    }
127}
test_span/layer.rs

test_span/
layer.rs
