1 files changed, 10 insertions, 7 deletions
diff --git a/compiler/rustc_thread_pool/src/latch.rs b/compiler/rustc_thread_pool/src/latch.rs
index 49ba62d3bea..18d654d9f78 100644
--- a/compiler/rustc_thread_pool/src/latch.rs
+++ b/compiler/rustc_thread_pool/src/latch.rs
@@ -3,6 +3,7 @@ use std::ops::Deref;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Condvar, Mutex};
 
+use crate::job::JobRef;
 use crate::registry::{Registry, WorkerThread};
 
 /// We define various kinds of latches, which are all a primitive signaling
@@ -166,11 +167,6 @@ impl<'r> SpinLatch<'r> {
     pub(super) fn cross(thread: &'r WorkerThread) -> SpinLatch<'r> {
         SpinLatch { cross: true, ..SpinLatch::new(thread) }
     }
-
-    #[inline]
-    pub(super) fn probe(&self) -> bool {
-        self.core_latch.probe()
-    }
 }
 
 impl<'r> AsCoreLatch for SpinLatch<'r> {
@@ -368,13 +364,20 @@ impl CountLatch {
         debug_assert!(old_counter != 0);
     }
 
-    pub(super) fn wait(&self, owner: Option<&WorkerThread>) {
+    pub(super) fn wait(
+        &self,
+        owner: Option<&WorkerThread>,
+        all_jobs_started: impl FnMut() -> bool,
+        is_job: impl FnMut(&JobRef) -> bool,
+    ) {
         match &self.kind {
             CountLatchKind::Stealing { latch, registry, worker_index } => unsafe {
                 let owner = owner.expect("owner thread");
                 debug_assert_eq!(registry.id(), owner.registry().id());
                 debug_assert_eq!(*worker_index, owner.index());
-                owner.wait_until(latch);
+                owner.wait_for_jobs::<_, true>(latch, all_jobs_started, is_job, |job| {
+                    owner.execute(job);
+                });
             },
             CountLatchKind::Blocking { latch } => latch.wait(),
         }