Capture backtraces of allocations for debugging

Largely based on https://github.com/bertptrs/tracing-mutex/pull/28 with
only minor modifications.

src/graph.rs

@@ -1,4 +1,5 @@
 use std::cell::Cell;
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::hash::Hash;
@@ -19,24 +20,24 @@ type Order = usize;
 /// visibly changed.
 ///
 /// [paper]: https://whileydave.com/publications/pk07_jea/
-#[derive(Default, Debug)]
-pub struct DiGraph<V>
+#[derive(Debug)]
+pub struct DiGraph<V, E>
 where
     V: Eq + Hash + Copy,
 {
-    nodes: HashMap<V, Node<V>>,
+    nodes: HashMap<V, Node<V, E>>,
     // Instead of reordering the orders in the graph whenever a node is deleted, we maintain a list
     // of unused ids that can be handed out later again.
     unused_order: Vec<Order>,
 }
 
 #[derive(Debug)]
-struct Node<V>
+struct Node<V, E>
 where
     V: Eq + Hash + Clone,
 {
     in_edges: HashSet<V>,
-    out_edges: HashSet<V>,
+    out_edges: HashMap<V, E>,
     // The "Ord" field is a Cell to ensure we can update it in an immutable context.
     // `std::collections::HashMap` doesn't let you have multiple mutable references to elements, but
     // this way we can use immutable references and still update `ord`. This saves quite a few
@@ -44,7 +45,7 @@ where
     ord: Cell<Order>,
 }
 
