Merge rayon-rs#942 rayon-rs#943

bors[bot] · cuviper · web-flow · commit 76e3a74ab2fc · 2022-06-10T18:30:12.000Z
942: Remove `mem::transmute` calls from job handling r=cuviper a=cuviper

Where `JobRef` was transmuting the `execute` function pointer to hide
its argument type, we can instead pass the opaque `*const ()` and let
the callee cast the pointer type.

Where `HeapJob` was transmuting its own `Box&lt;Self&gt;`, we can instead use
`Box::into_raw` and `Box::from_raw`. It can also call its function right
out of the recovered `Box`.


943: Use a single termination latch for the registry r=cuviper a=cuviper



Co-authored-by: Josh Stone &lt;cuviper@gmail.com&gt;
diff --git a/rayon-core/src/job.rs b/rayon-core/src/job.rs
@@ -20,7 +20,7 @@ pub(super) trait Job {
     /// Unsafe: this may be called from a different thread than the one
     /// which scheduled the job, so the implementer must ensure the
     /// appropriate traits are met, whether `Send`, `Sync`, or both.
-    unsafe fn execute(this: *const Self);
+    unsafe fn execute(this: *const ());
 }
 
 /// Effectively a Job trait object. Each JobRef **must** be executed
@@ -45,17 +45,15 @@ impl JobRef {
     where
         T: Job,
     {
-        let fn_ptr: unsafe fn(*const T) = <T as Job>::execute;
-
         // erase types:
         JobRef {
             pointer: data as *const (),
-            execute_fn: mem::transmute(fn_ptr),
+            execute_fn: <T as Job>::execute,
         }
     }
 
     #[inline]
-    pub(super) unsafe fn execute(&self) {
+    pub(super) unsafe fn execute(self) {
         (self.execute_fn)(self.pointer)
     }
 }
@@ -108,18 +106,11 @@ where
     F: FnOnce(bool) -> R + Send,
     R: Send,
 {
-    unsafe fn execute(this: *const Self) {
-        fn call<R>(func: impl FnOnce(bool) -> R) -> impl FnOnce() -> R {
-            move || func(true)
-        }
-
-        let this = &*this;
+    unsafe fn execute(this: *const ()) {
+        let this = &*(this as *const Self);
         let abort = unwind::AbortIfPanic;
         let func = (*this.func.get()).take().unwrap();
-        (*this.result.get()) = match unwind::halt_unwinding(call(func)) {
-            Ok(x) => JobResult::Ok(x),
-            Err(x) => JobResult::Panic(x),
-        };
+        (*this.result.get()) = JobResult::call(func);
         this.latch.set();
         mem::forget(abort);
     }
@@ -135,40 +126,43 @@ pub(super) struct HeapJob<BODY>
 where
     BODY: FnOnce() + Send,
 {
-    job: UnsafeCell<Option<BODY>>,
+    job: BODY,
 }
 
 impl<BODY> HeapJob<BODY>
 where
     BODY: FnOnce() + Send,
 {
-    pub(super) fn new(func: BODY) -> Self {
-        HeapJob {
-            job: UnsafeCell::new(Some(func)),
-        }
+    pub(super) fn new(job: BODY) -> Self {
+        HeapJob { job }
     }
 
     /// Creates a `JobRef` from this job -- note that this hides all
     /// lifetimes, so it is up to you to ensure that this JobRef
     /// doesn't outlive any data that it closes over.
-    pub(super) unsafe fn as_job_ref(self: Box<Self>) -> JobRef {
-        let this: *const Self = mem::transmute(self);
-        JobRef::new(this)
+    pub(super) unsafe fn into_job_ref(self: Box<Self>) -> JobRef {
+        JobRef::new(Box::into_raw(self))
     }
 }
 
 impl<BODY> Job for HeapJob<BODY>
 where
     BODY: FnOnce() + Send,
 {
-    unsafe fn execute(this: *const Self) {
-        let this: Box<Self> = mem::transmute(this);
-        let job = (*this.job.get()).take().unwrap();
-        job();
+    unsafe fn execute(this: *const ()) {
+        let this = Box::from_raw(this as *mut Self);
+        (this.job)();
     }
 }
 
 impl<T> JobResult<T> {
+    fn call(func: impl FnOnce(bool) -> T) -> Self {
+        match unwind::halt_unwinding(|| func(true)) {
+            Ok(x) => JobResult::Ok(x),
+            Err(x) => JobResult::Panic(x),
+        }
+    }
+
     /// Convert the `JobResult` for a job that has finished (and hence
     /// its JobResult is populated) into its return value.
     ///
@@ -204,10 +198,11 @@ impl JobFifo {
 }
 
 impl Job for JobFifo {
-    unsafe fn execute(this: *const Self) {
+    unsafe fn execute(this: *const ()) {
         // We "execute" a queue by executing its first job, FIFO.
+        let this = &*(this as *const Self);
         loop {
-            match (*this).inner.steal() {
+            match this.inner.steal() {
                 Steal::Success(job_ref) => break job_ref.execute(),
                 Steal::Empty => panic!("FIFO is empty"),
                 Steal::Retry => {}
diff --git a/rayon-core/src/latch.rs b/rayon-core/src/latch.rs
@@ -292,10 +292,11 @@ impl CountLatch {
         }
     }
 
+    /// Increments the latch counter by one and returns the previous value.
     #[inline]
-    pub(super) fn increment(&self) {
+    pub(super) fn increment(&self) -> usize {
         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
@@ -344,10 +345,12 @@ impl CountLockLatch {
         }
     }
 
+    /// Increments the latch counter by one and returns the previous value.
     #[inline]
-    pub(super) fn increment(&self) {
+    pub(super) fn increment(&self) -> usize {
         let old_counter = self.counter.fetch_add(1, Ordering::Relaxed);
         debug_assert!(old_counter != 0);
+        old_counter
     }
 
     pub(super) fn wait(&self) {
diff --git a/rayon-core/src/registry.rs b/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_count: AtomicUsize,
+    terminate: CountLatch,
 }
 
 /// ////////////////////////////////////////////////////////////////////////
@@ -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_count: AtomicUsize::new(1),
+            terminate: CountLatch::new(),
             panic_handler: builder.take_panic_handler(),
             start_handler: builder.take_start_handler(),
             exit_handler: builder.take_exit_handler(),
@@ -390,9 +390,8 @@ impl Registry {
         // drops) a `ThreadPool`; and, in that case, they cannot be
         // calling `inject()` later, since they dropped their
         // `ThreadPool`.
-        debug_assert_ne!(
-            self.terminate_count.load(Ordering::Acquire),
-            0,
+        debug_assert!(
+            !self.terminate.as_core_latch().probe(),
             "inject() sees state.terminate as true"
         );
 
@@ -523,7 +522,7 @@ impl Registry {
     /// terminate count and is responsible for invoking `terminate()`
     /// when finished.
     pub(super) fn increment_terminate_count(&self) {
-        let previous = self.terminate_count.fetch_add(1, Ordering::AcqRel);
+        let previous = self.terminate.increment();
         debug_assert!(previous != 0, "registry ref count incremented from zero");
         assert!(
             previous != std::usize::MAX,
@@ -535,9 +534,9 @@ impl Registry {
     /// dropped. The worker threads will gradually terminate, once any
     /// extant work is completed.
     pub(super) fn terminate(&self) {
-        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);
+        if self.terminate.set() {
+            for i in 0..self.num_threads() {
+                self.notify_worker_latch_is_set(i);
             }
         }
     }
@@ -563,15 +562,6 @@ 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>,
 }
@@ -581,7 +571,6 @@ impl ThreadInfo {
         ThreadInfo {
             primed: LockLatch::new(),
             stopped: LockLatch::new(),
-            terminate: CountLatch::new(),
             stealer,
         }
     }
@@ -828,12 +817,12 @@ unsafe fn main_loop(worker: Worker<JobRef>, registry: Arc<Registry>, index: usiz
         }
     }
 
-    let my_terminate_latch = &registry.thread_infos[index].terminate;
+    let terminate_latch = registry.terminate.as_core_latch();
     worker_thread.log(|| ThreadStart {
         worker: index,
-        terminate_addr: my_terminate_latch.as_core_latch().addr(),
+        terminate_addr: terminate_latch.addr(),
     });
-    worker_thread.wait_until(my_terminate_latch);
+    worker_thread.wait_until_cold(terminate_latch);
 
     // Should not be any work left in our queue.
     debug_assert!(worker_thread.take_local_job().is_none());
diff --git a/rayon-core/src/scope/mod.rs b/rayon-core/src/scope/mod.rs
@@ -543,7 +543,7 @@ impl<'scope> Scope<'scope> {
             let job_ref = Box::new(HeapJob::new(move || {
                 self.base.execute_job(move || body(self))
             }))
-            .as_job_ref();
+            .into_job_ref();
 
             // Since `Scope` implements `Sync`, we can't be sure that we're still in a
             // thread of this pool, so we can't just push to the local worker thread.
@@ -584,7 +584,7 @@ impl<'scope> ScopeFifo<'scope> {
             let job_ref = Box::new(HeapJob::new(move || {
                 self.base.execute_job(move || body(self))
             }))
-            .as_job_ref();
+            .into_job_ref();
 
             // If we're in the pool, use our scope's private fifo for this thread to execute
             // in a locally-FIFO order.  Otherwise, just use the pool's global injector.
@@ -699,7 +699,8 @@ impl ScopeLatch {
         }
     }
 
-    fn increment(&self) {
+    /// Increments the latch counter by one and returns the previous value.
+    fn increment(&self) -> usize {
         match self {
             ScopeLatch::Stealing { latch, .. } => latch.increment(),
             ScopeLatch::Blocking { latch } => latch.increment(),
diff --git a/rayon-core/src/spawn/mod.rs b/rayon-core/src/spawn/mod.rs
@@ -103,7 +103,7 @@ where
             registry.terminate(); // (*) permit registry to terminate now
         }
     }))
-    .as_job_ref()
+    .into_job_ref()
 }
 
 /// Fires off a task into the Rayon threadpool in the "static" or

Original file line number	Diff line number	Diff line change
`@@ -292,10 +292,11 @@ impl CountLatch {`
`292`	`292`	`}`
`293`	`293`	`}`
`294`	`294`
	`295`	`+ /// Increments the latch counter by one and returns the previous value.`
`295`	`296`	`#[inline]`
`296`		`- pub(super) fn increment(&self) {`
	`297`	`+ pub(super) fn increment(&self) -> usize {`
`297`	`298`	`debug_assert!(!self.core_latch.probe());`
`298`		`- self.counter.fetch_add(1, Ordering::Relaxed);`
	`299`	`+ self.counter.fetch_add(1, Ordering::Relaxed)`
`299`	`300`	`}`
`300`	`301`
`301`	`302`	`/// Decrements the latch counter by one. If this is the final`
`@@ -344,10 +345,12 @@ impl CountLockLatch {`
`344`	`345`	`}`
`345`	`346`	`}`
`346`	`347`
	`348`	`+ /// Increments the latch counter by one and returns the previous value.`
`347`	`349`	`#[inline]`
`348`		`- pub(super) fn increment(&self) {`
	`350`	`+ pub(super) fn increment(&self) -> usize {`
`349`	`351`	`let old_counter = self.counter.fetch_add(1, Ordering::Relaxed);`
`350`	`352`	`debug_assert!(old_counter != 0);`
	`353`	`+ old_counter`
`351`	`354`	`}`
`352`	`355`
`353`	`356`	`pub(super) fn wait(&self) {`
Original file line number	Diff line number	Diff line change
`@@ -103,7 +103,7 @@ where`
`103`	`103`	`registry.terminate(); // (*) permit registry to terminate now`
`104`	`104`	`}`
`105`	`105`	`}))`
`106`		`- .as_job_ref()`
	`106`	`+ .into_job_ref()`
`107`	`107`	`}`
`108`	`108`
`109`	`109`	`/// Fires off a task into the Rayon threadpool in the "static" or`