Rollup merge of #138672 - Zoxc:deferred-queries-in-deadlock-handler, r=oli-obk

Avoiding calling queries when collecting active queries

This PR changes active query collection to no longer call queries. Instead the fields needing queries have their computation delayed to when an cycle error is emitted or when printing the query backtrace in a panic.

This is done by splitting the fields in `QueryStackFrame` needing queries into a new `QueryStackFrameExtra` type. When collecting queries `QueryStackFrame` will contain a closure that can create `QueryStackFrameExtra`, which does make use of queries. Calling `lift` on a `QueryStackFrame` or `CycleError` will convert it to a variant containing `QueryStackFrameExtra` using those closures.

This also only calls queries needed to collect information on a cycle errors, instead of information on all active queries.

Calling queries when collecting active queries is a bit odd. Calling queries should not be done in the deadlock handler at all.

This avoids the out of memory scenario in #124901.

Author: Zalathar

compiler/rustc_interface/src/util.rs

@@ -18,7 +18,7 @@ use rustc_session::{EarlyDiagCtxt, Session, filesearch};
 use rustc_span::edit_distance::find_best_match_for_name;
 use rustc_span::edition::Edition;
 use rustc_span::source_map::SourceMapInputs;
-use rustc_span::{Symbol, sym};
+use rustc_span::{SessionGlobals, Symbol, sym};
 use rustc_target::spec::Target;
 use tracing::info;
 
@@ -188,26 +188,11 @@ pub(crate) fn run_in_thread_pool_with_globals<F: FnOnce(CurrentGcx) -> R + Send,
             // On deadlock, creates a new thread and forwards information in thread
             // locals to it. The new thread runs the deadlock handler.
 
-            // Get a `GlobalCtxt` reference from `CurrentGcx` as we cannot rely on having a
-            // `TyCtxt` TLS reference here.
-            let query_map = current_gcx2.access(|gcx| {
-                tls::enter_context(&tls::ImplicitCtxt::new(gcx), || {
-                    tls::with(|tcx| {
-                        match QueryCtxt::new(tcx).collect_active_jobs() {
-                            Ok(query_map) => query_map,
-                            Err(_) => {
-                                // There was an unexpected error collecting all active jobs, which we need
-                                // to find cycles to break.
-                                // We want to avoid panicking in the deadlock handler, so we abort instead.
-                                eprintln!("internal compiler error: failed to get query map in deadlock handler, aborting process");
-                                process::abort();
-                            }
-                        }
-                    })
-                })
-            });
-            let query_map = FromDyn::from(query_map);
+            let current_gcx2 = current_gcx2.clone();
             let registry = rayon_core::Registry::current();
+            let session_globals = rustc_span::with_session_globals(|session_globals| {
+                session_globals as *const SessionGlobals as usize
+            });
             thread::Builder::new()
                 .name("rustc query cycle handler".to_string())
                 .spawn(move || {
@@ -217,7 +202,24 @@ pub(crate) fn run_in_thread_pool_with_globals<F: FnOnce(CurrentGcx) -> R + Send,
                         // otherwise the compiler could just hang,
                         process::abort();
                     });
-                    break_query_cycles(query_map.into_inner(), &registry);
+
+                    // Get a `GlobalCtxt` reference from `CurrentGcx` as we cannot rely on having a
+                    // `TyCtxt` TLS reference here.
+                    current_gcx2.access(|gcx| {
+                        tls::enter_context(&tls::ImplicitCtxt::new(gcx), || {
+                            tls::with(|tcx| {
+                                // Accessing session globals is sound as they outlive `GlobalCtxt`.
+                                // They are needed to hash query keys containing spans or symbols.
+                                let query_map = rustc_span::set_session_globals_then(unsafe { &*(session_globals as *const SessionGlobals) }, || {
+                                    // Ensure there was no errors collecting all active jobs.
+                                    // We need the complete map to ensure we find a cycle to break.
+                                    QueryCtxt::new(tcx).collect_active_jobs().ok().expect("failed to collect active queries in deadlock handler")
+                                });
+                                break_query_cycles(query_map, &registry);
+                            })
+                        })
+                    });
+
                     on_panic.disable();
                 })
                 .unwrap();

