From 2bf24559252f9e4a186ea6227827c0c7c417471a Mon Sep 17 00:00:00 2001
From: Felix Rath <felixm.rath@gmail.com>
Date: Mon, 12 Aug 2024 03:16:39 +0200
Subject: [PATCH] Prevent double panic in query system, improve diagnostics

---
 compiler/rustc_query_impl/src/plumbing.rs         | 10 ++++++++--
 compiler/rustc_query_system/src/query/plumbing.rs | 11 +++++++++--
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs
index b9e700c1938..c064b2bd6c1 100644
--- a/compiler/rustc_query_impl/src/plumbing.rs
+++ b/compiler/rustc_query_impl/src/plumbing.rs
@@ -702,11 +702,17 @@ pub fn try_collect_active_jobs<'tcx>(tcx: TyCtxt<'tcx>, qmap: &mut QueryMap) {
                     let name = stringify!($name);
                     $crate::plumbing::create_query_frame(tcx, rustc_middle::query::descs::$name, key, kind, name)
                 };
-                tcx.query_system.states.$name.try_collect_active_jobs(
+                let res = tcx.query_system.states.$name.try_collect_active_jobs(
                     tcx,
                     make_query,
                     qmap,
-                ).unwrap();
+                );
+                // 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.
+                if res.is_none() {
+                    tracing::warn!("Failed to collect active jobs for query with name `{}`!", stringify!($name));
+                }
             }
 
             pub fn alloc_self_profile_query_strings<'tcx>(tcx: TyCtxt<'tcx>, string_cache: &mut QueryKeyStringCache) {
diff --git a/compiler/rustc_query_system/src/query/plumbing.rs b/compiler/rustc_query_system/src/query/plumbing.rs
index 8ef680cdb6c..6dbd6e89fe9 100644
--- a/compiler/rustc_query_system/src/query/plumbing.rs
+++ b/compiler/rustc_query_system/src/query/plumbing.rs
@@ -181,8 +181,15 @@ fn complete<C>(self, cache: &C, result: C::Value, dep_node_index: DepNodeIndex)
         cache.complete(key, result, dep_node_index);
 
         let job = {
-            let mut lock = state.active.lock_shard_by_value(&key);
-            lock.remove(&key).unwrap().expect_job()
+            let val = {
+                // don't keep the lock during the `unwrap()` of the retrieved value, or we taint the
+                // underlying shard.
+                // since unwinding also wants to look at this map, this can also prevent a double
+                // panic.
+                let mut lock = state.active.lock_shard_by_value(&key);
+                lock.remove(&key)
+            };
+            val.unwrap().expect_job()
         };
 
         job.signal_complete();