tracing_capture/
layer.rs

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

tracing_capture/
layer.rs
