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

tracing_durations_export/
lib.rs
