Merge #945

945: Revert "Use a single termination latch for the registry" r=cuviper a=cuviper

This reverts commit 1dfd5d9.

This was a mistake -- multiple threads can't wait on the same latch,
because the inner `CoreLatch` tracks thread-specific sleep state.


Co-authored-by: Josh Stone <[email protected]>

Author: bors[bot]

rayon-core/src/latch.rs

@@ -292,11 +292,10 @@ impl CountLatch {
         }
     }
 
-    /// Increments the latch counter by one and returns the previous value.
     #[inline]
-    pub(super) fn increment(&self) -> usize {
+    pub(super) fn increment(&self) {
         debug_assert!(!self.core_latch.probe());
-        self.counter.fetch_add(1, Ordering::Relaxed)
+        self.counter.fetch_add(1, Ordering::Relaxed);
     }
 
     /// Decrements the latch counter by one. If this is the final
@@ -345,12 +344,10 @@ impl CountLockLatch {
         }
     }
 
-    /// Increments the latch counter by one and returns the previous value.
     #[inline]
-    pub(super) fn increment(&self) -> usize {
+    pub(super) fn increment(&self) {
         let old_counter = self.counter.fetch_add(1, Ordering::Relaxed);
         debug_assert!(old_counter != 0);
-        old_counter
     }
 
     pub(super) fn wait(&self) {

rayon-core/src/registry.rs

@@ -152,7 +152,7 @@ pub(super) struct Registry {
     // - when `join()` or `scope()` is invoked, similarly, no adjustments are needed.
     //   These are always owned by some other job (e.g., one injected by `ThreadPool::install()`)
     //   and that job will keep the pool alive.
-    terminate: CountLatch,
+    terminate_count: AtomicUsize,
 }
 
 /// ////////////////////////////////////////////////////////////////////////
@@ -238,7 +238,7 @@ impl Registry {
             thread_infos: stealers.into_iter().map(ThreadInfo::new).collect(),
             sleep: Sleep::new(logger, n_threads),
             injected_jobs: Injector::new(),
-            terminate: CountLatch::new(),
+            terminate_count: AtomicUsize::new(1),
             panic_handler: builder.take_panic_handler(),
             start_handler: builder.take_start_handler(),
             exit_handler: builder.take_exit_handler(),
@@ -390,8 +390,9 @@ impl Registry {
         // drops) a `ThreadPool`; and, in that case, they cannot be
         // calling `inject()` later, since they dropped their
         // `ThreadPool`.
-        debug_assert!(
-            !self.terminate.as_core_latch().probe(),
+        debug_assert_ne!(
+            self.terminate_count.load(Ordering::Acquire),
+            0,
             "inject() sees state.terminate as true"
         );
 
@@ -522,7 +523,7 @@ impl Registry {
     /// terminate count and is responsible for invoking `terminate()`
     /// when finished.
     pub(super) fn increment_terminate_count(&self) {
-        let previous = self.terminate.increment();
+        let previous = self.terminate_count.fetch_add(1, Ordering::AcqRel);
         debug_assert!(previous != 0, "registry ref count incremented from zero");
         assert!(
             previous != std::usize::MAX,
@@ -534,9 +535,9 @@ impl Registry {
     /// dropped. The worker threads will gradually terminate, once any
     /// extant work is completed.
     pub(super) fn terminate(&self) {
-        if self.terminate.set() {
-            for i in 0..self.num_threads() {
-                self.notify_worker_latch_is_set(i);
+        if self.terminate_count.fetch_sub(1, Ordering::AcqRel) == 1 {
+            for (i, thread_info) in self.thread_infos.iter().enumerate() {
+                thread_info.terminate.set_and_tickle_one(self, i);
             }
         }
     }
@@ -562,6 +563,15 @@ struct ThreadInfo {
     /// until workers have stopped; only used for tests.
     stopped: LockLatch,
 
+    /// The latch used to signal that terminated has been requested.
+    /// This latch is *set* by the `terminate` method on the
+    /// `Registry`, once the registry's main "terminate" counter
+    /// reaches zero.
+    ///
+    /// NB. We use a `CountLatch` here because it has no lifetimes and is
+    /// meant for async use, but the count never gets higher than one.
+    terminate: CountLatch,
+
     /// the "stealer" half of the worker's deque
     stealer: Stealer<JobRef>,
 }
@@ -571,6 +581,7 @@ impl ThreadInfo {
         ThreadInfo {
             primed: LockLatch::new(),
             stopped: LockLatch::new(),
+            terminate: CountLatch::new(),
             stealer,
         }
     }
@@ -817,12 +828,12 @@ unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usiz
         }
     }
 
-    let terminate_latch = registry.terminate.as_core_latch();
+    let my_terminate_latch = &registry.thread_infos[index].terminate;
     worker_thread.log(|| ThreadStart {
         worker: index,
-        terminate_addr: terminate_latch.addr(),
+        terminate_addr: my_terminate_latch.as_core_latch().addr(),
     });
-    worker_thread.wait_until_cold(terminate_latch);
+    worker_thread.wait_until(my_terminate_latch);
 
     // Should not be any work left in our queue.
     debug_assert!(worker_thread.take_local_job().is_none());

rayon-core/src/scope/mod.rs

@@ -699,8 +699,7 @@ impl ScopeLatch {
         }
     }
 
-    /// Increments the latch counter by one and returns the previous value.
-    fn increment(&self) -> usize {
+    fn increment(&self) {
         match self {
             ScopeLatch::Stealing { latch, .. } => latch.increment(),
             ScopeLatch::Blocking { latch } => latch.increment(),