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

time_graph/
graph.rs
