Avoiding calling queries when collecting active queries

10 files changed, 313 insertions, 162 deletions

diff --git a/compiler/rustc_interface/src/util.rs b/compiler/rustc_interface/src/util.rs
index 333786f0ca3..83d80938b4e 100644
--- a/compiler/rustc_interface/src/util.rs
+++ b/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();
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 527c18addbe..6093d0c6e84 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -30,7 +30,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::{
diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs
index 4834444ed1d..a099f770417 100644
--- a/compiler/rustc_middle/src/query/plumbing.rs
+++ b/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>>,
             )*
         }
 
diff --git a/compiler/rustc_middle/src/values.rs b/compiler/rustc_middle/src/values.rs
index 9450ce7ec44..39fcc686c55 100644
--- a/compiler/rustc_middle/src/values.rs
+++ b/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;
                         }
diff --git a/compiler/rustc_query_impl/src/lib.rs b/compiler/rustc_query_impl/src/lib.rs
index a83c388c747..30a9e718d23 100644
--- a/compiler/rustc_query_impl/src/lib.rs
+++ b/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>>>()
         }
     }
 
diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs
index 55281cd5ac7..7f1d466b869 100644
--- a/compiler/rustc_query_impl/src/plumbing.rs
+++ b/compiler/rustc_query_impl/src/plumbing.rs
@@ -3,8 +3,10 @@
 //! manage the caches, and so forth.
 
 use std::num::NonZero;
+use std::sync::Arc;
 
 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 +26,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 +67,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 +86,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 +101,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 +172,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),
         });
@@ -298,39 +314,45 @@ macro_rules! should_ever_cache_on_disk {
 
 pub(crate) fn create_query_frame<
     'tcx,
-    K: Copy + Key + for<'a> HashStable<StableHashingContext<'a>>,
+    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 {
-    // 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();
-
-    // Avoid calling queries while formatting the description
-    let description = ty::print::with_no_queries!(do_describe(tcx, key));
-    let description = if tcx.sess.verbose_internals() {
-        format!("{description} [{name:?}]")
-    } else {
-        description
-    };
-    let span = if kind == dep_graph::dep_kinds::def_span || reduce_queries {
-        // The `def_span` query is used to calculate `default_span`,
-        // so exit to avoid infinite recursion.
-        None
-    } else {
-        Some(key.default_span(tcx))
-    };
+) -> QueryStackFrame<QueryStackDeferred<'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))
+
+    let extra = move || {
+        // 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();
+
+        // Avoid calling queries while formatting the description
+        let description = ty::print::with_no_queries!(do_describe(tcx, key));
+        let description = if tcx.sess.verbose_internals() {
+            format!("{description} [{name:?}]")
+        } else {
+            description
+        };
+        let span = if kind == dep_graph::dep_kinds::def_span || reduce_queries {
+            // The `def_span` query is used to calculate `default_span`,
+            // so exit to avoid infinite recursion.
+            None
+        } else {
+            Some(key.default_span(tcx))
+        };
+
+        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)
     };