-impl<V> DiGraph<V>
+impl<V, E> DiGraph<V, E>
 where
     V: Eq + Hash + Copy,
 {
@@ -55,7 +56,7 @@ where
     /// the node in the topological order.
     ///
     /// New nodes are appended to the end of the topological order when added.
-    fn add_node(&mut self, n: V) -> (&mut HashSet<V>, &mut HashSet<V>, Order) {
+    fn add_node(&mut self, n: V) -> (&mut HashSet<V>, &mut HashMap<V, E>, Order) {
         // need to compute next id before the call to entry() to avoid duplicate borrow of nodes
         let fallback_id = self.nodes.len();
 
@@ -89,7 +90,7 @@ where
                 // Return ordering to the pool of unused ones
                 self.unused_order.push(ord.get());
 
-                out_edges.into_iter().for_each(|m| {
+                out_edges.into_keys().for_each(|m| {
                     self.nodes.get_mut(&m).unwrap().in_edges.remove(&n);
                 });
 
@@ -106,7 +107,7 @@ where
     ///
     /// Nodes, both from and to, are created as needed when creating new edges. If the new edge
     /// would introduce a cycle, the edge is rejected and `false` is returned.
-    pub(crate) fn add_edge(&mut self, x: V, y: V) -> bool {
+    pub(crate) fn add_edge(&mut self, x: V, y: V, e: impl FnOnce() -> E) -> bool {
         if x == y {
             // self-edges are always considered cycles
             return false;
@@ -114,10 +115,13 @@ where
 
         let (_, out_edges, ub) = self.add_node(x);
 
-        if !out_edges.insert(y) {
+        match out_edges.entry(y) {
+            Entry::Occupied(_) => {
                 // Edge already exists, nothing to be done
                 return true;
             }
+            Entry::Vacant(entry) => entry.insert(e()),
+        };
 
         let (in_edges, _, lb) = self.add_node(y);
 
@@ -157,14 +161,14 @@ where
     /// Forwards depth-first-search
     fn dfs_f<'a>(
         &'a self,
-        n: &'a Node<V>,
+        n: &'a Node<V, E>,
         ub: Order,
         visited: &mut HashSet<V>,
-        delta_f: &mut Vec<&'a Node<V>>,
+        delta_f: &mut Vec<&'a Node<V, E>>,
     ) -> bool {
         delta_f.push(n);
 
-        n.out_edges.iter().all(|w| {
+        n.out_edges.keys().all(|w| {
             let node = &self.nodes[w];
             let ord = node.ord.get();
 
@@ -185,10 +189,10 @@ where
     /// Backwards depth-first-search
     fn dfs_b<'a>(
         &'a self,
-        n: &'a Node<V>,
+        n: &'a Node<V, E>,
         lb: Order,
         visited: &mut HashSet<V>,
-        delta_b: &mut Vec<&'a Node<V>>,
+        delta_b: &mut Vec<&'a Node<V, E>>,
     ) {
         delta_b.push(n);
 
@@ -202,7 +206,7 @@ where
         }
     }
 
-    fn reorder(&self, mut delta_f: Vec<&Node<V>>, mut delta_b: Vec<&Node<V>>) {
+    fn reorder(&self, mut delta_f: Vec<&Node<V, E>>, mut delta_b: Vec<&Node<V, E>>) {
         self.sort(&mut delta_f);
         self.sort(&mut delta_b);
 
@@ -223,12 +227,25 @@ where
         }
     }
 
-    fn sort(&self, ids: &mut [&Node<V>]) {
+    fn sort(&self, ids: &mut [&Node<V, E>]) {
         // Can use unstable sort because mutex ids should not be equal
         ids.sort_unstable_by_key(|v| &v.ord);
     }
 }
 
+// Manual `Default` impl as derive causes unnecessarily strong bounds.
+impl<V, E> Default for DiGraph<V, E>
+where
+    V: Eq + Hash + Copy,
+{
+    fn default() -> Self {
+        Self {
+            nodes: Default::default(),
+            unused_order: Default::default(),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use rand::seq::SliceRandom;
@@ -236,12 +253,14 @@ mod tests {
 
     use super::*;
 
+    fn nop() {}
+
     #[test]
     fn test_no_self_cycle() {
         // Regression test for https://github.com/bertptrs/tracing-mutex/issues/7
         let mut graph = DiGraph::default();
 
-        assert!(!graph.add_edge(1, 1));
+        assert!(!graph.add_edge(1, 1, nop));
     }
 
     #[test]
@@ -249,16 +268,16 @@ mod tests {
         let mut graph = DiGraph::default();
 
         // Add some safe edges
-        assert!(graph.add_edge(0, 1));
-        assert!(graph.add_edge(1, 2));
-        assert!(graph.add_edge(2, 3));
-        assert!(graph.add_edge(4, 2));
+        assert!(graph.add_edge(0, 1, nop));
+        assert!(graph.add_edge(1, 2, nop));
+        assert!(graph.add_edge(2, 3, nop));
+        assert!(graph.add_edge(4, 2, nop));
 
         // Try to add an edge that introduces a cycle
-        assert!(!graph.add_edge(3, 1));
+        assert!(!graph.add_edge(3, 1, nop));
 
         // Add an edge that should reorder 0 to be after 4
-        assert!(graph.add_edge(4, 0));
+        assert!(graph.add_edge(4, 0, nop));
     }
 
     /// Fuzz the DiGraph implementation by adding a bunch of valid edges.
@@ -287,7 +306,7 @@ mod tests {
         let mut graph = DiGraph::default();
 
         for (x, y) in edges {
-            assert!(graph.add_edge(x, y));
+            assert!(graph.add_edge(x, y, nop));
         }
     }
 }

src/lib.rs

@@ -46,6 +46,7 @@
 //!
 //! [paper]: https://whileydave.com/publications/pk07_jea/
 #![cfg_attr(docsrs, feature(doc_cfg))]
+use std::backtrace::Backtrace;
 use std::cell::RefCell;
 use std::fmt;
 use std::marker::PhantomData;
@@ -53,6 +54,7 @@ use std::ops::Deref;
 use std::ops::DerefMut;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
+use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::MutexGuard;
 use std::sync::OnceLock;
@@ -141,7 +143,7 @@ impl MutexId {
             if let Some(&previous) = locks.borrow().last() {
                 let mut graph = get_dependency_graph();
 
-                !graph.add_edge(previous, self.value())
+                !graph.add_edge(previous, self.value(), MutexDep::capture)
             } else {
                 false
             }
@@ -259,9 +261,18 @@ impl<'a> Drop for BorrowedMutex<'a> {
     }
 }
 
+#[derive(Clone)]
+struct MutexDep(Arc<Backtrace>);
+
+impl MutexDep {
+    pub fn capture() -> Self {
+        Self(Arc::new(Backtrace::capture()))
+    }
+}
+
 /// Get a reference to the current dependency graph
-fn get_dependency_graph() -> impl DerefMut<Target = DiGraph<usize>> {
-    static DEPENDENCY_GRAPH: OnceLock<Mutex<DiGraph<usize>>> = OnceLock::new();
+fn get_dependency_graph() -> impl DerefMut<Target = DiGraph<usize, MutexDep>> {
+    static DEPENDENCY_GRAPH: OnceLock<Mutex<DiGraph<usize, MutexDep>>> = OnceLock::new();
 
     DEPENDENCY_GRAPH
         .get_or_init(Default::default)
@@ -292,11 +303,11 @@ mod tests {
         let c = LazyMutexId::new();
 
         let mut graph = get_dependency_graph();
-        assert!(graph.add_edge(a.value(), b.value()));
-        assert!(graph.add_edge(b.value(), c.value()));
+        assert!(graph.add_edge(a.value(), b.value(), MutexDep::capture));
+        assert!(graph.add_edge(b.value(), c.value(), MutexDep::capture));
 
         // Creating an edge c → a should fail as it introduces a cycle.
-        assert!(!graph.add_edge(c.value(), a.value()));
+        assert!(!graph.add_edge(c.value(), a.value(), MutexDep::capture));
 
         // Drop graph handle so we can drop vertices without deadlocking
         drop(graph);
@@ -304,7 +315,7 @@ mod tests {
         drop(b);
 
         // If b's destructor correctly ran correctly we can now add an edge from c to a.
-        assert!(get_dependency_graph().add_edge(c.value(), a.value()));
+        assert!(get_dependency_graph().add_edge(c.value(), a.value(), MutexDep::capture));
     }
 
     /// Test creating a cycle, then panicking.