1#![allow(clippy::needless_doctest_main)]
2#![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#[derive(Clone, Debug)]
94pub struct GraphLayer {
95 graph: Arc<Mutex<CallGraph>>,
96 top_node: Option<&'static str>,
97}
98
99impl GraphLayer {
100 #[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 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#[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 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 pub fn new() -> Self {
178 Default::default()
179 }
180
181 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 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 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 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 *weight += 1;
236 } else {
237 locked.add_edge(node_a, node_b, 1);
239 }
240 }
241}