+
     let hash = || {
         tcx.with_stable_hashing_context(|mut hcx| {
             let mut hasher = StableHasher::new();
@@ -341,7 +363,11 @@ 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)
+    // SAFETY: None of the captures in `extra` have destructors that access 'tcx
+    // as they don't have destructors.
+    let info = unsafe { QueryStackDeferred::new(Arc::new(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 +714,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 +797,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: &[
diff --git a/compiler/rustc_query_system/src/query/config.rs b/compiler/rustc_query_system/src/query/config.rs
index 371b896400a..e508eadb73b 100644
--- a/compiler/rustc_query_system/src/query/config.rs
+++ b/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;
 
diff --git a/compiler/rustc_query_system/src/query/job.rs b/compiler/rustc_query_system/src/query/job.rs
index 402c7831472..de35cd79ea2 100644
--- a/compiler/rustc_query_system/src/query/job.rs
+++ b/compiler/rustc_query_system/src/query/job.rs
@@ -1,3 +1,4 @@
+use std::fmt::Debug;
 use std::hash::Hash;
 use std::io::Write;
 use std::iter;
@@ -12,6 +13,7 @@ use rustc_hir::def::DefKind;
 use rustc_session::Session;
 use rustc_span::{DUMMY_SP, Span};
 
+use super::QueryStackFrameExtra;
 use crate::dep_graph::DepContext;
 use crate::error::CycleStack;
 use crate::query::plumbing::CycleError;
@@ -19,45 +21,54 @@ use crate::query::{QueryContext, QueryStackFrame};
 
 /// Represents a span and a query key.
 #[derive(Clone, Debug)]
-pub struct QueryInfo {
+pub struct QueryInfo<I> {
     /// The span corresponding to the reason for which this query was required.
     pub span: Span,
-    pub query: QueryStackFrame,
+    pub query: QueryStackFrame<I>,
 }
 
-pub type QueryMap = FxHashMap<QueryJobId, QueryJobInfo>;
+impl<I> QueryInfo<I> {
+    pub(crate) fn lift<Qcx: QueryContext<QueryInfo = I>>(
+        &self,
+        qcx: Qcx,
+    ) -> QueryInfo<QueryStackFrameExtra> {
+        QueryInfo { span: self.span, query: self.query.lift(qcx) }
+    }
+}
+
+pub type QueryMap<I> = FxHashMap<QueryJobId, QueryJobInfo<I>>;
 
 /// A value uniquely identifying an active query job.
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
 pub struct QueryJobId(pub NonZero<u64>);
 
 impl QueryJobId {
-    fn query(self, map: &QueryMap) -> QueryStackFrame {
+    fn query<I: Clone>(self, map: &QueryMap<I>) -> QueryStackFrame<I> {
         map.get(&self).unwrap().query.clone()
     }
 
-    fn span(self, map: &QueryMap) -> Span {
+    fn span<I>(self, map: &QueryMap<I>) -> Span {
         map.get(&self).unwrap().job.span
     }
 
-    fn parent(self, map: &QueryMap) -> Option<QueryJobId> {
+    fn parent<I>(self, map: &QueryMap<I>) -> Option<QueryJobId> {
         map.get(&self).unwrap().job.parent
     }
 
-    fn latch(self, map: &QueryMap) -> Option<&QueryLatch> {
+    fn latch<I>(self, map: &QueryMap<I>) -> Option<&QueryLatch<I>> {
         map.get(&self).unwrap().job.latch.as_ref()
     }
 }
 
 #[derive(Clone, Debug)]
-pub struct QueryJobInfo {
-    pub query: QueryStackFrame,
-    pub job: QueryJob,
+pub struct QueryJobInfo<I> {
+    pub query: QueryStackFrame<I>,
+    pub job: QueryJob<I>,
 }
 
 /// Represents an active query job.
-#[derive(Clone, Debug)]
-pub struct QueryJob {
+#[derive(Debug)]
+pub struct QueryJob<I> {
     pub id: QueryJobId,
 
     /// The span corresponding to the reason for which this query was required.
@@ -67,17 +78,23 @@ pub struct QueryJob {
     pub parent: Option<QueryJobId>,
 
     /// The latch that is used to wait on this job.
-    latch: Option<QueryLatch>,
+    latch: Option<QueryLatch<I>>,
 }
 
-impl QueryJob {
+impl<I> Clone for QueryJob<I> {
+    fn clone(&self) -> Self {
+        Self { id: self.id, span: self.span, parent: self.parent, latch: self.latch.clone() }
+    }
+}
+
+impl<I> QueryJob<I> {
     /// Creates a new query job.
     #[inline]
     pub fn new(id: QueryJobId, span: Span, parent: Option<QueryJobId>) -> Self {
         QueryJob { id, span, parent, latch: None }
     }
 
-    pub(super) fn latch(&mut self) -> QueryLatch {
+    pub(super) fn latch(&mut self) -> QueryLatch<I> {
         if self.latch.is_none() {
             self.latch = Some(QueryLatch::new());
         }
@@ -97,12 +114,12 @@ impl QueryJob {
 }
 
 impl QueryJobId {
-    pub(super) fn find_cycle_in_stack(
+    pub(super) fn find_cycle_in_stack<I: Clone>(
         &self,
-        query_map: QueryMap,
+        query_map: QueryMap<I>,
         current_job: &Option<QueryJobId>,
         span: Span,
-    ) -> CycleError {
+    ) -> CycleError<I> {
         // Find the waitee amongst `current_job` parents
         let mut cycle = Vec::new();
         let mut current_job = Option::clone(current_job);
@@ -136,7 +153,7 @@ impl QueryJobId {
 
     #[cold]
     #[inline(never)]
-    pub fn find_dep_kind_root(&self, query_map: QueryMap) -> (QueryJobInfo, usize) {
+    pub fn find_dep_kind_root<I: Clone>(&self, query_map: QueryMap<I>) -> (QueryJobInfo<I>, usize) {
         let mut depth = 1;
         let info = query_map.get(&self).unwrap();
         let dep_kind = info.query.dep_kind;
@@ -156,25 +173,31 @@ impl QueryJobId {
 }
 
 #[derive(Debug)]
-struct QueryWaiter {
+struct QueryWaiter<I> {
     query: Option<QueryJobId>,
     condvar: Condvar,
     span: Span,
-    cycle: Mutex<Option<CycleError>>,
+    cycle: Mutex<Option<CycleError<I>>>,
 }
 
 #[derive(Debug)]
-struct QueryLatchInfo {
+struct QueryLatchInfo<I> {
     complete: bool,
-    waiters: Vec<Arc<QueryWaiter>>,
+    waiters: Vec<Arc<QueryWaiter<I>>>,
 }
 
-#[derive(Clone, Debug)]
-pub(super) struct QueryLatch {
-    info: Arc<Mutex<QueryLatchInfo>>,
+#[derive(Debug)]
+pub(super) struct QueryLatch<I> {
+    info: Arc<Mutex<QueryLatchInfo<I>>>,
 }
 
-impl QueryLatch {
+impl<I> Clone for QueryLatch<I> {
+    fn clone(&self) -> Self {
+        Self { info: Arc::clone(&self.info) }
+    }
+}
+
+impl<I> QueryLatch<I> {
     fn new() -> Self {
         QueryLatch {
             info: Arc::new(Mutex::new(QueryLatchInfo { complete: false, waiters: Vec::new() })),
@@ -182,7 +205,11 @@ impl QueryLatch {
     }
 
     /// Awaits for the query job to complete.
-    pub(super) fn wait_on(&self, query: Option<QueryJobId>, span: Span) -> Result<(), CycleError> {
+    pub(super) fn wait_on(
+        &self,
+        query: Option<QueryJobId>,
+        span: Span,
+    ) -> Result<(), CycleError<I>> {
         let waiter =
             Arc::new(QueryWaiter { query, span, cycle: Mutex::new(None), condvar: Condvar::new() });
         self.wait_on_inner(&waiter);
@@ -197,7 +224,7 @@ impl QueryLatch {
     }
 
     /// Awaits the caller on this latch by blocking the current thread.
-    fn wait_on_inner(&self, waiter: &Arc<QueryWaiter>) {
+    fn wait_on_inner(&self, waiter: &Arc<QueryWaiter<I>>) {
         let mut info = self.info.lock();
         if !info.complete {
             // We push the waiter on to the `waiters` list. It can be accessed inside
@@ -232,7 +259,7 @@ impl QueryLatch {
 
     /// Removes a single waiter from the list of waiters.
     /// This is used to break query cycles.
-    fn extract_waiter(&self, waiter: usize) -> Arc<QueryWaiter> {
+    fn extract_waiter(&self, waiter: usize) -> Arc<QueryWaiter<I>> {
         let mut info = self.info.lock();
         debug_assert!(!info.complete);
         // Remove the waiter from the list of waiters
@@ -252,7 +279,11 @@ type Waiter = (QueryJobId, usize);
 /// For visits of resumable waiters it returns Some(Some(Waiter)) which has the
 /// required information to resume the waiter.
 /// If all `visit` calls returns None, this function also returns None.
-fn visit_waiters<F>(query_map: &QueryMap, query: QueryJobId, mut visit: F) -> Option<Option<Waiter>>
+fn visit_waiters<I, F>(
+    query_map: &QueryMap<I>,
+    query: QueryJobId,
+    mut visit: F,
+) -> Option<Option<Waiter>>
 where
     F: FnMut(Span, QueryJobId) -> Option<Option<Waiter>>,
 {
@@ -282,8 +313,8 @@ where
 /// `span` is the reason for the `query` to execute. This is initially DUMMY_SP.
 /// If a cycle is detected, this initial value is replaced with the span causing
 /// the cycle.
-fn cycle_check(
-    query_map: &QueryMap,
+fn cycle_check<I>(
+    query_map: &QueryMap<I>,
     query: QueryJobId,
     span: Span,
     stack: &mut Vec<(Span, QueryJobId)>,
@@ -322,8 +353,8 @@ fn cycle_check(
 /// Finds out if there's a path to the compiler root (aka. code which isn't in a query)
 /// from `query` without going through any of the queries in `visited`.
 /// This is achieved with a depth first search.
-fn connected_to_root(
-    query_map: &QueryMap,
+fn connected_to_root<I>(
+    query_map: &QueryMap<I>,
     query: QueryJobId,
     visited: &mut FxHashSet<QueryJobId>,
 ) -> bool {
@@ -344,7 +375,7 @@ fn connected_to_root(
 }
 
 // Deterministically pick an query from a list
-fn pick_query<'a, T, F>(query_map: &QueryMap, queries: &'a [T], f: F) -> &'a T
+fn pick_query<'a, I: Clone, T, F>(query_map: &QueryMap<I>, queries: &'a [T], f: F) -> &'a T
 where
     F: Fn(&T) -> (Span, QueryJobId),
 {
@@ -369,10 +400,10 @@ where
 /// the function return true.
 /// If a cycle was not found, the starting query is removed from `jobs` and
 /// the function returns false.
-fn remove_cycle(
-    query_map: &QueryMap,
+fn remove_cycle<I: Clone>(
+    query_map: &QueryMap<I>,
     jobs: &mut Vec<QueryJobId>,
-    wakelist: &mut Vec<Arc<QueryWaiter>>,
+    wakelist: &mut Vec<Arc<QueryWaiter<I>>>,
 ) -> bool {
     let mut visited = FxHashSet::default();
     let mut stack = Vec::new();
@@ -473,7 +504,10 @@ fn remove_cycle(
 /// uses a query latch and then resuming that waiter.
 /// There may be multiple cycles involved in a deadlock, so this searches
 /// all active queries for cycles before finally resuming all the waiters at once.
-pub fn break_query_cycles(query_map: QueryMap, registry: &rayon_core::Registry) {
+pub fn break_query_cycles<I: Clone + Debug>(
+    query_map: QueryMap<I>,
+    registry: &rayon_core::Registry,
+) {
     let mut wakelist = Vec::new();
     let mut jobs: Vec<QueryJobId> = query_map.keys().cloned().collect();
 
@@ -520,7 +554,7 @@ pub fn report_cycle<'a>(
 ) -> Diag<'a> {
     assert!(!stack.is_empty());
 
-    let span = stack[0].query.default_span(stack[1 % stack.len()].span);
+    let span = stack[0].query.info.default_span(stack[1 % stack.len()].span);
 
     let mut cycle_stack = Vec::new();
 
@@ -529,31 +563,31 @@ pub fn report_cycle<'a>(
 
     for i in 1..stack.len() {
         let query = &stack[i].query;
-        let span = query.default_span(stack[(i + 1) % stack.len()].span);
-        cycle_stack.push(CycleStack { span, desc: query.description.to_owned() });
+        let span = query.info.default_span(stack[(i + 1) % stack.len()].span);
+        cycle_stack.push(CycleStack { span, desc: query.info.description.to_owned() });
     }
 
     let mut cycle_usage = None;
     if let Some((span, ref query)) = *usage {
         cycle_usage = Some(crate::error::CycleUsage {
-            span: query.default_span(span),
-            usage: query.description.to_string(),
+            span: query.info.default_span(span),
+            usage: query.info.description.to_string(),
         });
     }
 
-    let alias = if stack.iter().all(|entry| matches!(entry.query.def_kind, Some(DefKind::TyAlias)))
-    {
-        Some(crate::error::Alias::Ty)
-    } else if stack.iter().all(|entry| entry.query.def_kind == Some(DefKind::TraitAlias)) {
-        Some(crate::error::Alias::Trait)
-    } else {
-        None
-    };
+    let alias =
+        if stack.iter().all(|entry| matches!(entry.query.info.def_kind, Some(DefKind::TyAlias))) {
+            Some(crate::error::Alias::Ty)
+        } else if stack.iter().all(|entry| entry.query.info.def_kind == Some(DefKind::TraitAlias)) {
+            Some(crate::error::Alias::Trait)
+        } else {
+            None
+        };
 
     let cycle_diag = crate::error::Cycle {
         span,
         cycle_stack,
-        stack_bottom: stack[0].query.description.to_owned(),
+        stack_bottom: stack[0].query.info.description.to_owned(),
         alias,
         cycle_usage,
         stack_count,
@@ -589,6 +623,7 @@ pub fn print_query_stack<Qcx: QueryContext>(
         let Some(query_info) = query_map.get(&query) else {
             break;
         };
+        let query_extra = qcx.lift_query_info(&query_info.query.info);
         if Some(count_printed) < limit_frames || limit_frames.is_none() {
             // Only print to stderr as many stack frames as `num_frames` when present.
             // FIXME: needs translation
@@ -596,7 +631,7 @@ pub fn print_query_stack<Qcx: QueryContext>(
             #[allow(rustc::untranslatable_diagnostic)]
             dcx.struct_failure_note(format!(
                 "#{} [{:?}] {}",
-                count_printed, query_info.query.dep_kind, query_info.query.description
+                count_printed, query_info.query.dep_kind, query_extra.description
             ))
             .with_span(query_info.job.span)
             .emit();
@@ -609,7 +644,7 @@ pub fn print_query_stack<Qcx: QueryContext>(
                 "#{} [{}] {}",
                 count_total,
                 qcx.dep_context().dep_kind_info(query_info.query.dep_kind).name,
-                query_info.query.description
+                query_extra.description
             );
         }
 
diff --git a/compiler/rustc_query_system/src/query/mod.rs b/compiler/rustc_query_system/src/query/mod.rs
index 0d0c66aa978..d9ee6532831 100644
--- a/compiler/rustc_query_system/src/query/mod.rs
+++ b/compiler/rustc_query_system/src/query/mod.rs
@@ -1,4 +1,9 @@
 mod plumbing;
+use std::fmt::Debug;
+use std::marker::PhantomData;
+use std::mem::transmute;
+use std::sync::Arc;
+
 pub use self::plumbing::*;
 
 mod job;
@@ -11,6 +16,7 @@ mod caches;
 pub use self::caches::{DefIdCache, DefaultCache, QueryCache, SingleCache, VecCache};
 
 mod config;
+use rustc_data_structures::sync::{DynSend, DynSync};
 use rustc_errors::DiagInner;
 use rustc_hashes::Hash64;
 use rustc_hir::def::DefKind;
@@ -25,31 +31,59 @@ use crate::dep_graph::{DepKind, DepNodeIndex, HasDepContext, SerializedDepNodeIn
 ///
 /// This is mostly used in case of cycles for error reporting.
 #[derive(Clone, Debug)]
-pub struct QueryStackFrame {
-    pub description: String,
-    span: Option<Span>,
-    pub def_id: Option<DefId>,
-    pub def_kind: Option<DefKind>,
-    /// A def-id that is extracted from a `Ty` in a query key
-    pub def_id_for_ty_in_cycle: Option<DefId>,
+pub struct QueryStackFrame<I> {
+    /// This field initially stores a `QueryStackDeferred` during collection,
+    /// but can later be changed to `QueryStackFrameExtra` containing concrete information
+    /// by calling `lift`. This is done so that collecting query does not need to invoke
+    /// queries, instead `lift` will call queries in a more appropriate location.
+    pub info: I,
+
     pub dep_kind: DepKind,
     /// This hash is used to deterministically pick
     /// a query to remove cycles in the parallel compiler.
     hash: Hash64,
+    pub def_id: Option<DefId>,
+    /// A def-id that is extracted from a `Ty` in a query key
+    pub def_id_for_ty_in_cycle: Option<DefId>,
 }
 
-impl QueryStackFrame {
+impl<I> QueryStackFrame<I> {
     #[inline]
     pub fn new(
-        description: String,
-        span: Option<Span>,
-        def_id: Option<DefId>,
-        def_kind: Option<DefKind>,
+        info: I,
         dep_kind: DepKind,
-        def_id_for_ty_in_cycle: Option<DefId>,
         hash: impl FnOnce() -> Hash64,
+        def_id: Option<DefId>,
+        def_id_for_ty_in_cycle: Option<DefId>,
     ) -> Self {
-        Self { description, span, def_id, def_kind, def_id_for_ty_in_cycle, dep_kind, hash: hash() }
+        Self { info, def_id, dep_kind, hash: hash(), def_id_for_ty_in_cycle }
+    }
+
+    fn lift<Qcx: QueryContext<QueryInfo = I>>(
+        &self,
+        qcx: Qcx,
+    ) -> QueryStackFrame<QueryStackFrameExtra> {
+        QueryStackFrame {
+            info: qcx.lift_query_info(&self.info),
+            dep_kind: self.dep_kind,
+            hash: self.hash,
+            def_id: self.def_id,
+            def_id_for_ty_in_cycle: self.def_id_for_ty_in_cycle,
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct QueryStackFrameExtra {
+    pub description: String,
+    span: Option<Span>,
+    pub def_kind: Option<DefKind>,
+}
+
+impl QueryStackFrameExtra {
+    #[inline]
+    pub fn new(description: String, span: Option<Span>, def_kind: Option<DefKind>) -> Self {
+        Self { description, span, def_kind }
     }
 
     // FIXME(eddyb) Get more valid `Span`s on queries.
@@ -62,7 +96,37 @@ impl QueryStackFrame {
     }
 }
 
-/// Track a 'side effects' for a particular query.
+/// Track a 'side effect' for a particular query.
+/// This is used to hold a closure which can create `QueryStackFrameExtra`.
+#[derive(Clone)]
+pub struct QueryStackDeferred<'tcx> {
+    _dummy: PhantomData<&'tcx ()>,
+
+    // `extract` may contain references to 'tcx, but we can't tell drop checking that it won't
+    // access it in the destructor.
+    extract: Arc<dyn Fn() -> QueryStackFrameExtra + DynSync + DynSend>,
+}
+
+impl<'tcx> QueryStackDeferred<'tcx> {
+    /// SAFETY: `extract` may not access 'tcx in its destructor.
+    pub unsafe fn new(
+        extract: Arc<dyn Fn() -> QueryStackFrameExtra + DynSync + DynSend + 'tcx>,
+    ) -> Self {
+        Self { _dummy: PhantomData, extract: unsafe { transmute(extract) } }
+    }
+
+    pub fn extract(&self) -> QueryStackFrameExtra {
+        (self.extract)()
+    }
+}
+
+impl<'tcx> Debug for QueryStackDeferred<'tcx> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("QueryStackDeferred")
+    }
+}
+
+/// Tracks 'side effects' for a particular query.
 /// This struct is saved to disk along with the query result,
 /// and loaded from disk if we mark the query as green.
 /// This allows us to 'replay' changes to global state
@@ -81,12 +145,16 @@ pub enum QuerySideEffect {
 }
 
 pub trait QueryContext: HasDepContext {
+    type QueryInfo: Clone;
+
     fn next_job_id(self) -> QueryJobId;
 
     /// Get the query information from the TLS context.
     fn current_query_job(self) -> Option<QueryJobId>;
 
-    fn collect_active_jobs(self) -> Result<QueryMap, QueryMap>;
+    fn collect_active_jobs(self) -> Result<QueryMap<Self::QueryInfo>, QueryMap<Self::QueryInfo>>;
+
+    fn lift_query_info(self, info: &Self::QueryInfo) -> QueryStackFrameExtra;
 
     /// Load a side effect associated to the node in the previous session.
     fn load_side_effect(
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index 3a9d80280c2..6ea8e3b9200 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -16,7 +16,7 @@ use rustc_errors::{Diag, FatalError, StashKey};
 use rustc_span::{DUMMY_SP, Span};
 use tracing::instrument;
 
-use super::QueryConfig;
+use super::{QueryConfig, QueryStackFrameExtra};
 use crate::HandleCycleError;
 use crate::dep_graph::{DepContext, DepGraphData, DepNode, DepNodeIndex, DepNodeParams};
 use crate::ich::StableHashingContext;
@@ -29,23 +29,23 @@ fn equivalent_key<K: Eq, V>(k: &K) -> impl Fn(&(K, V)) -> bool + '_ {
     move |x| x.0 == *k
 }
 
-pub struct QueryState<K> {
-    active: Sharded<hashbrown::HashTable<(K, QueryResult)>>,
+pub struct QueryState<K, I> {
+    active: Sharded<hashbrown::HashTable<(K, QueryResult<I>)>>,
 }
 
 /// Indicates the state of a query for a given key in a query map.
-enum QueryResult {
+enum QueryResult<I> {
     /// An already executing query. The query job can be used to await for its completion.
-    Started(QueryJob),
+    Started(QueryJob<I>),
 
     /// The query panicked. Queries trying to wait on this will raise a fatal error which will
     /// silently panic.
     Poisoned,
 }
 
-impl QueryResult {
+impl<I> QueryResult<I> {
     /// Unwraps the query job expecting that it has started.
-    fn expect_job(self) -> QueryJob {
+    fn expect_job(self) -> QueryJob<I> {
         match self {
             Self::Started(job) => job,
             Self::Poisoned => {
@@ -55,7 +55,7 @@ impl QueryResult {
     }
 }
 
-impl<K> QueryState<K>
+impl<K, I> QueryState<K, I>
 where
     K: Eq + Hash + Copy + Debug,
 {
@@ -66,8 +66,8 @@ where
     pub fn try_collect_active_jobs<Qcx: Copy>(
         &self,
         qcx: Qcx,
-        make_query: fn(Qcx, K) -> QueryStackFrame,
-        jobs: &mut QueryMap,
+        make_query: fn(Qcx, K) -> QueryStackFrame<I>,
+        jobs: &mut QueryMap<I>,
     ) -> Option<()> {
         let mut active = Vec::new();
 
@@ -76,7 +76,7 @@ where
         for shard in self.active.try_lock_shards() {
             for (k, v) in shard?.iter() {
                 if let QueryResult::Started(ref job) = *v {
-                    active.push((*k, job.clone()));
+                    active.push((*k, (*job).clone()));
                 }
             }
         }
@@ -92,19 +92,19 @@ where
     }
 }
 
-impl<K> Default for QueryState<K> {
-    fn default() -> QueryState<K> {
+impl<K, I> Default for QueryState<K, I> {
+    fn default() -> QueryState<K, I> {
         QueryState { active: Default::default() }
     }
 }
 
 /// A type representing the responsibility to execute the job in the `job` field.
 /// This will poison the relevant query if dropped.
-struct JobOwner<'tcx, K>
+struct JobOwner<'tcx, K, I>
 where
     K: Eq + Hash + Copy,
 {
-    state: &'tcx QueryState<K>,
+    state: &'tcx QueryState<K, I>,
     key: K,
 }
 
@@ -146,7 +146,7 @@ where
         }
         Stash => {
             let guar = if let Some(root) = cycle_error.cycle.first()
-                && let Some(span) = root.query.span
+                && let Some(span) = root.query.info.span
             {
                 error.stash(span, StashKey::Cycle).unwrap()
             } else {
@@ -157,7 +157,7 @@ where
     }
 }
 
-impl<'tcx, K> JobOwner<'tcx, K>
+impl<'tcx, K, I> JobOwner<'tcx, K, I>
 where
     K: Eq + Hash + Copy,
 {
@@ -194,7 +194,7 @@ where
     }
 }
 
-impl<'tcx, K> Drop for JobOwner<'tcx, K>
+impl<'tcx, K, I> Drop for JobOwner<'tcx, K, I>
 where
     K: Eq + Hash + Copy,
 {
@@ -222,10 +222,19 @@ where
 }
 
 #[derive(Clone, Debug)]
-pub struct CycleError {
+pub struct CycleError<I = QueryStackFrameExtra> {
     /// The query and related span that uses the cycle.
-    pub usage: Option<(Span, QueryStackFrame)>,
-    pub cycle: Vec<QueryInfo>,
+    pub usage: Option<(Span, QueryStackFrame<I>)>,
+    pub cycle: Vec<QueryInfo<I>>,
+}
+
+impl<I> CycleError<I> {
+    fn lift<Qcx: QueryContext<QueryInfo = I>>(&self, qcx: Qcx) -> CycleError<QueryStackFrameExtra> {
+        CycleError {
+            usage: self.usage.as_ref().map(|(span, frame)| (*span, frame.lift(qcx))),
+            cycle: self.cycle.iter().map(|info| info.lift(qcx)).collect(),
+        }
+    }
 }
 
 /// Checks whether there is already a value for this key in the in-memory
@@ -262,10 +271,10 @@ where
 {
     // Ensure there was no errors collecting all active jobs.
     // We need the complete map to ensure we find a cycle to break.
-    let query_map = qcx.collect_active_jobs().expect("failed to collect active queries");
+    let query_map = qcx.collect_active_jobs().ok().expect("failed to collect active queries");
 
     let error = try_execute.find_cycle_in_stack(query_map, &qcx.current_query_job(), span);
-    (mk_cycle(query, qcx, error), None)
+    (mk_cycle(query, qcx, error.lift(qcx)), None)
 }
 
 #[inline(always)]
@@ -274,7 +283,7 @@ fn wait_for_query<Q, Qcx>(
     qcx: Qcx,
     span: Span,
     key: Q::Key,
-    latch: QueryLatch,
+    latch: QueryLatch<Qcx::QueryInfo>,
     current: Option<QueryJobId>,
 ) -> (Q::Value, Option<DepNodeIndex>)
 where
@@ -314,7 +323,7 @@ where
 
             (v, Some(index))
         }
-        Err(cycle) => (mk_cycle(query, qcx, cycle), None),
+        Err(cycle) => (mk_cycle(query, qcx, cycle.lift(qcx)), None),
     }
 }
 
@@ -392,7 +401,7 @@ where
 fn execute_job<Q, Qcx, const INCR: bool>(
     query: Q,
     qcx: Qcx,
-    state: &QueryState<Q::Key>,
+    state: &QueryState<Q::Key, Qcx::QueryInfo>,
     key: Q::Key,
     key_hash: u64,
     id: QueryJobId,