compiler/rustc_middle/src/query/mod.rs

@@ -31,7 +31,9 @@ use rustc_index::IndexVec;
 use rustc_lint_defs::LintId;
 use rustc_macros::rustc_queries;
 use rustc_query_system::ich::StableHashingContext;
-use rustc_query_system::query::{QueryCache, QueryMode, QueryState, try_get_cached};
+use rustc_query_system::query::{
+    QueryCache, QueryMode, QueryStackDeferred, QueryState, try_get_cached,
+};
 use rustc_session::Limits;
 use rustc_session::config::{EntryFnType, OptLevel, OutputFilenames, SymbolManglingVersion};
 use rustc_session::cstore::{

compiler/rustc_middle/src/query/plumbing.rs

@@ -488,7 +488,7 @@ macro_rules! define_callbacks {
         #[derive(Default)]
         pub struct QueryStates<'tcx> {
             $(
-                pub $name: QueryState<$($K)*>,
+                pub $name: QueryState<$($K)*, QueryStackDeferred<'tcx>>,
             )*
         }
 

compiler/rustc_middle/src/values.rs

@@ -88,7 +88,7 @@ impl<'tcx> Value<TyCtxt<'tcx>> for Representability {
             if info.query.dep_kind == dep_kinds::representability
                 && let Some(field_id) = info.query.def_id
                 && let Some(field_id) = field_id.as_local()
-                && let Some(DefKind::Field) = info.query.def_kind
+                && let Some(DefKind::Field) = info.query.info.def_kind
             {
                 let parent_id = tcx.parent(field_id.to_def_id());
                 let item_id = match tcx.def_kind(parent_id) {
@@ -216,7 +216,7 @@ impl<'tcx, T> Value<TyCtxt<'tcx>> for Result<T, &'_ ty::layout::LayoutError<'_>>
                             continue;
                         };
                         let frame_span =
-                            frame.query.default_span(cycle[(i + 1) % cycle.len()].span);
+                            frame.query.info.default_span(cycle[(i + 1) % cycle.len()].span);
                         if frame_span.is_dummy() {
                             continue;
                         }

compiler/rustc_query_impl/src/lib.rs

@@ -26,8 +26,8 @@ use rustc_middle::ty::TyCtxt;
 use rustc_query_system::dep_graph::SerializedDepNodeIndex;
 use rustc_query_system::ich::StableHashingContext;
 use rustc_query_system::query::{
-    CycleError, HashResult, QueryCache, QueryConfig, QueryMap, QueryMode, QueryState,
-    get_query_incr, get_query_non_incr,
+    CycleError, HashResult, QueryCache, QueryConfig, QueryMap, QueryMode, QueryStackDeferred,
+    QueryState, get_query_incr, get_query_non_incr,
 };
 use rustc_query_system::{HandleCycleError, Value};
 use rustc_span::{ErrorGuaranteed, Span};
@@ -84,7 +84,10 @@ where
     }
 
     #[inline(always)]
-    fn query_state<'a>(self, qcx: QueryCtxt<'tcx>) -> &'a QueryState<Self::Key>
+    fn query_state<'a>(
+        self,
+        qcx: QueryCtxt<'tcx>,
+    ) -> &'a QueryState<Self::Key, QueryStackDeferred<'tcx>>
     where
         QueryCtxt<'tcx>: 'a,
     {
@@ -93,7 +96,7 @@ where
         unsafe {
             &*(&qcx.tcx.query_system.states as *const QueryStates<'tcx>)
                 .byte_add(self.dynamic.query_state)
-                .cast::<QueryState<Self::Key>>()
+                .cast::<QueryState<Self::Key, QueryStackDeferred<'tcx>>>()
         }
     }
 

compiler/rustc_query_impl/src/plumbing.rs

@@ -5,6 +5,7 @@
 use std::num::NonZero;
 
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
+use rustc_data_structures::sync::{DynSend, DynSync};
 use rustc_data_structures::unord::UnordMap;
 use rustc_hashes::Hash64;
 use rustc_index::Idx;
@@ -24,8 +25,8 @@ use rustc_middle::ty::{self, TyCtxt};
 use rustc_query_system::dep_graph::{DepNodeParams, HasDepContext};
 use rustc_query_system::ich::StableHashingContext;
 use rustc_query_system::query::{
-    QueryCache, QueryConfig, QueryContext, QueryJobId, QueryMap, QuerySideEffect, QueryStackFrame,
-    force_query,
+    QueryCache, QueryConfig, QueryContext, QueryJobId, QueryMap, QuerySideEffect,
+    QueryStackDeferred, QueryStackFrame, QueryStackFrameExtra, force_query,
 };
 use rustc_query_system::{QueryOverflow, QueryOverflowNote};
 use rustc_serialize::{Decodable, Encodable};
@@ -65,7 +66,9 @@ impl<'tcx> HasDepContext for QueryCtxt<'tcx> {
     }
 }
 
