tracing_callgraph/
lib.rs

1#![allow(clippy::needless_doctest_main)]
2//! A [tracing](https://github.com/tokio-rs/tracing/) [Layer][`GraphLayer`] for generating a call graphs.
3//!
4//! # Overview
5//!
6//! [`tracing`] is a framework for instrumenting Rust programs to collect
7//! scoped, structured, and async-aware diagnostics. `tracing-callgraph` provides helpers
8//! for consuming `tracing` instrumentation that can later be visualized as a
9//! call graph in [Graphviz](http://www.graphviz.org/) `dot` representation.
10//!
11//! ## Layer Setup
12//!
13//! ```rust
14//! use tracing_callgraph::GraphLayer;
15//! use tracing_subscriber::{registry::Registry, prelude::*};
16//!
17//! fn setup_global_subscriber() -> impl Drop {
18//!     let (graph_layer, _guard) = GraphLayer::with_file("./output.dot").unwrap();
19//!     let subscriber = Registry::default().with(graph_layer);
20//!
21//!     tracing::subscriber::set_global_default(subscriber).expect("Could not set global default");
22//!     _guard
23//! }
24//!
25//! #[tracing::instrument]
26//! fn outer_a() {
27//!     inner()
28//! }
29//!
30//! #[tracing::instrument]
31//! fn outer_b() {
32//!     inner()
33//! }
34//!
35//! #[tracing::instrument]
36//! fn inner() {}
37//!
38//! fn main() {
39//!     let _ = setup_global_subscriber();
40//!     outer_a();
41//!     outer_b();
42//! }
43//! ```
44//!
45#![warn(
46    missing_debug_implementations,
47    missing_docs,
48    rust_2018_idioms,
49    unreachable_pub,
50    bad_style,
51    const_err,
52    dead_code,
53    improper_ctypes,
54    non_shorthand_field_patterns,
55    no_mangle_generic_items,
56    overflowing_literals,
57    path_statements,
58    patterns_in_fns_without_body,
59    private_in_public,
60    unconditional_recursion,
61    unused,
62    unused_allocation,
63    unused_comparisons,
64    unused_parens,
65    while_true
66)]
67
68pub use error::Error;
69
70use error::Kind;
71use petgraph::{dot::Dot, graphmap::GraphMap, Directed};
72use std::{
73    fs::File,
74    io::{BufWriter, Write},
75    path::Path,
76    sync::{Arc, Mutex},
77};
78use tracing::{span, Subscriber};
79use tracing_subscriber::{layer::Context, registry::LookupSpan, Layer};
80
81mod error;
82
83type CallGraph = GraphMap<&'static str, usize, Directed>;
84
85/// A `Layer` that records span open events as directed edges in a call graph.
86///
87/// # Dropping and Flushing
88///
89/// To ensure all data is flushed when the program exits, `GraphLayer` exposes
90/// the [`flush_on_drop`] function, which returns a [`FlushGuard`]. The [`FlushGuard`]
91/// will flush the writer when it is dropped. If necessary, it can also be used to manually
92/// flush the writer.
93#[derive(Clone, Debug)]
94pub struct GraphLayer {
95    graph: Arc<Mutex<CallGraph>>,
96    top_node: Option<&'static str>,
97}
98
99impl GraphLayer {
100    /// Add a top node to the graph.
101    #[allow(clippy::clone_double_ref)]
102    pub fn enable_top_node(mut self, name: &'static str) -> Self {
103        self = self.disable_top_node();
104        self.top_node = Some(name.clone());
105        self.graph.lock().unwrap().add_node(name);
106        self
107    }
108
109    /// Remove the top node to the graph.
110    pub fn disable_top_node(mut self) -> Self {
111        if let Some(name) = self.top_node.take() {
112            self.graph.lock().unwrap().remove_node(name);
113        }
114        self
115    }
116}
117
118/// An RAII guard for flushing a writer.
119#[must_use]
120#[derive(Debug)]
121pub struct FlushGuard<W>
122where
123    W: Write + 'static,
124{
125    graph: Arc<Mutex<CallGraph>>,
126    writer: W,
127}
128
129impl<W> FlushGuard<W>
130where
131    W: Write + 'static,
132{
133    /// Flush the internal writer, ensuring that the graph is written.
134    pub fn flush(&mut self) -> Result<(), Error> {
135        let graph = match self.graph.lock() {
136            Ok(graph) => graph,
137            Err(e) => {
138                if !std::thread::panicking() {
139                    panic!("{}", e);
140                } else {
141                    return Ok(());
142                }
143            }
144        };
145        writeln!(self.writer, "{:?}", Dot::new(&*graph))
146            .map_err(Kind::FlushFile)
147            .map_err(Error)?;
148
149        self.writer.flush().map_err(Kind::FlushFile).map_err(Error)
150    }
151}
152
153impl<W> Drop for FlushGuard<W>
154where
155    W: Write + 'static,
156{
157    fn drop(&mut self) {
158        match self.flush() {
159            Ok(_) => (),
160            Err(e) => e.report(),
161        }
162    }
163}
164
165impl Default for GraphLayer {
166    fn default() -> Self {
167        let graph = CallGraph::new();
168        Self {
169            graph: Arc::new(Mutex::new(graph)),
170            top_node: None,
171        }
172    }
173}
174
175impl GraphLayer {
176    /// Returns a new [`GraphLayer`] which constructs the call graph.
177    pub fn new() -> Self {
178        Default::default()
179    }
180
181    /// Returns a [`FlushGuard`] which will flush the `GraphLayer`'s writer when
182    /// it is dropped, or when `flush` is manually invoked on the guard.
183    pub fn flush_on_drop<W>(&self, writer: W) -> FlushGuard<W>
184    where
185        W: Write + 'static,
186    {
187        FlushGuard {
188            graph: self.graph.clone(),
189            writer,
190        }
191    }
192}
193
194impl GraphLayer {
195    /// Constructs a `GraphLayer` that constructs the call graph, and a
196    /// `FlushGuard` which writes the graph to a `dot` file when dropped.
197    pub fn with_file(path: impl AsRef<Path>) -> Result<(Self, FlushGuard<BufWriter<File>>), Error> {
198        let path = path.as_ref();
199        let file = File::create(path)
200            .map_err(|source| Kind::CreateFile {
201                path: path.into(),
202                source,
203            })
204            .map_err(Error)?;
205        let writer = BufWriter::new(file);
206        let layer = Self::new();
207        let guard = layer.flush_on_drop(writer);
208        Ok((layer, guard))
209    }
210}
211
212impl<S> Layer<S> for GraphLayer
213where
214    S: Subscriber + for<'span> LookupSpan<'span>,
215{
216    fn on_enter(&self, id: &span::Id, ctx: Context<'_, S>) {
217        let mut locked = self.graph.lock().unwrap();
218
219        // Add node
220        let first = ctx.span(id).expect("expected: span id exists in registry");
221        let node_b = first.name();
222        locked.add_node(node_b);
223
224        // Find parent node
225        let node_a = if let Some(parent) = first.parent() {
226            parent.name()
227        } else if let Some(name) = self.top_node {
228            name
229        } else {
230            return;
231        };
232
233        if let Some(weight) = locked.edge_weight_mut(node_a, node_b) {
234            // Increase edge weight
235            *weight += 1;
236        } else {
237            // Add edge
238            locked.add_edge(node_a, node_b, 1);
239        }
240    }
241}