Streamline active job collection.

In `gather_active_jobs`:
- Remove the `require_complete` parameter. We now always use
  `try_lock_shards` and let the caller abort on incompleteness
- Remove `active`. The comment about `make_frame` is out of date, and
  `create_deferred_query_stack_frame` *is* safe to call with the query
  state locked.
- Don't stop on the first failed shard; continue to get as much of the
  job map as possible.

In `collect_active_jobs_from_all_queries`:
- Remove the  `require_complete` parameter.
- Change the return type. Having the same thing in both the `Ok` and
  `Err` cases was strange.

Author: nnethercote

compiler/rustc_interface/src/util.rs

@@ -254,9 +254,13 @@ internal compiler error: query cycle handler thread panicked, aborting process";
                                     || {
                                         // Ensure there were no errors collecting all active jobs.
                                         // We need the complete map to ensure we find a cycle to break.
-                                        collect_active_jobs_from_all_queries(tcx, false).expect(
+                                        let (job_map, collect_ok) =
+                                            collect_active_jobs_from_all_queries(tcx);
+                                        assert!(
+                                            collect_ok,
                                             "failed to collect active queries in deadlock handler",
-                                        )
+                                        );
+                                        job_map
                                     },
                                 );
                                 break_query_cycles(job_map, &registry);

compiler/rustc_query_impl/src/execution.rs

@@ -1,7 +1,7 @@
 use std::hash::Hash;
 use std::mem;
 
-use rustc_data_structures::hash_table::{Entry, HashTable};
+use rustc_data_structures::hash_table::Entry;
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_data_structures::sync::{DynSend, DynSync};
 use rustc_data_structures::{outline, sharded, sync};
@@ -19,7 +19,9 @@ use rustc_span::{DUMMY_SP, Span};
 use crate::collect_active_jobs_from_all_queries;
 use crate::dep_graph::{DepNode, DepNodeIndex};
 use crate::job::{QueryJobInfo, QueryJobMap, find_cycle_in_stack, report_cycle};
-use crate::plumbing::{current_query_job, next_job_id, start_query};
+use crate::plumbing::{
+    create_deferred_query_stack_frame, current_query_job, next_job_id, start_query,
+};
 
 #[inline]
 fn equivalent_key<K: Eq, V>(k: K) -> impl Fn(&(K, V)) -> bool {
@@ -42,6 +44,8 @@ pub(crate) fn all_inactive<'tcx, K>(state: &QueryState<'tcx, K>) -> bool {
 }
 
 /// Internal plumbing for collecting the set of active jobs for this query.
+/// Returns `false` if we failed to collect all jobs, which can happen if some
+/// query cache shards could not be locked.
 ///
 /// Should only be called from `collect_active_jobs_from_all_queries`.
 ///
