Use pointers instead of &self in Latch::set

`Latch::set` can invalidate its own `&self`, because it releases the
owning thread to continue execution, which may then invalidate the latch
by deallocation, reuse, etc. We've known about this problem when it
comes to accessing latch fields too late, but the possibly dangling
reference was still a problem, like rust-lang/rust#55005.

The result of that was rust-lang/rust#98017, omitting the LLVM attribute
`dereferenceable` on references to `!Freeze` types -- those containing
`UnsafeCell`. However, miri's Stacked Borrows implementation is finer-
grained than that, only relaxing for the cell itself in the `!Freeze`
type. For rayon, that solves the dangling reference in atomic calls, but
remains a problem for other fields of a `Latch`.

This easiest fix for rayon is to use a raw pointer instead of `&self`.
We still end up with some temporary references for stuff like atomics,
but those should be fine with the rules above.

Author: cuviper

rayon-core/src/job.rs

@@ -112,7 +112,7 @@ where
         let abort = unwind::AbortIfPanic;
         let func = (*this.func.get()).take().unwrap();
         (*this.result.get()) = JobResult::call(func);
-        this.latch.set();
+        Latch::set(&this.latch);
         mem::forget(abort);
     }
 }

rayon-core/src/latch.rs

@@ -37,10 +37,15 @@ pub(super) trait Latch {
     ///
     /// Setting a latch triggers other threads to wake up and (in some
     /// cases) complete. This may, in turn, cause memory to be
-    /// allocated and so forth.  One must be very careful about this,
+    /// deallocated and so forth. One must be very careful about this,
     /// and it's typically better to read all the fields you will need
     /// to access *before* a latch is set!
-    fn set(&self);
+    ///
+    /// This function operates on `*const Self` instead of `&self` to allow it
+    /// to become dangling during this call. The caller must ensure that the
+    /// pointer is valid upon entry, and not invalidated during the call by any
+    /// actions other than `set` itself.
+    unsafe fn set(this: *const Self);
 }
 
 pub(super) trait AsCoreLatch {
@@ -123,8 +128,8 @@ impl CoreLatch {
     /// doing some wakeups; those are encapsulated in the surrounding
     /// latch code.
     #[inline]
-    fn set(&self) -> bool {
-        let old_state = self.state.swap(SET, Ordering::AcqRel);
+    unsafe fn set(this: *const Self) -> bool {
+        let old_state = (*this).state.swap(SET, Ordering::AcqRel);
         old_state == SLEEPING
     }
 
@@ -186,29 +191,29 @@ impl<'r> AsCoreLatch for SpinLatch<'r> {
 
 impl<'r> Latch for SpinLatch<'r> {
     #[inline]
-    fn set(&self) {
+    unsafe fn set(this: *const Self) {
         let cross_registry;
 
-        let registry: &Registry = if self.cross {
+        let registry: &Registry = if (*this).cross {
             // Ensure the registry stays alive while we notify it.
             // Otherwise, it would be possible that we set the spin
             // latch and the other thread sees it and exits, causing
             // the registry to be deallocated, all before we get a
             // chance to invoke `registry.notify_worker_latch_is_set`.
-            cross_registry = Arc::clone(self.registry);
+            cross_registry = Arc::clone((*this).registry);
             &cross_registry
         } else {
             // If this is not a "cross-registry" spin-latch, then the
             // thread which is performing `set` is itself ensuring
             // that the registry stays alive. However, that doesn't
             // include this *particular* `Arc` handle if the waiting
             // thread then exits, so we must completely dereference it.
-            self.registry
+            (*this).registry
         };
-        let target_worker_index = self.target_worker_index;
+        let target_worker_index = (*this).target_worker_index;
 
-        // NOTE: Once we `set`, the target may proceed and invalidate `&self`!
-        if self.core_latch.set() {
+        // NOTE: Once we `set`, the target may proceed and invalidate `this`!
+        if CoreLatch::set(&(*this).core_latch) {
             // Subtle: at this point, we can no longer read from
             // `self`, because the thread owning this spin latch may
             // have awoken and deallocated the latch. Therefore, we
@@ -255,10 +260,10 @@ impl LockLatch {
 
 impl Latch for LockLatch {
     #[inline]
-    fn set(&self) {
-        let mut guard = self.m.lock().unwrap();
+    unsafe fn set(this: *const Self) {
+        let mut guard = (*this).m.lock().unwrap();
         *guard = true;
-        self.v.notify_all();
+        (*this).v.notify_all();
     }
 }
 
@@ -307,9 +312,9 @@ impl CountLatch {
     /// count, then the latch is **set**, and calls to `probe()` will
     /// return true. Returns whether the latch was set.
     #[inline]
-    pub(super) fn set(&self) -> bool {
-        if self.counter.fetch_sub(1, Ordering::SeqCst) == 1 {
-            self.core_latch.set();
+    pub(super) unsafe fn set(this: *const Self) -> bool {
+        if (*this).counter.fetch_sub(1, Ordering::SeqCst) == 1 {
+            CoreLatch::set(&(*this).core_latch);
             true
         } else {
             false
@@ -320,8 +325,12 @@ impl CountLatch {
     /// the latch is set, then the specific worker thread is tickled,
     /// which should be the one that owns this latch.
     #[inline]
-    pub(super) fn set_and_tickle_one(&self, registry: &Registry, target_worker_index: usize) {
-        if self.set() {
+    pub(super) unsafe fn set_and_tickle_one(
+        this: *const Self,
+        registry: &Registry,
+        target_worker_index: usize,
+    ) {
+        if Self::set(this) {
             registry.notify_worker_latch_is_set(target_worker_index);
         }
     }
@@ -362,9 +371,9 @@ impl CountLockLatch {
 
 impl Latch for CountLockLatch {
     #[inline]
-    fn set(&self) {
-        if self.counter.fetch_sub(1, Ordering::SeqCst) == 1 {
-            self.lock_latch.set();
+    unsafe fn set(this: *const Self) {
+        if (*this).counter.fetch_sub(1, Ordering::SeqCst) == 1 {
+            LockLatch::set(&(*this).lock_latch);
         }
     }
 }
@@ -374,7 +383,7 @@ where
     L: Latch,
 {
     #[inline]
-    fn set(&self) {
-        L::set(self);
+    unsafe fn set(this: *const Self) {
+        L::set(&**this);
     }
 }

rayon-core/src/registry.rs

@@ -575,7 +575,7 @@ impl Registry {
     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);
+                unsafe { CountLatch::set_and_tickle_one(&thread_info.terminate, self, i) };
             }
         }
     }
@@ -869,7 +869,7 @@ unsafe fn main_loop(
     let registry = &*worker_thread.registry;
 
     // let registry know we are ready to do work
-    registry.thread_infos[index].primed.set();
+    Latch::set(&registry.thread_infos[index].primed);
 
     // Worker threads should not panic. If they do, just abort, as the
     // internal state of the threadpool is corrupted. Note that if
@@ -892,7 +892,7 @@ unsafe fn main_loop(
     debug_assert!(worker_thread.take_local_job().is_none());
 
     // let registry know we are done
-    registry.thread_infos[index].stopped.set();
+    Latch::set(&registry.thread_infos[index].stopped);
 
     // Normal termination, do not abort.
     mem::forget(abort_guard);

rayon-core/src/scope/mod.rs

@@ -540,10 +540,10 @@ impl<'scope> Scope<'scope> {
         BODY: FnOnce(&Scope<'scope>) + Send + 'scope,
     {
         let scope_ptr = ScopePtr(self);
-        let job = HeapJob::new(move || {
+        let job = HeapJob::new(move || unsafe {
             // SAFETY: this job will execute before the scope ends.
-            let scope = unsafe { scope_ptr.as_ref() };
-            scope.base.execute_job(move || body(scope))
+            let scope = scope_ptr.as_ref();
+            ScopeBase::execute_job(&scope.base, move || body(scope))
         });
         let job_ref = self.base.heap_job_ref(job);
 
@@ -562,12 +562,12 @@ impl<'scope> Scope<'scope> {
         BODY: Fn(&Scope<'scope>, BroadcastContext<'_>) + Send + Sync + 'scope,
     {
         let scope_ptr = ScopePtr(self);
-        let job = ArcJob::new(move || {
+        let job = ArcJob::new(move || unsafe {
             // SAFETY: this job will execute before the scope ends.
-            let scope = unsafe { scope_ptr.as_ref() };
+            let scope = scope_ptr.as_ref();
             let body = &body;
             let func = move || BroadcastContext::with(move |ctx| body(scope, ctx));
-            scope.base.execute_job(func);
+            ScopeBase::execute_job(&scope.base, func)
         });
         self.base.inject_broadcast(job)
     }
@@ -600,10 +600,10 @@ impl<'scope> ScopeFifo<'scope> {
         BODY: FnOnce(&ScopeFifo<'scope>) + Send + 'scope,
     {
         let scope_ptr = ScopePtr(self);
-        let job = HeapJob::new(move || {
+        let job = HeapJob::new(move || unsafe {
             // SAFETY: this job will execute before the scope ends.
-            let scope = unsafe { scope_ptr.as_ref() };
-            scope.base.execute_job(move || body(scope))
+            let scope = scope_ptr.as_ref();
+            ScopeBase::execute_job(&scope.base, move || body(scope))
         });
         let job_ref = self.base.heap_job_ref(job);
 
@@ -628,12 +628,12 @@ impl<'scope> ScopeFifo<'scope> {
         BODY: Fn(&ScopeFifo<'scope>, BroadcastContext<'_>) + Send + Sync + 'scope,
     {
         let scope_ptr = ScopePtr(self);
-        let job = ArcJob::new(move || {
+        let job = ArcJob::new(move || unsafe {
             // SAFETY: this job will execute before the scope ends.
-            let scope = unsafe { scope_ptr.as_ref() };
+            let scope = scope_ptr.as_ref();
             let body = &body;
             let func = move || BroadcastContext::with(move |ctx| body(scope, ctx));
-            scope.base.execute_job(func);
+            ScopeBase::execute_job(&scope.base, func)
         });
         self.base.inject_broadcast(job)
     }
@@ -688,36 +688,36 @@ impl<'scope> ScopeBase<'scope> {
     where
         FUNC: FnOnce() -> R,
     {
-        let result = self.execute_job_closure(func);
+        let result = unsafe { Self::execute_job_closure(self, func) };
         self.job_completed_latch.wait(owner);
         self.maybe_propagate_panic();
         result.unwrap() // only None if `op` panicked, and that would have been propagated
     }
 
     /// Executes `func` as a job, either aborting or executing as
     /// appropriate.
-    fn execute_job<FUNC>(&self, func: FUNC)
+    unsafe fn execute_job<FUNC>(this: *const Self, func: FUNC)
     where
         FUNC: FnOnce(),
     {
-        let _: Option<()> = self.execute_job_closure(func);
+        let _: Option<()> = Self::execute_job_closure(this, func);
     }
 
     /// Executes `func` as a job in scope. Adjusts the "job completed"
     /// counters and also catches any panic and stores it into
     /// `scope`.
-    fn execute_job_closure<FUNC, R>(&self, func: FUNC) -> Option<R>
+    unsafe fn execute_job_closure<FUNC, R>(this: *const Self, func: FUNC) -> Option<R>
     where
         FUNC: FnOnce() -> R,
     {
         match unwind::halt_unwinding(func) {
             Ok(r) => {
-                self.job_completed_latch.set();
+                Latch::set(&(*this).job_completed_latch);
                 Some(r)
             }
             Err(err) => {
-                self.job_panicked(err);
-                self.job_completed_latch.set();
+                (*this).job_panicked(err);
+                Latch::set(&(*this).job_completed_latch);
                 None
             }
         }
@@ -797,14 +797,14 @@ impl ScopeLatch {
 }
 
 impl Latch for ScopeLatch {
-    fn set(&self) {
-        match self {
+    unsafe fn set(this: *const Self) {
+        match &*this {
             ScopeLatch::Stealing {
                 latch,
                 registry,
                 worker_index,
-            } => latch.set_and_tickle_one(registry, *worker_index),
-            ScopeLatch::Blocking { latch } => latch.set(),
+            } => CountLatch::set_and_tickle_one(latch, registry, *worker_index),
+            ScopeLatch::Blocking { latch } => Latch::set(latch),
         }
     }
 }