From 401568afc48cbd55785385a1a62278e7b14a3c49 Mon Sep 17 00:00:00 2001 From: ywxt Date: Mon, 10 Nov 2025 14:41:03 +0800 Subject: [PATCH] Lock shards while collecting active jobs. Co-authored-by: Zoxc --- compiler/rustc_interface/src/util.rs | 2 +- compiler/rustc_query_impl/src/plumbing.rs | 37 ++++++++++++++++--- compiler/rustc_query_system/src/query/job.rs | 2 +- compiler/rustc_query_system/src/query/mod.rs | 4 +- .../rustc_query_system/src/query/plumbing.rs | 26 ++++++++++++- 5 files changed, 61 insertions(+), 10 deletions(-) diff --git a/compiler/rustc_interface/src/util.rs b/compiler/rustc_interface/src/util.rs index ade7ec38fb353..483db2036aefe 100644 --- a/compiler/rustc_interface/src/util.rs +++ b/compiler/rustc_interface/src/util.rs @@ -252,7 +252,7 @@ pub(crate) fn run_in_thread_pool_with_globals< 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().expect("failed to collect active queries in deadlock handler") + QueryCtxt::new(tcx).try_collect_active_jobs().expect("failed to collect active queries in deadlock handler") }); break_query_cycles(query_map, ®istry); }) diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs index 4e601a6c59442..768de9672c084 100644 --- a/compiler/rustc_query_impl/src/plumbing.rs +++ b/compiler/rustc_query_impl/src/plumbing.rs @@ -46,6 +46,15 @@ impl<'tcx> QueryCtxt<'tcx> { pub fn new(tcx: TyCtxt<'tcx>) -> Self { QueryCtxt { tcx } } + + fn collect_active_jobs(self) -> QueryMap> { + let mut jobs = QueryMap::default(); + + for collect in super::COLLECT_ACTIVE_JOBS.iter() { + collect(self.tcx, &mut jobs) + } + jobs + } } impl<'tcx> std::ops::Deref for QueryCtxt<'tcx> { @@ -91,7 +100,7 @@ impl<'tcx> QueryContext for QueryCtxt<'tcx> { /// 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( + fn try_collect_active_jobs( self, ) -> Result>, QueryMap>> { let mut jobs = QueryMap::default(); @@ -163,11 +172,7 @@ impl<'tcx> QueryContext for QueryCtxt<'tcx> { } fn depth_limit_error(self, job: QueryJobId) { - // FIXME: `collect_active_jobs` expects no locks to be held, which doesn't hold for this call. - let query_map = match self.collect_active_jobs() { - Ok(query_map) => query_map, - Err(query_map) => query_map, - }; + let query_map = self.collect_active_jobs(); let (info, depth) = job.find_dep_kind_root(query_map); let suggested_limit = match self.recursion_limit() { @@ -757,6 +762,22 @@ macro_rules! define_queries { res } + pub(crate) fn collect_active_jobs<'tcx>( + tcx: TyCtxt<'tcx>, + qmap: &mut QueryMap>, + ) { + let make_query = |tcx, key| { + let kind = rustc_middle::dep_graph::dep_kinds::$name; + let name = stringify!($name); + $crate::plumbing::create_query_frame(tcx, rustc_middle::query::descs::$name, key, kind, name) + }; + tcx.query_system.states.$name.collect_active_jobs( + tcx, + make_query, + qmap, + ); + } + pub(crate) fn alloc_self_profile_query_strings<'tcx>( tcx: TyCtxt<'tcx>, string_cache: &mut QueryKeyStringCache @@ -819,6 +840,10 @@ macro_rules! define_queries { ] = &[$(query_impl::$name::try_collect_active_jobs),*]; + const COLLECT_ACTIVE_JOBS: &[ + for<'tcx> fn(TyCtxt<'tcx>, &mut QueryMap>) + ] = &[$(query_impl::$name::collect_active_jobs),*]; + const ALLOC_SELF_PROFILE_QUERY_STRINGS: &[ for<'tcx> fn(TyCtxt<'tcx>, &mut QueryKeyStringCache) ] = &[$(query_impl::$name::alloc_self_profile_query_strings),*]; diff --git a/compiler/rustc_query_system/src/query/job.rs b/compiler/rustc_query_system/src/query/job.rs index fd1ea997ebe56..a68f5abc88390 100644 --- a/compiler/rustc_query_system/src/query/job.rs +++ b/compiler/rustc_query_system/src/query/job.rs @@ -616,7 +616,7 @@ pub fn print_query_stack( let mut count_total = 0; // Make use of a partial query map if we fail to take locks collecting active queries. - let query_map = match qcx.collect_active_jobs() { + let query_map = match qcx.try_collect_active_jobs() { Ok(query_map) => query_map, Err(query_map) => query_map, }; diff --git a/compiler/rustc_query_system/src/query/mod.rs b/compiler/rustc_query_system/src/query/mod.rs index 855769dacc3e1..7de5cac87519e 100644 --- a/compiler/rustc_query_system/src/query/mod.rs +++ b/compiler/rustc_query_system/src/query/mod.rs @@ -161,7 +161,9 @@ pub trait QueryContext: HasDepContext { /// Get the query information from the TLS context. fn current_query_job(self) -> Option; - fn collect_active_jobs(self) -> Result, QueryMap>; + fn try_collect_active_jobs( + self, + ) -> Result, QueryMap>; fn lift_query_info(self, info: &Self::QueryInfo) -> QueryStackFrameExtra; diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs index e74de5edc42d7..bb1a6d15d6796 100644 --- a/compiler/rustc_query_system/src/query/plumbing.rs +++ b/compiler/rustc_query_system/src/query/plumbing.rs @@ -90,6 +90,30 @@ where Some(()) } + + pub fn collect_active_jobs( + &self, + qcx: Qcx, + make_query: fn(Qcx, K) -> QueryStackFrame, + jobs: &mut QueryMap, + ) { + let mut active = Vec::new(); + + for shard in self.active.lock_shards() { + for (k, v) in shard.iter() { + if let QueryResult::Started(ref job) = *v { + active.push((*k, (*job).clone())); + } + } + } + + // Call `make_query` while we're not holding a `self.active` lock as `make_query` may call + // queries leading to a deadlock. + for (key, job) in active { + let query = make_query(qcx, key); + jobs.insert(job.id, QueryJobInfo { query, job }); + } + } } impl Default for QueryState { @@ -271,7 +295,7 @@ 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().ok().expect("failed to collect active queries"); + let query_map = qcx.try_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.lift(qcx)), None)