Reduce lock depth.

Author: cjgillot

compiler/rustc_query_system/src/dep_graph/graph.rs

@@ -3,7 +3,6 @@ use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::profiling::{EventId, QueryInvocationId, SelfProfilerRef};
 use rustc_data_structures::sharded::{self, Sharded};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
-use rustc_data_structures::steal::Steal;
 use rustc_data_structures::sync::{AtomicU32, AtomicU64, Lock, Lrc};
 use rustc_data_structures::unord::UnordMap;
 use rustc_index::IndexVec;
@@ -195,8 +194,11 @@ impl DepGraph {
     }
 
     pub fn with_query(&self, f: impl Fn(&DepGraphQuery)) {
-        if let Some(data) = &self.data {
-            data.current.encoder.borrow().with_query(f)
+        if let Some(data) = &self.data
+            && let Some(record_graph) = &data.current.record_graph
+        {
+            let record_graph = record_graph.lock();
+            f(&*record_graph)
         }
     }
 
@@ -984,7 +986,7 @@ impl DepGraph {
 
     pub fn finish_encoding(&self, profiler: &SelfProfilerRef) -> FileEncodeResult {
         if let Some(data) = &self.data {
-            data.current.encoder.steal().finish(profiler)
+            data.current.encoder.lock().finish(profiler)
         } else {
             Ok(0)
         }
@@ -1070,7 +1072,8 @@ rustc_index::newtype_index! {
 /// manipulating both, we acquire `new_node_to_index` or `prev_index_to_index`
 /// first, and `data` second.
 pub(super) struct CurrentDepGraph {
-    encoder: Steal<GraphEncoder>,
+    encoder: Lock<GraphEncoder>,
+    record_graph: Option<Lock<DepGraphQuery>>,
     new_node_to_index: Sharded<FxHashMap<DepNode, DepNodeIndex>>,
     prev_index_to_index: Lock<IndexVec<SerializedDepNodeIndex, Option<DepNodeIndex>>>,
 
@@ -1146,12 +1149,9 @@ impl CurrentDepGraph {
             .map(EventId::from_label);
 
         CurrentDepGraph {
-            encoder: Steal::new(GraphEncoder::new::<D>(
-                encoder,
-                prev_graph_node_count,
-                record_graph,
-                record_stats,
-            )),
+            encoder: Lock::new(GraphEncoder::new::<D>(encoder, record_stats)),
+            record_graph: record_graph
+                .then(|| Lock::new(DepGraphQuery::new(prev_graph_node_count))),
             new_node_to_index: Sharded::new(|| {
                 FxHashMap::with_capacity_and_hasher(
                     new_node_count_estimate / sharded::shards(),
@@ -1192,8 +1192,13 @@ impl CurrentDepGraph {
         let dep_node_index = match self.new_node_to_index.lock_shard_by_value(&key).entry(key) {
             Entry::Occupied(entry) => *entry.get(),
             Entry::Vacant(entry) => {
-                let dep_node_index =
-                    self.encoder.borrow().send(profiler, key, current_fingerprint, edges);
+                let dep_node_index = self.encoder.lock().send(
+                    profiler,
+                    key,
+                    current_fingerprint,
+                    edges,
+                    &self.record_graph,
+                );
                 entry.insert(dep_node_index);
                 dep_node_index
             }
@@ -1224,8 +1229,13 @@ impl CurrentDepGraph {
                 let dep_node_index = match prev_index_to_index[prev_index] {
                     Some(dep_node_index) => dep_node_index,
                     None => {
-                        let dep_node_index =
-                            self.encoder.borrow().send(profiler, key, fingerprint, edges);
+                        let dep_node_index = self.encoder.lock().send(
+                            profiler,
+                            key,
+                            fingerprint,
+                            edges,
+                            &self.record_graph,
+                        );
                         prev_index_to_index[prev_index] = Some(dep_node_index);
                         dep_node_index
                     }
@@ -1287,7 +1297,8 @@ impl CurrentDepGraph {
                     .map(|i| prev_index_to_index[i].unwrap())
                     .collect();
                 let fingerprint = prev_graph.fingerprint_by_index(prev_index);
-                let dep_node_index = self.encoder.borrow().send(profiler, key, fingerprint, edges);
+                let dep_node_index =
+                    self.encoder.lock().send(profiler, key, fingerprint, edges, &self.record_graph);
                 prev_index_to_index[prev_index] = Some(dep_node_index);
                 #[cfg(debug_assertions)]
                 self.record_edge(dep_node_index, key, fingerprint);

compiler/rustc_query_system/src/dep_graph/serialized.rs

@@ -482,9 +482,9 @@ impl<D: Deps> NodeEncoder for EncoderState<D> {
         let edge_count = node.edges.len();
         self.total_edge_count += edge_count;
 
-        if let Some(record_graph) = &record_graph {
+        if let Some(record_graph) = record_graph {
             // Do not ICE when a query is called from within `with_query`.
-            if let Some(record_graph) = &mut record_graph.try_lock() {
+            if let Some(mut record_graph) = record_graph.try_lock() {
                 record_graph.push(index, node.node, &node.edges);
             }
         }
@@ -537,36 +537,21 @@ impl<D: Deps> NodeEncoder for EncoderState<D> {
 }
 
 pub struct GraphEncoder {
-    status: Lock<Box<dyn NodeEncoder>>,
-    record_graph: Option<Lock<DepGraphQuery>>,
+    status: Box<dyn NodeEncoder>,
 }
 
 impl GraphEncoder {
-    pub fn new<D: Deps>(
-        encoder: FileEncoder,
-        prev_node_count: usize,
-        record_graph: bool,
-        record_stats: bool,
-    ) -> Self {
-        let record_graph = record_graph.then(|| Lock::new(DepGraphQuery::new(prev_node_count)));
+    pub fn new<D: Deps>(encoder: FileEncoder, record_stats: bool) -> Self {
         let status = Box::new(EncoderState::<D>::new(encoder, record_stats)) as _;
-        let status = Lock::new(status);
-        GraphEncoder { status, record_graph }
-    }
-
-    pub(crate) fn with_query(&self, f: impl Fn(&DepGraphQuery)) {
-        if let Some(record_graph) = &self.record_graph {
-            f(&record_graph.lock())
-        }
+        GraphEncoder { status }
     }
 
     pub(crate) fn print_incremental_info(
         &self,
         total_read_count: u64,
         total_duplicate_read_count: u64,
     ) {
-        let status = self.status.lock();
-        if let Some((total_node_count, total_edge_count, record_stats)) = status.get_stats() {
+        if let Some((total_node_count, total_edge_count, record_stats)) = self.status.get_stats() {
             let mut stats: Vec<_> = record_stats.values().collect();
             stats.sort_by_key(|s| -(s.node_counter as i64));
 
@@ -613,19 +598,20 @@ impl GraphEncoder {
     }
 
     pub(crate) fn send(
-        &self,
+        &mut self,
         profiler: &SelfProfilerRef,
         node: DepNode,
         fingerprint: Fingerprint,
         edges: EdgesVec,
+        record_graph: &Option<Lock<DepGraphQuery>>,
     ) -> DepNodeIndex {
         let _prof_timer = profiler.generic_activity("incr_comp_encode_dep_graph");
         let node = NodeInfo { node, fingerprint, edges };
-        self.status.lock().encode_node(&node, &self.record_graph)
+        self.status.encode_node(&node, record_graph)
     }
 
-    pub fn finish(self, profiler: &SelfProfilerRef) -> FileEncodeResult {
+    pub fn finish(&mut self, profiler: &SelfProfilerRef) -> FileEncodeResult {
         let _prof_timer = profiler.generic_activity("incr_comp_encode_dep_graph");
-        self.status.into_inner().finish(profiler)
+        self.status.finish(profiler)
     }
 }