tracing_calltree/
layer.rs

1use crate::builder::CallTreeInner;
2use crate::sample::TimingSample;
3use crate::tree::{CallTreeNode, NodeKey};
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6use tracing::Subscriber;
7use tracing::span::{Attributes, Id, Record};
8use tracing_subscriber::layer::{Context, Layer};
9use tracing_subscriber::registry::LookupSpan;
10
11#[derive(Clone)]
12pub struct CallTreeLayer {
13    inner: Arc<CallTreeInner>,
14}
15
16impl CallTreeLayer {
17    pub(crate) fn new(inner: Arc<CallTreeInner>) -> Self {
18        Self { inner }
19    }
20}
21
22#[derive(Clone)]
23struct SpanCallTreeState {
24    nearest_profiled_node: Option<Arc<CallTreeNode>>,
25    timing: Option<SpanTimingState>,
26}
27
28#[derive(Clone)]
29struct SpanTimingState {
30    node: Arc<CallTreeNode>,
31    first_enter: Option<Instant>,
32    active_started: Option<Instant>,
33    active_elapsed: Duration,
34    enter_depth: u32,
35}
36
37impl SpanTimingState {
38    fn new(node: Arc<CallTreeNode>) -> Self {
39        Self {
40            node,
41            first_enter: None,
42            active_started: None,
43            active_elapsed: Duration::ZERO,
44            enter_depth: 0,
45        }
46    }
47}
48
49impl<S> Layer<S> for CallTreeLayer
50where
51    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
52{
53    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
54        let Some(span) = ctx.span(id) else {
55            return;
56        };
57
58        let parent_profiled = parent_profiled_node(attrs, &ctx);
59        let metadata = span.metadata();
60        let is_profiled = (self.inner.config.filter)(metadata);
61
62        let mut nearest_profiled_node = parent_profiled.clone();
63        let timing = if is_profiled {
64            let depth = parent_profiled
65                .as_ref()
66                .map_or(1, |parent| parent.depth().saturating_add(1));
67
68            if depth > self.inner.config.max_depth {
69                None
70            } else {
71                let key = NodeKey::from_metadata(metadata);
72                self.inner
73                    .state
74                    .get_or_create_node(
75                        parent_profiled.as_ref(),
76                        key,
77                        depth,
78                        self.inner.config.max_nodes,
79                    )
80                    .map(|node| {
81                        nearest_profiled_node = Some(node.clone());
82                        SpanTimingState::new(node)
83                    })
84            }
85        } else {
86            None
87        };
88
89        span.extensions_mut().insert(SpanCallTreeState {
90            nearest_profiled_node,
91            timing,
92        });
93    }
94
95    fn on_record(&self, _span: &Id, _values: &Record<'_>, _ctx: Context<'_, S>) {}
96
97    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
98        let Some(span) = ctx.span(id) else {
99            return;
100        };
101
102        let mut extensions = span.extensions_mut();
103        let Some(state) = extensions.get_mut::<SpanCallTreeState>() else {
104            return;
105        };
106        let Some(timing) = state.timing.as_mut() else {
107            return;
108        };
109
110        if timing.enter_depth == 0 {
111            let now = Instant::now();
112            timing.first_enter.get_or_insert(now);
113            timing.active_started = Some(now);
114        }
115
116        timing.enter_depth = timing.enter_depth.saturating_add(1);
117    }
118
119    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
120        let Some(span) = ctx.span(id) else {
121            return;
122        };
123
124        let mut extensions = span.extensions_mut();
125        let Some(state) = extensions.get_mut::<SpanCallTreeState>() else {
126            return;
127        };
128        let Some(timing) = state.timing.as_mut() else {
129            return;
130        };
131
132        if timing.enter_depth == 0 {
133            return;
134        }
135
136        timing.enter_depth -= 1;
137        if timing.enter_depth == 0 {
138            if let Some(started) = timing.active_started.take() {
139                timing.active_elapsed += Instant::now().saturating_duration_since(started);
140            }
141        }
142    }
143
144    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
145        let Some(span) = ctx.span(&id) else {
146            return;
147        };
148
149        let state = span.extensions_mut().remove::<SpanCallTreeState>();
150        let Some(mut timing) = state.and_then(|state| state.timing) else {
151            return;
152        };
153
154        let Some(first_enter) = timing.first_enter else {
155            return;
156        };
157
158        let now = Instant::now();
159        if timing.enter_depth > 0 {
160            if let Some(started) = timing.active_started.take() {
161                timing.active_elapsed += now.saturating_duration_since(started);
162            }
163        }
164
165        timing.node.record_sample(
166            TimingSample::new(
167                now.saturating_duration_since(first_enter),
168                timing.active_elapsed,
169            ),
170            self.inner.config.window_size,
171            self.inner.state.dropped_samples_counter(),
172        );
173    }
174}
175
176fn parent_profiled_node<S>(
177    attrs: &Attributes<'_>,
178    ctx: &Context<'_, S>,
179) -> Option<Arc<CallTreeNode>>
180where
181    S: Subscriber + for<'lookup> LookupSpan<'lookup>,
182{
183    let parent = attrs
184        .parent()
185        .and_then(|parent| ctx.span(parent))
186        .or_else(|| {
187            attrs
188                .is_contextual()
189                .then(|| ctx.lookup_current())
190                .flatten()
191        })?;
192
193    let extensions = parent.extensions();
194    extensions
195        .get::<SpanCallTreeState>()
196        .and_then(|state| state.nearest_profiled_node.clone())
197}
tracing_calltree/layer.rs

tracing_calltree/
layer.rs
