Refactor ActiveJobGuard

compiler/rustc_query_impl/src/execution.rs

@@ -1,5 +1,5 @@
 use std::hash::Hash;
-use std::mem;
+use std::mem::ManuallyDrop;
 
 use rustc_data_structures::hash_table::{Entry, HashTable};
 use rustc_data_structures::stack::ensure_sufficient_stack;
@@ -26,17 +26,6 @@ fn equivalent_key<K: Eq, V>(k: K) -> impl Fn(&(K, V)) -> bool {
     move |x| x.0 == k
 }
 
-/// Obtains the enclosed [`QueryJob`], or panics if this query evaluation
-/// was poisoned by a panic.
-fn expect_job<'tcx>(status: ActiveKeyStatus<'tcx>) -> QueryJob<'tcx> {
-    match status {
-        ActiveKeyStatus::Started(job) => job,
-        ActiveKeyStatus::Poisoned => {
-            panic!("job for query failed to start and was poisoned")
-        }
-    }
-}
-
 pub(crate) fn all_inactive<'tcx, K>(state: &QueryState<'tcx, K>) -> bool {
     state.active.lock_shards().all(|shard| shard.is_empty())
 }
@@ -104,20 +93,6 @@ where
     Some(())
 }
 
-/// Guard object representing the responsibility to execute a query job and
-/// mark it as completed.
-///
-/// This will poison the relevant query key if it is dropped without calling
-/// [`Self::complete`].
-struct ActiveJobGuard<'tcx, K>
-where
-    K: Eq + Hash + Copy,
-{
-    state: &'tcx QueryState<'tcx, K>,
-    key: K,
-    key_hash: u64,
-}
-
 #[cold]
 #[inline(never)]
 fn mk_cycle<'tcx, C: QueryCache>(
@@ -148,39 +123,65 @@ fn mk_cycle<'tcx, C: QueryCache>(
     }
 }
 
+/// Guard object representing the responsibility to execute a query job and
+/// mark it as completed.
+///
+/// This will poison the relevant query key if it is dropped without calling
+/// [`Self::complete`].
+struct ActiveJobGuard<'tcx, K>
+where
+    K: Eq + Hash + Copy,
+{
+    state: &'tcx QueryState<'tcx, K>,
+    key: K,
+    key_hash: u64,
+}
+
 impl<'tcx, K> ActiveJobGuard<'tcx, K>
 where
     K: Eq + Hash + Copy,
 {
     /// Completes the query by updating the query cache with the `result`,
     /// signals the waiter, and forgets the guard so it won't poison the query.
-    fn complete<C>(self, cache: &C, result: C::Value, dep_node_index: DepNodeIndex)
+    fn complete<C>(self, cache: &C, value: C::Value, dep_node_index: DepNodeIndex)
     where
         C: QueryCache<Key = K>,
     {
-        // Forget ourself so our destructor won't poison the query.
-        // (Extract fields by value first to make sure we don't leak anything.)
-        let Self { state, key, key_hash }: Self = self;
-        mem::forget(self);
-
         // Mark as complete before we remove the job from the active state
         // so no other thread can re-execute this query.
-        cache.complete(key, result, dep_node_index);
-
-        let job = {
-            // 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 shard = state.active.lock_shard_by_hash(key_hash);
-            match shard.find_entry(key_hash, equivalent_key(key)) {
-                Err(_) => None,
-                Ok(occupied) => Some(occupied.remove().0.1),
+        cache.complete(self.key, value, dep_node_index);
+
+        let mut this = ManuallyDrop::new(self);
+
+        // Drop everything without poisoning the query.
+        this.drop_and_maybe_poison(/* poison */ false);
+    }
+
+    fn drop_and_maybe_poison(&mut self, poison: bool) {
+        let status = {
+            let mut shard = self.state.active.lock_shard_by_hash(self.key_hash);
+            match shard.find_entry(self.key_hash, equivalent_key(self.key)) {
+                Err(_) => {
+                    // Note: we must not panic while holding the lock, because unwinding also looks
+                    // at this map, which can result in a double panic. So drop it first.
+                    drop(shard);
+                    panic!();
+                }
+                Ok(occupied) => {
+                    let ((key, status), vacant) = occupied.remove();
+                    if poison {
+                        vacant.insert((key, ActiveKeyStatus::Poisoned));
+                    }
+                    status
+                }
             }
         };
-        let job = expect_job(job.expect("active query job entry"));
 
-        job.signal_complete();
+        // Also signal the completion of the job, so waiters will continue execution.
+        match status {
+            ActiveKeyStatus::Started(job) => job.signal_complete(),
+            ActiveKeyStatus::Poisoned => panic!(),
+        }
     }
 }
 
@@ -192,21 +193,7 @@ where
     #[cold]
     fn drop(&mut self) {
         // Poison the query so jobs waiting on it panic.
-        let Self { state, key, key_hash } = *self;
-        let job = {
-            let mut shard = state.active.lock_shard_by_hash(key_hash);
-            match shard.find_entry(key_hash, equivalent_key(key)) {
-                Err(_) => panic!(),
-                Ok(occupied) => {
-                    let ((key, value), vacant) = occupied.remove();
-                    vacant.insert((key, ActiveKeyStatus::Poisoned));
-                    expect_job(value)
-                }
-            }
-        };
-        // Also signal the completion of the job, so waiters
-        // will continue execution.
-        job.signal_complete();
+        self.drop_and_maybe_poison(/* poison */ true);
     }
 }