-impl QueryContext for QueryCtxt<'_> {
+impl<'tcx> QueryContext for QueryCtxt<'tcx> {
+    type QueryInfo = QueryStackDeferred<'tcx>;
+
     #[inline]
     fn next_job_id(self) -> QueryJobId {
         QueryJobId(
@@ -82,7 +85,9 @@ impl QueryContext for QueryCtxt<'_> {
     /// Returns a query map representing active query jobs.
     /// It returns an incomplete map as an error if it fails
     /// to take locks.
-    fn collect_active_jobs(self) -> Result<QueryMap, QueryMap> {
+    fn collect_active_jobs(
+        self,
+    ) -> Result<QueryMap<QueryStackDeferred<'tcx>>, QueryMap<QueryStackDeferred<'tcx>>> {
         let mut jobs = QueryMap::default();
         let mut complete = true;
 
@@ -95,6 +100,13 @@ impl QueryContext for QueryCtxt<'_> {
         if complete { Ok(jobs) } else { Err(jobs) }
     }
 
+    fn lift_query_info(
+        self,
+        info: &QueryStackDeferred<'tcx>,
+    ) -> rustc_query_system::query::QueryStackFrameExtra {
+        info.extract()
+    }
+
     // Interactions with on_disk_cache
     fn load_side_effect(
         self,
@@ -159,7 +171,10 @@ impl QueryContext for QueryCtxt<'_> {
 
         self.sess.dcx().emit_fatal(QueryOverflow {
             span: info.job.span,
-            note: QueryOverflowNote { desc: info.query.description, depth },
+            note: QueryOverflowNote {
+                desc: self.lift_query_info(&info.query.info).description,
+                depth,
+            },
             suggested_limit,
             crate_name: self.crate_name(LOCAL_CRATE),
         });
@@ -296,16 +311,17 @@ macro_rules! should_ever_cache_on_disk {
     };
 }
 
-pub(crate) fn create_query_frame<
-    'tcx,
-    K: Copy + Key + for<'a> HashStable<StableHashingContext<'a>>,
->(
-    tcx: TyCtxt<'tcx>,
-    do_describe: fn(TyCtxt<'tcx>, K) -> String,
-    key: K,
-    kind: DepKind,
-    name: &'static str,
-) -> QueryStackFrame {
+fn create_query_frame_extra<'tcx, K: Key + Copy + 'tcx>(
+    (tcx, key, kind, name, do_describe): (
+        TyCtxt<'tcx>,
+        K,
+        DepKind,
+        &'static str,
+        fn(TyCtxt<'tcx>, K) -> String,
+    ),
+) -> QueryStackFrameExtra {
+    let def_id = key.key_as_def_id();
+
     // If reduced queries are requested, we may be printing a query stack due
     // to a panic. Avoid using `default_span` and `def_kind` in that case.
     let reduce_queries = with_reduced_queries();
@@ -324,13 +340,28 @@ pub(crate) fn create_query_frame<
     } else {
         Some(key.default_span(tcx))
     };
-    let def_id = key.key_as_def_id();
+
     let def_kind = if kind == dep_graph::dep_kinds::def_kind || reduce_queries {
         // Try to avoid infinite recursion.
         None
     } else {
         def_id.and_then(|def_id| def_id.as_local()).map(|def_id| tcx.def_kind(def_id))
     };
+    QueryStackFrameExtra::new(description, span, def_kind)
+}
+
+pub(crate) fn create_query_frame<
+    'tcx,
+    K: Copy + DynSend + DynSync + Key + for<'a> HashStable<StableHashingContext<'a>> + 'tcx,
+>(
+    tcx: TyCtxt<'tcx>,
+    do_describe: fn(TyCtxt<'tcx>, K) -> String,
+    key: K,
+    kind: DepKind,
+    name: &'static str,
+) -> QueryStackFrame<QueryStackDeferred<'tcx>> {
+    let def_id = key.key_as_def_id();
+
     let hash = || {
         tcx.with_stable_hashing_context(|mut hcx| {
             let mut hasher = StableHasher::new();
@@ -341,7 +372,10 @@ pub(crate) fn create_query_frame<
     };
     let def_id_for_ty_in_cycle = key.def_id_for_ty_in_cycle();
 
-    QueryStackFrame::new(description, span, def_id, def_kind, kind, def_id_for_ty_in_cycle, hash)
+    let info =
+        QueryStackDeferred::new((tcx, key, kind, name, do_describe), create_query_frame_extra);
+
+    QueryStackFrame::new(info, kind, hash, def_id, def_id_for_ty_in_cycle)
 }
 
 pub(crate) fn encode_query_results<'a, 'tcx, Q>(
@@ -688,7 +722,10 @@ macro_rules! define_queries {
                 }
             }
 
-            pub(crate) fn try_collect_active_jobs<'tcx>(tcx: TyCtxt<'tcx>, qmap: &mut QueryMap) -> Option<()> {
+            pub(crate) fn try_collect_active_jobs<'tcx>(
+                tcx: TyCtxt<'tcx>,
+                qmap: &mut QueryMap<QueryStackDeferred<'tcx>>,
+            ) -> Option<()> {
                 let make_query = |tcx, key| {
                     let kind = rustc_middle::dep_graph::dep_kinds::$name;
                     let name = stringify!($name);
@@ -768,7 +805,9 @@ macro_rules! define_queries {
 
         // These arrays are used for iteration and can't be indexed by `DepKind`.
 
-        const TRY_COLLECT_ACTIVE_JOBS: &[for<'tcx> fn(TyCtxt<'tcx>, &mut QueryMap) -> Option<()>] =
+        const TRY_COLLECT_ACTIVE_JOBS: &[
+            for<'tcx> fn(TyCtxt<'tcx>, &mut QueryMap<QueryStackDeferred<'tcx>>) -> Option<()>
+        ] =
             &[$(query_impl::$name::try_collect_active_jobs),*];
 
         const ALLOC_SELF_PROFILE_QUERY_STRINGS: &[

compiler/rustc_query_system/src/query/config.rs

@@ -6,6 +6,7 @@ use std::hash::Hash;
 use rustc_data_structures::fingerprint::Fingerprint;
 use rustc_span::ErrorGuaranteed;
 
+use super::QueryStackFrameExtra;
 use crate::dep_graph::{DepKind, DepNode, DepNodeParams, SerializedDepNodeIndex};
 use crate::error::HandleCycleError;
 use crate::ich::StableHashingContext;
@@ -27,7 +28,7 @@ pub trait QueryConfig<Qcx: QueryContext>: Copy {
     fn format_value(self) -> fn(&Self::Value) -> String;
 
     // Don't use this method to access query results, instead use the methods on TyCtxt
-    fn query_state<'a>(self, tcx: Qcx) -> &'a QueryState<Self::Key>
+    fn query_state<'a>(self, tcx: Qcx) -> &'a QueryState<Self::Key, Qcx::QueryInfo>
     where
         Qcx: 'a;
 
@@ -57,7 +58,7 @@ pub trait QueryConfig<Qcx: QueryContext>: Copy {
     fn value_from_cycle_error(
         self,
         tcx: Qcx::DepContext,
-        cycle_error: &CycleError,
+        cycle_error: &CycleError<QueryStackFrameExtra>,
         guar: ErrorGuaranteed,
     ) -> Self::Value;