time_graph/
graph.rs

1use std::sync::Mutex;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::collections::BTreeMap;
4use std::time::Duration;
5use std::cell::RefCell;
6
7use once_cell::sync::Lazy;
8use quanta::Clock;
9use petgraph::graph::{Graph, NodeIndex};
10
11use crate::{CallSite, CallSiteId};
12
13/// Global clock to record start/end times
14static CLOCK: Lazy<Clock> = Lazy::new(Clock::new);
15
16/// Global call graph, including recorded timings and calls count
17static CALL_GRAPH: Lazy<Mutex<LightCallGraph>> = Lazy::new(|| {
18    Mutex::new(LightCallGraph::new())
19});
20
21/// Should we collect data?
22static COLLECTION_ENABLED: AtomicBool = AtomicBool::new(false);
23
24thread_local! {
25    /// For each thread, which span is currently executing? This will become the
26    /// parent of new spans.
27    pub static LOCAL_CURRENT_SPAN: RefCell<Option<CallSiteId>> = const { RefCell::new(None) };
28}
29
30/// A [`Span`] records a single execution of code associated with a
31/// [`CallSite`].
32///
33/// This is not usually constructed manually but with either the
34/// [`macro@spanned`] or [`instrument`](attr.instrument.html) macros.
35pub struct Span {
36    callsite: &'static CallSite,
37}
38
39impl Span {
40    /// Create a new [`Span`] associated with the given `callsite`.
41    pub fn new(callsite: &'static CallSite) -> Span {
42        Span {
43            callsite: callsite,
44        }
45    }
46
47    /// Enter the span, the span will automatically be exited when the
48    /// [`SpanGuard`] is dropped.
49    #[must_use]
50    pub fn enter(&self) -> SpanGuard<'_> {
51        if !COLLECTION_ENABLED.load(Ordering::Acquire) {
52            return SpanGuard {
53                span: self,
54                parent: None,
55                start: 0,
56            };
57        }
58
59        let id = self.callsite.id();
60        let parent = LOCAL_CURRENT_SPAN.with(|parent| {
61            let mut parent = parent.borrow_mut();
62
63            let previous = *parent;
64            *parent = Some(id);
65            return previous;
66        });
67
68        SpanGuard {
69            span: self,
70            parent: parent,
71            start: CLOCK.raw(),
72        }
73    }
74}
75
76/// When a [`SpanGuard`] is dropped, it saves the execution time of the
77/// corresponding span in the global call graph.
78pub struct SpanGuard<'a> {
79    span: &'a Span,
80    parent: Option<CallSiteId>,
81    start: u64,
82}
83
84impl<'a> Drop for SpanGuard<'a>  {
85    fn drop(&mut self) {
86        if !COLLECTION_ENABLED.load(Ordering::Acquire) {
87            return;
88        }
89        let elapsed = CLOCK.delta(self.start, CLOCK.raw());
90
91        LOCAL_CURRENT_SPAN.with(|parent| {
92            let mut parent = parent.borrow_mut();
93            *parent = self.parent;
94        });
95
96
97        let mut graph = CALL_GRAPH.lock().expect("poisoned mutex");
98        let callsite = self.span.callsite.id();
99        graph.add_node(callsite);
100        graph.increase_timing(callsite, elapsed);
101
102        if let Some(parent) = self.parent {
103            graph.add_node(parent);
104            graph.increase_call_count(parent, callsite);
105        }
106    }
107}
108
109/// Call graph node identifying their call site with its `CallSiteId`.
110struct LightGraphNode {
111    callsite: CallSiteId,
112    elapsed: Duration,
113    called: u32,
114}
115
116impl LightGraphNode {
117    fn new(callsite: CallSiteId) -> LightGraphNode {
118        LightGraphNode {
119            callsite: callsite,
120            elapsed: Duration::new(0, 0),
121            called: 0,
122        }
123    }
124}
125
126/// Simple Call graph, identifying call site with their `CallSiteId`.
127///
128/// The graph nodes are spans with associated timings, while the edges represent
129/// the number of calls from one node to the other.
130struct LightCallGraph {
131    graph: Graph<LightGraphNode, usize>
132}
133
134impl LightCallGraph {
135    fn new() -> LightCallGraph {
136        LightCallGraph {
137            graph: Graph::new(),
138        }
139    }
140
141    pub fn clear(&mut self) {
142        self.graph.clear()
143    }
144
145    /// Find a node in the graph with its `CallSiteId`.
146    fn find(&mut self, callsite: CallSiteId) -> Option<NodeIndex> {
147        self.graph.node_indices().find(|&id| self.graph[id].callsite == callsite)
148    }
149
150    /// Add a node for the given callsite to the graph, do nothing if there is
151    /// already such a node
152    pub fn add_node(&mut self, callsite: CallSiteId) {
153        if self.find(callsite).is_none() {
154            self.graph.add_node(LightGraphNode::new(callsite));
155        }
156    }
157
158    /// Increase the number of time the `parent` span called the `child` span
159    /// by one.
160    pub fn increase_call_count(&mut self, parent: CallSiteId, child: CallSiteId) {
161        let parent = self.find(parent).expect("missing node for parent");
162        let child = self.find(child).expect("missing node for child");
163        if let Some(edge) = self.graph.find_edge(parent, child) {
164            let count = self
165                .graph
166                .edge_weight_mut(edge)
167                .expect("failed to get edge weights");
168            *count += 1;
169        } else {
170            // initialize edge count to 1
171            self.graph.add_edge(parent, child, 1);
172        }
173    }
174
175    /// Increase the timing associated with a span by `time`, and the number of
176    /// time this span has been called by one.
177    pub fn increase_timing(&mut self, span: CallSiteId, time: Duration) {
178        let id = self.find(span).expect("missing node");
179        self.graph[id].elapsed += time;
180        self.graph[id].called += 1;
181    }
182}
183
184/// Clear the global call graph from all data
185pub fn clear_collected_data() {
186    CALL_GRAPH.lock().expect("poisoned mutex").clear();
187}
188
189/// Enable/disable data collection
190pub fn enable_data_collection(enabled: bool) {
191    COLLECTION_ENABLED.store(enabled, Ordering::Release);
192}
193
194/// Get a copy of the call graph as currently known
195pub fn get_full_graph() -> FullCallGraph {
196    let graph = CALL_GRAPH.lock().expect("poisoned mutex");
197
198    let mut all_callsites = BTreeMap::new();
199    crate::traverse_registered_callsite(|callsite| {
200        all_callsites.insert(callsite.id(), callsite);
201    });
202
203    let graph = graph.graph.map(|index, node| {
204        TimedSpan::new(node, index.index(), all_callsites[&node.callsite])
205    }, |_, &edge| edge);
206
207    return FullCallGraph {
208        graph: graph
209    };
210}
211
212/// [`TimedSpan`] contains all data related to a single function or span inside
213/// the global call graph.
214pub struct TimedSpan {
215    /// Unique identifier of this function/span in the call graph
216    pub id: usize,
217    /// [`CallSite`] associated with this function/span
218    pub callsite: &'static CallSite,
219    /// Total elapsed time inside this function/span
220    pub elapsed: Duration,
221    /// Number of times this function/span have been called
222    pub called: u32,
223}
224
225impl TimedSpan {
226    fn new(node: &LightGraphNode, id: usize, callsite: &'static CallSite) -> TimedSpan {
227        TimedSpan {
228            id: id,
229            callsite: callsite,
230            elapsed: node.elapsed,
231            called: node.called,
232        }
233    }
234}
235
236impl std::fmt::Display for TimedSpan {
237    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
238        write!(
239            f,
240            "{} ran for {:?}, called {} times",
241            self.callsite.full_name(), self.elapsed, self.called
242        )
243    }
244}
245
246/// Full call graph including execution time and number of calls between
247/// functions/spans.
248///
249/// This graph is a directed graph linking different `SpanTiming` by the
250/// number of time a given span was the child of another one.
251///
252/// # Examples
253///
254/// Code that looks like this
255/// ```no_run
256/// #[time_graph::instrument]
257/// fn start() {
258///     inside();
259///     inside();
260///     inner();
261/// }
262///
263/// #[time_graph::instrument]
264/// fn inside() {
265///    inner();
266/// }
267///
268/// #[time_graph::instrument]
269/// fn inner() {
270///     // do stuff
271/// }
272/// ```
273///
274/// Will result in a graph like this, where the number near the edge correspond
275/// to the number of time a given span called another one.
276/// ```bash no_run
277///             | start, called 1 |
278///                /           |
279///              /  2          |
280///            /               |  1
281///   | inside, called 2 |     |
282///                 \          |
283///                 2 \        |
284///                     \      |
285///                  | inner, called 3 |
286/// ```
287pub struct FullCallGraph {
288    graph: Graph<TimedSpan, usize>
289}
290
291/// A set of calls from one function/span to another
292pub struct Calls {
293    /// the outer/calling function/span
294    pub caller: usize,
295    /// the inner/called function/span
296    pub callee: usize,
297    /// number of time the inner function/span have been called by the outer one
298    pub count: usize,
299}
300
301impl FullCallGraph {
302    /// Get the full list of spans/functions known by this graph
303    pub fn spans(&self) -> impl Iterator<Item = &TimedSpan> {
304        self.graph.raw_nodes().iter().map(|node| &node.weight)
305    }
306
307    /// Get the list of calls between spans in this graph
308    pub fn calls(&self) -> impl Iterator<Item = Calls> + '_ {
309        self.graph.raw_edges().iter().map(|edge| Calls {
310            caller: edge.target().index(),
311            callee: edge.source().index(),
312            count: edge.weight,
313        })
314    }
315
316    /// Get the full graph in [graphviz](https://graphviz.org/) dot format.
317    ///
318    /// The exact output is unstable and should not be relied on.
319    pub fn as_dot(&self) -> String {
320        petgraph::dot::Dot::new(&self.graph).to_string()
321    }
322
323    /// Get a per span summary table of this graph.
324    ///
325    /// The exact output is unstable and should not be relied on.
326    ///
327    /// This function is only available if the `"table"` cargo feature is enabled
328    ///
329    /// # Panic
330    ///
331    /// This function will panic if the graph is cyclical, i.e. if two or more
332    /// span are mutually recursive.
333    #[cfg(feature = "table")]
334    pub fn as_table(&self) -> String {
335        self.as_table_impl(false)
336    }
337
338    /// Same as `as_table`, but using the short names of the spans instead of
339    /// the full name.
340    #[cfg(feature = "table")]
341    pub fn as_short_table(&self) -> String {
342        self.as_table_impl(true)
343    }
344
345    #[cfg(feature = "table")]
346    fn as_table_impl(&self, short_names: bool) -> String {
347        use petgraph::Direction;
348
349        use crate::table::Row;
350        use crate::table::TableCell;
351
352        let mut names = BTreeMap::new();
353        for node in self.graph.node_weights() {
354            if short_names {
355                names.insert(node.id, node.callsite.name().to_string());
356            } else {
357                names.insert(node.id, node.callsite.full_name());
358            }
359        }
360
361        let mut table = crate::table::Table::new();
362        table.style = crate::table::TableStyle::extended();
363
364        table.add_row(Row::new(vec![
365            "id",
366            // pad "span name" to make the table look nicer with short names
367            "span name                                   ",
368            "call count",
369            "called by",
370            "total",
371            "mean",
372        ]));
373
374        for &node_id in petgraph::algo::kosaraju_scc(&self.graph)
375            .iter()
376            .rev()
377            .flatten()
378        {
379            let node = &self.graph[node_id];
380
381            let mut called_by = vec![];
382            for other in self.graph.neighbors_directed(node_id, Direction::Incoming) {
383                called_by.push(self.graph[other].id.to_string());
384            }
385            let called_by = if !called_by.is_empty() {
386                called_by.join(", ")
387            } else {
388                "—".into()
389            };
390
391            let mean = node.elapsed / node.called;
392            let warn = if mean < Duration::from_nanos(1500) { " ⚠️ " } else { "" };
393
394            table.add_row(Row::new(vec![
395                TableCell::new_right_aligned(node.id),
396                TableCell::new(&names[&node.id]),
397                TableCell::new_right_aligned(node.called),
398                TableCell::new_right_aligned(called_by),
399                TableCell::new_right_aligned(format!("{:.2?}", node.elapsed)),
400                TableCell::new_right_aligned(format!("{:.2?}{}", mean, warn)),
401            ]));
402        }
403
404        return table.render();
405    }
406
407    /// Get all the data in this graph in JSON.
408    ///
409    /// The exact output is unstable and should not be relied on.
410    ///
411    /// This function is only available if the `"json"` cargo feature is enabled
412    #[cfg(feature = "json")]
413    pub fn as_json(&self) -> String {
414        let mut spans = json::JsonValue::new_object();
415        for span in self.spans() {
416            spans[&span.callsite.full_name()] = json::object! {
417                "id" => span.id,
418                "elapsed" => format!("{:?}", span.elapsed),
419                "called" => span.called,
420            };
421        }
422
423        let mut all_calls = json::JsonValue::new_array();
424        for call in self.calls() {
425            all_calls.push(json::object! {
426                "caller" => call.caller,
427                "callee" => call.callee,
428                "count" => call.count,
429            }).expect("failed to add edge information to JSON");
430        }
431
432        return json::stringify(json::object! {
433            "timings" => spans,
434            "calls" => all_calls,
435        });
436    }
437}
time_graph/graph.rs

time_graph/
graph.rs