@@ -51,57 +55,36 @@ pub(crate) fn all_inactive<'tcx, K>(state: &QueryState<'tcx, K>) -> bool {
 pub(crate) fn gather_active_jobs<'tcx, C>(
     query: &'tcx QueryVTable<'tcx, C>,
     tcx: TyCtxt<'tcx>,
-    require_complete: bool,
     job_map_out: &mut QueryJobMap<'tcx>, // Out-param; job info is gathered into this map
-) -> Option<()>
+) -> bool
 where
     C: QueryCache<Key: QueryKey + DynSend + DynSync>,
     QueryVTable<'tcx, C>: DynSync,
 {
-    let mut active = Vec::new();
-
-    // Helper to gather active jobs from a single shard.
-    let mut gather_shard_jobs = |shard: &HashTable<(C::Key, ActiveKeyStatus<'tcx>)>| {
-        for (k, v) in shard.iter() {
-            if let ActiveKeyStatus::Started(ref job) = *v {
-                active.push((*k, job.clone()));
-            }
-        }
-    };
-
-    // Lock shards and gather jobs from each shard.
-    if require_complete {
-        for shard in query.state.active.lock_shards() {
-            gather_shard_jobs(&shard);
-        }
-    } else {
-        // We use try_lock_shards here since we are called from the
-        // deadlock handler, and this shouldn't be locked.
-        for shard in query.state.active.try_lock_shards() {
-            // This can be called during unwinding, and the function has a `try_`-prefix, so
-            // don't `unwrap()` here, just manually check for `None` and do best-effort error
-            // reporting.
-            match shard {
-                None => {
-                    tracing::warn!(
-                        "Failed to collect active jobs for query with name `{}`!",
-                        query.name
-                    );
-                    return None;
+    let mut ok = true;
+
+    for shard in query.state.active.try_lock_shards() {
+        match shard {
+            Some(shard) => {
+                for (key, status) in shard.iter() {
+                    if let ActiveKeyStatus::Started(job) = status {
+                        // This function to call with query state locked because it is very simple.
+                        let frame = create_deferred_query_stack_frame(tcx, query, *key);
+                        job_map_out.insert(job.id, QueryJobInfo { frame, job: job.clone() });
+                    }
                 }
-                Some(shard) => gather_shard_jobs(&shard),
+            }
+            None => {
+                tracing::warn!(
+                    "Failed to collect active jobs for query with name `{}`!",
+                    query.name
+                );
+                ok = false;
             }
         }
     }
 
-    // Call `make_frame` while we're not holding a `state.active` lock as `make_frame` may call
-    // queries leading to a deadlock.
-    for (key, job) in active {
-        let frame = crate::plumbing::create_deferred_query_stack_frame(tcx, query, key);
-        job_map_out.insert(job.id, QueryJobInfo { frame, job });
-    }
-
-    Some(())
+    ok
 }
 
 /// Guard object representing the responsibility to execute a query job and
@@ -220,9 +203,8 @@ fn cycle_error<'tcx, C: QueryCache>(
 ) -> (C::Value, Option<DepNodeIndex>) {
     // Ensure there was no errors collecting all active jobs.
     // We need the complete map to ensure we find a cycle to break.
-    let job_map = collect_active_jobs_from_all_queries(tcx, false)
-        .ok()
-        .expect("failed to collect active queries");
+    let (job_map, collect_ok) = collect_active_jobs_from_all_queries(tcx);
+    assert!(collect_ok, "failed to collect active queries for cycle error");
 
     let error = find_cycle_in_stack(try_execute, job_map, &current_query_job(), span);
     (mk_cycle(query, tcx, error.lift()), None)

compiler/rustc_query_impl/src/job.rs

@@ -384,9 +384,8 @@ pub fn print_query_stack<'tcx>(
     let mut count_printed = 0;
     let mut count_total = 0;
 
-    // Make use of a partial query job map if we fail to take locks collecting active queries.
-    let job_map: QueryJobMap<'_> = collect_active_jobs_from_all_queries(tcx, false)
-        .unwrap_or_else(|partial_job_map| partial_job_map);
+    // Make use of a partial query job map if we fail to take some locks collecting active queries.
+    let (job_map, _collect_ok) = collect_active_jobs_from_all_queries(tcx);
 
     if let Some(ref mut file) = file {
         let _ = writeln!(file, "\n\nquery stack during panic:");

compiler/rustc_query_impl/src/plumbing.rs

@@ -36,8 +36,9 @@ use crate::job::find_dep_kind_root;
 use crate::{GetQueryVTable, collect_active_jobs_from_all_queries};
 
 fn depth_limit_error<'tcx>(tcx: TyCtxt<'tcx>, job: QueryJobId) {
-    let job_map =
-        collect_active_jobs_from_all_queries(tcx, true).expect("failed to collect active queries");
+    let (job_map, collect_ok) = collect_active_jobs_from_all_queries(tcx);
+    assert!(collect_ok, "failed to collect active queries for depth limit error");
+
     let (info, depth) = find_dep_kind_root(job, job_map);
 
     let suggested_limit = match tcx.recursion_limit() {
@@ -512,36 +513,25 @@ macro_rules! define_queries {
             }
         }
 
-        /// Returns a map of currently active query jobs, collected from all queries.
-        ///
-        /// If `require_complete` is `true`, this function locks all shards of the
-        /// query results to produce a complete map, which always returns `Ok`.
-        /// Otherwise, it may return an incomplete map as an error if any shard
-        /// lock cannot be acquired.
-        ///
-        /// Prefer passing `false` to `require_complete` to avoid potential deadlocks,
-        /// especially when called from within a deadlock handler, unless a
-        /// complete map is needed and no deadlock is possible at this call site.
+        /// Returns a map of currently active query jobs, collected from all queries, and a success
+        /// indicator. If the indicator is false, the map will be incomplete because one or more
+        /// shard locks could not be acquired.
         pub fn collect_active_jobs_from_all_queries<'tcx>(
             tcx: TyCtxt<'tcx>,
-            require_complete: bool,
-        ) -> Result<QueryJobMap<'tcx>, QueryJobMap<'tcx>> {
-            let mut job_map_out = QueryJobMap::default();
-            let mut complete = true;
+        ) -> (QueryJobMap<'tcx>, bool) {
+            let mut job_map = QueryJobMap::default();
+            let mut all_ok = true;
 
             $(
-                let res = crate::execution::gather_active_jobs(
+                let query_ok = crate::execution::gather_active_jobs(
                     &tcx.query_system.query_vtables.$name,
                     tcx,
-                    require_complete,
-                    &mut job_map_out,
+                    &mut job_map,
                 );
-                if res.is_none() {
-                    complete = false;
-                }
+                all_ok = all_ok && query_ok;
             )*
 
-            if complete { Ok(job_map_out) } else { Err(job_map_out) }
+            (job_map, all_ok)
         }
 
         /// All self-profiling events generated by the query engine use