tracing_capture/
layer.rs

1//! `CaptureLayer` and related types.
2
3use std::{
4    fmt, ops,
5    sync::{Arc, RwLock},
6};
7
8use id_arena::Arena;
9use tracing_core::{
10    span::{Attributes, Id, Record},
11    Event, Metadata, Subscriber,
12};
13use tracing_subscriber::{
14    layer::{Context, Filter},
15    registry::LookupSpan,
16    Layer,
17};
18use tracing_tunnel::TracedValues;
19
20use crate::{
21    CapturedEvent, CapturedEventId, CapturedEventInner, CapturedEvents, CapturedSpan,
22    CapturedSpanId, CapturedSpanInner, CapturedSpans, SpanStats,
23};
24
25/// Storage of captured tracing information.
26///
27/// `Storage` instances are not created directly; instead, they are wrapped in [`SharedStorage`]
28/// and can be accessed via [`lock()`](SharedStorage::lock()).
29#[derive(Debug)]
30pub struct Storage {
31    pub(crate) spans: Arena<CapturedSpanInner>,
32    pub(crate) events: Arena<CapturedEventInner>,
33    root_span_ids: Vec<CapturedSpanId>,
34    root_event_ids: Vec<CapturedEventId>,
35}
36
37impl Storage {
38    pub(crate) fn new() -> Self {
39        Self {
40            spans: Arena::new(),
41            events: Arena::new(),
42            root_span_ids: vec![],
43            root_event_ids: vec![],
44        }
45    }
46
47    pub(crate) fn span(&self, id: CapturedSpanId) -> CapturedSpan<'_> {
48        CapturedSpan {
49            inner: &self.spans[id],
50            storage: self,
51        }
52    }
53
54    pub(crate) fn event(&self, id: CapturedEventId) -> CapturedEvent<'_> {
55        CapturedEvent {
56            inner: &self.events[id],
57            storage: self,
58        }
59    }
60
61    /// Iterates over captured spans in the order of capture.
62    pub fn all_spans(&self) -> CapturedSpans<'_> {
63        CapturedSpans::from_arena(self)
64    }
65
66    /// Iterates over root spans (i.e., spans that do not have a captured parent span)
67    /// in the order of capture.
68    pub fn root_spans(&self) -> CapturedSpans<'_> {
69        CapturedSpans::from_slice(self, &self.root_span_ids)
70    }
71
72    /// Iterates over all captured events in the order of capture.
73    pub fn all_events(&self) -> CapturedEvents<'_> {
74        CapturedEvents::from_arena(self)
75    }
76
77    /// Iterates over root events (i.e., events that do not have a captured parent span)
78    /// in the order of capture.
79    pub fn root_events(&self) -> CapturedEvents<'_> {
80        CapturedEvents::from_slice(self, &self.root_event_ids)
81    }
82
83    pub(crate) fn push_span(
84        &mut self,
85        metadata: &'static Metadata<'static>,
86        values: TracedValues<&'static str>,
87        parent_id: Option<CapturedSpanId>,
88    ) -> CapturedSpanId {
89        let span_id = self.spans.alloc_with_id(|id| CapturedSpanInner {
90            metadata,
91            values,
92            stats: SpanStats::default(),
93            id,
94            parent_id,
95            child_ids: vec![],
96            event_ids: vec![],
97            follows_from_ids: vec![],
98        });
99        if let Some(parent_id) = parent_id {
100            let span = self.spans.get_mut(parent_id).unwrap();
101            span.child_ids.push(span_id);
102        } else {
103            self.root_span_ids.push(span_id);
104        }
105        span_id
106    }
107
108    fn on_span_enter(&mut self, id: CapturedSpanId) {
109        let span = self.spans.get_mut(id).unwrap();
110        span.stats.entered += 1;
111    }
112
113    fn on_span_exit(&mut self, id: CapturedSpanId) {
114        let span = self.spans.get_mut(id).unwrap();
115        span.stats.exited += 1;
116    }
117
118    fn on_span_closed(&mut self, id: CapturedSpanId) {
119        let span = self.spans.get_mut(id).unwrap();
120        span.stats.is_closed = true;
121    }
122
123    fn on_record(&mut self, id: CapturedSpanId, values: TracedValues<&'static str>) {
124        let span = self.spans.get_mut(id).unwrap();
125        span.values.extend(values);
126    }
127
128    fn on_follows_from(&mut self, id: CapturedSpanId, follows_id: CapturedSpanId) {
129        let span = self.spans.get_mut(id).unwrap();
130        span.follows_from_ids.push(follows_id);
131    }
132
133    pub(crate) fn push_event(
134        &mut self,
135        metadata: &'static Metadata<'static>,
136        values: TracedValues<&'static str>,
137        parent_id: Option<CapturedSpanId>,
138    ) -> CapturedEventId {
139        let event_id = self.events.alloc_with_id(|id| CapturedEventInner {
140            metadata,
141            values,
142            id,
143            parent_id,
144        });
145        if let Some(parent_id) = parent_id {
146            let span = self.spans.get_mut(parent_id).unwrap();
147            span.event_ids.push(event_id);
148        } else {
149            self.root_event_ids.push(event_id);
150        }
151        event_id
152    }
153}
154
155/// Shared wrapper for tracing [`Storage`].
156#[derive(Debug, Clone)]
157pub struct SharedStorage {
158    inner: Arc<RwLock<Storage>>,
159}
160
161impl Default for SharedStorage {
162    fn default() -> Self {
163        Self {
164            inner: Arc::new(RwLock::new(Storage::new())),
165        }
166    }
167}
168
169#[allow(clippy::missing_panics_doc)] // lock poisoning propagation
170impl SharedStorage {
171    /// Locks the underlying [`Storage`] for exclusive access. While the lock is held,
172    /// capturing cannot progress; beware of deadlocks!
173    pub fn lock(&self) -> impl ops::Deref<Target = Storage> + '_ {
174        self.inner
175            .read()
176            .expect("failed accessing shared tracing data storage")
177    }
178}
179
180/// Tracing [`Layer`] that captures (optionally filtered) spans and events.
181///
182/// The layer can optionally filter spans and events in addition to global [`Subscriber`] filtering.
183/// This could be used instead of per-layer filtering if it's not supported by the `Subscriber`.
184/// Keep in mind that without filtering, `CaptureLayer` can capture a lot of
185/// unnecessary spans / events.
186///
187/// Captured events are [tied](CapturedSpan::events()) to the nearest captured span
188/// in the span hierarchy. If no entered spans are captured when the event is emitted,
189/// the event will be captured in [`Storage::root_events()`].
190///
191/// # Examples
192///
193/// See [crate-level docs](index.html) for an example of usage.
194pub struct CaptureLayer<S> {
195    filter: Option<Box<dyn Filter<S> + Send + Sync>>,
196    storage: Arc<RwLock<Storage>>,
197}
198
199impl<S> fmt::Debug for CaptureLayer<S> {
200    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
201        formatter
202            .debug_struct("CaptureLayer")
203            .field("filter", &self.filter.as_ref().map(|_| "Filter"))
204            .field("storage", &self.storage)
205            .finish()
206    }
207}
208
209impl<S> CaptureLayer<S>
210where
211    S: Subscriber + for<'a> LookupSpan<'a>,
212{
213    /// Creates a new layer that will use the specified `storage` to store captured data.
214    /// Captured spans are not filtered; like any [`Layer`], filtering can be set up
215    /// on the layer or subscriber level.
216    pub fn new(storage: &SharedStorage) -> Self {
217        Self {
218            filter: None,
219            storage: Arc::clone(&storage.inner),
220        }
221    }
222
223    /// Specifies filtering for this layer. Unlike with [per-layer filtering](Layer::with_filter()),
224    /// the resulting layer will perform filtering for all [`Subscriber`]s, not just [`Registry`].
225    ///
226    /// [`Registry`]: tracing_subscriber::Registry
227    #[must_use]
228    pub fn with_filter<F>(mut self, filter: F) -> Self
229    where
230        F: Filter<S> + Send + Sync + 'static,
231    {
232        self.filter = Some(Box::new(filter));
233        self
234    }
235
236    fn enabled(&self, metadata: &Metadata<'_>, ctx: &Context<'_, S>) -> bool {
237        self.filter
238            .as_deref()
239            .map_or(true, |filter| filter.enabled(metadata, ctx))
240    }
241
242    fn lock(&self) -> impl ops::DerefMut<Target = Storage> + '_ {
243        self.storage
244            .write()
245            .expect("failed locking shared tracing data storage for write")
246    }
247}
248
249impl<S> Layer<S> for CaptureLayer<S>
250where
251    S: Subscriber + for<'a> LookupSpan<'a>,
252{
253    fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
254        if !self.enabled(attrs.metadata(), &ctx) {
255            return;
256        }
257
258        let parent_id = if let Some(mut scope) = ctx.span_scope(id) {
259            scope.find_map(|span| span.extensions().get::<CapturedSpanId>().copied())
260        } else {
261            None
262        };
263        let values = TracedValues::from_values(attrs.values());
264        let arena_id = self.lock().push_span(attrs.metadata(), values, parent_id);
265        ctx.span(id).unwrap().extensions_mut().insert(arena_id);
266    }
267
268    fn on_record(&self, id: &Id, values: &Record<'_>, ctx: Context<'_, S>) {
269        let span = ctx.span(id).unwrap();
270        if let Some(id) = span.extensions().get::<CapturedSpanId>().copied() {
271            self.lock().on_record(id, TracedValues::from_record(values));
272        };
273    }
274
275    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
276        if !self.enabled(event.metadata(), &ctx) {
277            return;
278        }
279
280        let parent_id = if let Some(mut scope) = ctx.event_scope(event) {
281            scope.find_map(|span| span.extensions().get::<CapturedSpanId>().copied())
282        } else {
283            None
284        };
285        self.lock()
286            .push_event(event.metadata(), TracedValues::from_event(event), parent_id);
287    }
288
289    fn on_enter(&self, id: &Id, ctx: Context<'_, S>) {
290        let span = ctx.span(id).unwrap();
291        if let Some(id) = span.extensions().get::<CapturedSpanId>().copied() {
292            self.lock().on_span_enter(id);
293        };
294    }
295
296    fn on_exit(&self, id: &Id, ctx: Context<'_, S>) {
297        let span = ctx.span(id).unwrap();
298        if let Some(id) = span.extensions().get::<CapturedSpanId>().copied() {
299            self.lock().on_span_exit(id);
300        };
301    }
302
303    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
304        let span = ctx.span(&id).unwrap();
305        if let Some(id) = span.extensions().get::<CapturedSpanId>().copied() {
306            self.lock().on_span_closed(id);
307        };
308    }
309
310    fn on_follows_from(&self, id: &Id, follows_id: &Id, ctx: Context<'_, S>) {
311        let span = ctx.span(id).unwrap();
312        let follows = ctx.span(follows_id).unwrap();
313        if let Some(id) = span.extensions().get::<CapturedSpanId>().copied() {
314            if let Some(follows_id) = follows.extensions().get::<CapturedSpanId>().copied() {
315                self.lock().on_follows_from(id, follows_id);
316            }
317        };
318    }
319}
tracing_capture/layer.rs

tracing_capture/
layer.rs
