tracing_durations_export/
lib.rs

1//! Record and visualize which spans are active in parallel.
2//!
3//! ## Usage
4//!
5//! ```rust
6//! use std::fs::File;
7//! use std::io::BufWriter;
8//! use tracing_durations_export::{DurationsLayer, DurationsLayerBuilder, DurationsLayerDropGuard};
9//! use tracing_subscriber::layer::SubscriberExt;
10//! use tracing_subscriber::{registry::Registry, fmt};
11//!
12//! fn setup_global_subscriber() -> DurationsLayerDropGuard {
13//!     let fmt_layer = fmt::Layer::default();
14//!     let (duration_layer, guard) = DurationsLayerBuilder::default()
15//!         .durations_file("traces.ndjson")
16//!         // Available with the `plot` feature
17//!         // .plot_file("traces.svg")
18//!         .build()
19//!         .unwrap();
20//!     let subscriber = Registry::default()
21//!         .with(fmt_layer)
22//!         .with(duration_layer);
23//!
24//!     tracing::subscriber::set_global_default(subscriber).unwrap();
25//!
26//!     guard
27//! }
28//!
29//! // your code here ...
30//! ```
31//!
32//! The output file will look something like below, where each section where a span is active is one line.
33//!
34//! ```ndjson
35//! [...]
36//! {"id":6,"name":"read_cache","start":{"secs":0,"nanos":122457871},"end":{"secs":0,"nanos":122463135},"parents":[5],"fields":{"id":"2"}}
37//! {"id":5,"name":"cached_network_request","start":{"secs":0,"nanos":122433854},"end":{"secs":0,"nanos":122499689},"parents":[],"fields":{"id":"2","api":"https://example.net/cached"}}
38//! {"id":9007474132647937,"name":"parse_cache","start":{"secs":0,"nanos":122625724},"end":{"secs":0,"nanos":125791908},"parents":[],"fields":{}}
39//! {"id":5,"name":"cached_network_request","start":{"secs":0,"nanos":125973025},"end":{"secs":0,"nanos":126007737},"parents":[],"fields":{"id":"2","api":"https://example.net/cached"}}
40//! {"id":5,"name":"cached_network_request","start":{"secs":0,"nanos":126061739},"end":{"secs":0,"nanos":126066912},"parents":[],"fields":{"id":"2","api":"https://example.net/cached"}}
41//! {"id":2251799813685254,"name":"read_cache","start":{"secs":0,"nanos":126157156},"end":{"secs":0,"nanos":126193547},"parents":[2251799813685253],"fields":{"id":"3"}}
42//! {"id":2251799813685253,"name":"cached_network_request","start":{"secs":0,"nanos":126144140},"end":{"secs":0,"nanos":126213181},"parents":[],"fields":{"api":"https://example.net/cached","id":"3"}}
43//! {"id":27021597764222977,"name":"make_network_request","start":{"secs":0,"nanos":128343009},"end":{"secs":0,"nanos":128383121},"parents":[13510798882111491],"fields":{"api":"https://example.net/cached","id":"0"}}```
44//! [...]
45//! ```
46//!
47//! Note that 0 is the time of the first span, not the start of the process.
48
49use fs::File;
50use once_cell::sync::Lazy;
51use serde::Serialize;
52use std::collections::hash_map::RandomState;
53use std::collections::HashMap;
54use std::fmt::Debug;
55use std::io::{BufWriter, Write};
56use std::marker::PhantomData;
57use std::path::PathBuf;
58use std::sync::{Arc, Mutex};
59use std::time::{Duration, Instant};
60use std::{io, iter};
61use tracing::field::Field;
62use tracing::{span, Subscriber};
63use tracing_subscriber::layer::Context;
64use tracing_subscriber::registry::LookupSpan;
65use tracing_subscriber::Layer;
66
67#[cfg(feature = "plot")]
68pub mod plot;
69
70/// A zero timestamp initialized by the first span
71static START: Lazy<Instant> = Lazy::new(Instant::now);
72
73/// A recorded active section of a span.
74#[derive(Serialize)]
75// Remove bound on `RandomState`
76#[serde(bound(serialize = ""))]
77pub struct SpanInfo<'a, RS = RandomState> {
78    pub id: u64,
79    pub name: &'static str,
80    pub start: Duration,
81    pub end: Duration,
82    pub parents: Option<&'a [u64]>,
83    pub is_main_thread: bool,
84    pub fields: Option<&'a HashMap<&'static str, String, RS>>,
85}
86
87pub struct DurationsLayerBuilder {
88    /// See [`DurationsLayerBuilder::with_fields`].
89    with_fields: bool,
90    /// See [`DurationsLayerBuilder::with_parents`].
91    with_parents: bool,
92    /// See [`DurationsLayerBuilder::durations_file`].
93    durations_file: Option<PathBuf>,
94    /// See [`DurationsLayerBuilder::plot_file`].
95    #[cfg(feature = "plot")]
96    plot_file: Option<PathBuf>,
97    #[cfg(feature = "plot")]
98    plot_config: plot::PlotConfig,
99    #[cfg(feature = "plot")]
100    plot_layout: plot::PlotLayout,
101}
102
103impl Default for DurationsLayerBuilder {
104    fn default() -> Self {
105        Self {
106            with_fields: true,
107            with_parents: true,
108            durations_file: None,
109            #[cfg(feature = "plot")]
110            plot_file: None,
111            #[cfg(feature = "plot")]
112            plot_config: plot::PlotConfig::default(),
113            #[cfg(feature = "plot")]
114            plot_layout: plot::PlotLayout::default(),
115        }
116    }
117}
118
119impl DurationsLayerBuilder {
120    /// This function needs to be called on the (tokio) main thread for accurate reporting.
121    pub fn build<S>(self) -> io::Result<(DurationsLayer<S>, DurationsLayerDropGuard)> {
122        let out = self
123            .durations_file
124            .map(|file| File::create(file).map(BufWriter::new))
125            .transpose()?;
126        let layer = DurationsLayer {
127            main_thead_id: std::thread::current().id(),
128            start_index: Mutex::default(),
129            fields: Mutex::default(),
130            is_main_thread: Mutex::new(Default::default()),
131            out: Arc::new(Mutex::new(out)),
132            #[cfg(feature = "plot")]
133            plot_data: Arc::new(Mutex::default()),
134            #[cfg(feature = "plot")]
135            plot_file: self.plot_file,
136            with_fields: self.with_fields,
137            with_parents: self.with_parents,
138            #[cfg(feature = "plot")]
139            plot_config: self.plot_config,
140            #[cfg(feature = "plot")]
141            plot_layout: self.plot_layout,
142            _inner: PhantomData,
143        };
144        let guard = layer.drop_guard();
145        Ok((layer, guard))
146    }
147
148    /// Whether to record the fields passed to the span (default: `true`).
149    ///
150    /// # Example
151    ///
152    /// Span:
153    /// ```rust
154    /// # use tracing::info_span;
155    /// info_span!("make_request", host = "example.org", object = 10);
156    /// ```
157    ///
158    /// With `true`:
159    /// ```json
160    /// {"id":4,"start":{"secs":0,"nanos":446},"end":{"secs":0,"nanos":448},"name":"make_request","parents":[1,3],"fields":{"host":"example.org","object":"10"}}
161    /// ```
162    ///
163    /// With `false`:
164    /// ```json
165    /// {"id":4,"start":{"secs":0,"nanos":446},"end":{"secs":0,"nanos":448},"name":"make_request","parents":[1,3]}
166    /// ```
167    pub fn with_fields(self, enabled: bool) -> Self {
168        Self {
169            with_fields: enabled,
170            ..self
171        }
172    }
173
174    /// Whether to record the ids of the parent spans (default: `true`).
175    ///
176    /// # Example
177    ///
178    /// Span:
179    /// ```rust
180    /// # use tracing::info_span;
181    /// info_span!("make_request", host = "example.org", object = 10);
182    /// ```
183    ///
184    /// With `true`:
185    /// ```json
186    /// {"id":4,"start":{"secs":0,"nanos":446},"end":{"secs":0,"nanos":448},"name":"make_request","parents":[1,3],"fields":{"host":"example.org","object":"10"}}
187    /// ```
188    ///
189    /// With `false`:
190    /// ```json
191    /// {"id":4,"start":{"secs":0,"nanos":446},"end":{"secs":0,"nanos":448},"name":"make_request","fields":{"host":"example.org","object":"10"}}
192    /// ```
193    pub fn with_parents(self, enabled: bool) -> Self {
194        Self {
195            with_parents: enabled,
196            ..self
197        }
198    }
199
200    /// Record all span active durations as ndjson.
201    ///
202    /// Example output line, see [module level documentation](`crate`) for more details.
203    ///
204    /// ```ndjson
205    /// {"id":6,"name":"read_cache","start":{"secs":0,"nanos":122457871},"end":{"secs":0,"nanos":122463135},"parents":[3,4],"fields":{"id":"2"}}
206    /// ```
207    ///
208    /// The file is flushed when [`DurationsLayerDropGuard`] is dropped.
209    pub fn durations_file(self, file: impl Into<PathBuf>) -> Self {
210        Self {
211            durations_file: Some(file.into()),
212            ..self
213        }
214    }
215
216    /// Plot the result and save them as svg.
217    ///
218    /// TODO(konstin): Figure out how to embed an svg in rustdoc.
219    ///
220    /// The file is written when [`DurationsLayerDropGuard`] is dropped.
221    #[cfg(feature = "plot")]
222    pub fn plot_file(self, file: impl Into<PathBuf>) -> Self {
223        Self {
224            plot_file: Some(file.into()),
225            ..self
226        }
227    }
228
229    #[cfg(feature = "plot")]
230    pub fn plot_config(self, plot_config: plot::PlotConfig) -> Self {
231        Self {
232            plot_config,
233            ..self
234        }
235    }
236}
237
238type CollectedFields<RS> = HashMap<&'static str, String, RS>;
239
240#[derive(Default)]
241struct FieldsCollector<RS = RandomState>(CollectedFields<RS>);
242
243impl tracing::field::Visit for FieldsCollector {
244    fn record_str(&mut self, field: &Field, value: &str) {
245        self.0.insert(field.name(), value.to_string());
246    }
247
248    fn record_debug(&mut self, field: &Field, value: &dyn Debug) {
249        self.0.insert(field.name(), format!("{value:?}"));
250    }
251}
252
253/// On drop, flush the output writer and, if applicable, write the plot.
254pub struct DurationsLayerDropGuard {
255    out: Arc<Mutex<Option<BufWriter<File>>>>,
256    #[cfg(feature = "plot")]
257    plot_file: Option<PathBuf>,
258    #[cfg(feature = "plot")]
259    plot_data: Arc<Mutex<Vec<plot::OwnedSpanInfo>>>,
260    #[cfg(feature = "plot")]
261    plot_config: plot::PlotConfig,
262    #[cfg(feature = "plot")]
263    plot_layout: plot::PlotLayout,
264}
265
266impl Drop for DurationsLayerDropGuard {
267    fn drop(&mut self) {
268        if let Some(out) = self.out.lock().expect("There was a prior panic").as_mut() {
269            if let Err(err) = out.flush() {
270                eprintln!("`DurationLayer` failed to flush out file: {err}");
271            }
272        }
273
274        #[cfg(feature = "plot")]
275        {
276            if let Some(plot_file) = &self.plot_file {
277                let end = self
278                    .plot_data
279                    .lock()
280                    .unwrap()
281                    .iter()
282                    .map(|span| span.end)
283                    .max();
284                // This is some only if the plot option was and any spans were recorded
285                if let Some(end) = end {
286                    let svg = plot::plot(
287                        &self.plot_data.lock().expect("There was a prior panic"),
288                        end,
289                        &self.plot_config,
290                        &self.plot_layout,
291                    );
292                    if let Err(err) = svg::save(plot_file, &svg) {
293                        eprintln!("`DurationLayer` failed to write plot: {err}");
294                    }
295                }
296            }
297        }
298    }
299}
300
301/// `tracing` layer to record which spans are active in parallel as ndjson.
302pub struct DurationsLayer<S, RS = RandomState> {
303    main_thead_id: std::thread::ThreadId,
304    // Each of the 3 fields below has different initialization:
305    //
306    // TODO(konstin): Attach this as span extension instead?
307    start_index: Mutex<HashMap<span::Id, Duration, RS>>,
308    // TODO(konstin): Attach this as span extension instead?
309    fields: Mutex<HashMap<span::Id, CollectedFields<RS>>>,
310    // TODO(konstin): Attach this as span extension instead?
311    is_main_thread: Mutex<HashMap<span::Id, bool>>,
312    out: Arc<Mutex<Option<BufWriter<File>>>>,
313    #[cfg(feature = "plot")]
314    plot_data: Arc<Mutex<Vec<plot::OwnedSpanInfo>>>,
315    #[cfg(feature = "plot")]
316    plot_file: Option<PathBuf>,
317    with_fields: bool,
318    with_parents: bool,
319    #[cfg(feature = "plot")]
320    plot_config: plot::PlotConfig,
321    #[cfg(feature = "plot")]
322    plot_layout: plot::PlotLayout,
323    _inner: PhantomData<S>,
324}
325
326impl<S> DurationsLayer<S> {
327    fn drop_guard(&self) -> DurationsLayerDropGuard {
328        DurationsLayerDropGuard {
329            out: self.out.clone(),
330            #[cfg(feature = "plot")]
331            plot_file: self.plot_file.clone(),
332            #[cfg(feature = "plot")]
333            plot_data: self.plot_data.clone(),
334            #[cfg(feature = "plot")]
335            plot_config: self.plot_config.clone(),
336            #[cfg(feature = "plot")]
337            plot_layout: self.plot_layout.clone(),
338        }
339    }
340}
341
342impl<S> Layer<S> for DurationsLayer<S>
343where
344    S: Subscriber + for<'span> LookupSpan<'span>,
345{
346    /// Record the fields
347    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, _ctx: Context<'_, S>) {
348        // We only get the fields here (i think they aren't stored with the span?), so we have to record them here
349        if self.with_fields {
350            let mut visitor = FieldsCollector::default();
351            attrs.record(&mut visitor);
352            self.fields
353                .lock()
354                .expect("There was a prior panic")
355                .insert(id.clone(), visitor.0);
356        }
357        self.is_main_thread
358            .lock()
359            .expect("There was a prior panic")
360            .insert(
361                id.clone(),
362                self.main_thead_id == std::thread::current().id(),
363            );
364    }
365
366    /// Record the start timestamp
367    fn on_enter(&self, id: &span::Id, _ctx: Context<'_, S>) {
368        self.start_index
369            .lock()
370            .unwrap()
371            .insert(id.clone(), START.elapsed());
372    }
373
374    /// Write a record to the ndjson writer
375    fn on_exit(&self, id: &span::Id, ctx: Context<'_, S>) {
376        let span = ctx.span(id).unwrap();
377        let parents = if self.with_parents {
378            let parents = iter::successors(span.parent(), |span| span.parent())
379                .map(|span| span.id().into_u64())
380                .collect::<Vec<_>>();
381            Some(parents)
382        } else {
383            None
384        };
385        let attributes = self.fields.lock().expect("There was a prior panic");
386        let fields = attributes.get(id);
387        debug_assert!(
388            !self.with_fields || fields.is_some(),
389            "Expected fields to be record for span {} {}",
390            span.name(),
391            id.into_u64()
392        );
393
394        let is_main_thread = self.main_thead_id == std::thread::current().id();
395        let span_info = SpanInfo {
396            id: id.into_u64(),
397            name: span.name(),
398            start: self.start_index.lock().expect("There was a prior panic")[id],
399            end: START.elapsed(),
400            parents: parents.as_deref(),
401            is_main_thread,
402            fields,
403        };
404        // https://github.com/rust-lang/rust-clippy/pull/12892
405        #[allow(clippy::needless_borrows_for_generic_args)]
406        if let Some(mut writer) = self.out.lock().expect("There was a prior panic").as_mut() {
407            // ndjson, write the json and then a newline
408            serde_json::to_writer(&mut writer, &span_info).unwrap();
409            writeln!(&mut writer).unwrap();
410        }
411
412        #[cfg(feature = "plot")]
413        {
414            if self.plot_file.is_some() {
415                self.plot_data
416                    .lock()
417                    .expect("There was a prior panic")
418                    .push(plot::OwnedSpanInfo {
419                        id: id.into_u64(),
420                        name: span.name().to_string(),
421                        start: self.start_index.lock().expect("There was a prior panic")[id],
422                        end: START.elapsed(),
423                        parents,
424                        is_main_thread,
425                        fields: fields.map(|fields| {
426                            fields
427                                .iter()
428                                .map(|(key, value)| (key.to_string(), value.to_string()))
429                                .collect()
430                        }),
431                    })
432            }
433        }
434    }
435}
tracing_durations_export/lib.rs

tracing_durations_export/
lib.rs
