diff --git a/rayon-core/src/registry.rs b/rayon-core/src/registry.rs index 3fa462593..143c06558 100644 --- a/rayon-core/src/registry.rs +++ b/rayon-core/src/registry.rs @@ -8,11 +8,11 @@ use crate::{ Yield, }; use crossbeam_deque::{Injector, Steal, Stealer, Worker}; -use std::cell::Cell; +use std::cell::{Cell, UnsafeCell}; use std::fmt; use std::hash::{DefaultHasher, Hasher}; use std::io; -use std::mem; +use std::mem::{self, ManuallyDrop}; use std::ptr; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Once}; @@ -299,10 +299,9 @@ impl Registry { // Rather than starting a new thread, we're just taking over the current thread // *without* running the main loop, so we can still return from here. // The WorkerThread is leaked, but we never shutdown the global pool anyway. - let worker_thread = Box::into_raw(Box::new(WorkerThread::from(thread))); unsafe { - WorkerThread::set_current(worker_thread); + WorkerThread::set_current(WorkerThread::from(thread)); Latch::set(®istry.thread_infos[index].primed); } continue; @@ -668,9 +667,16 @@ pub(super) struct WorkerThread { // worker is fully unwound. Using an unsafe pointer avoids the need // for a RefCell etc. thread_local! { - static WORKER_THREAD_STATE: Cell<*const WorkerThread> = const { Cell::new(ptr::null()) }; + static WORKER_THREAD_STATE: ManuallyDrop>> = + const { ManuallyDrop::new(UnsafeCell::new(None)) }; } +const _: () = { + if size_of::>() != size_of::() { + panic!() + } +}; + impl From for WorkerThread { fn from(thread: ThreadBuilder) -> Self { Self { @@ -686,10 +692,8 @@ impl From for WorkerThread { impl Drop for WorkerThread { fn drop(&mut self) { - // Undo `set_current` WORKER_THREAD_STATE.with(|t| { - assert!(t.get().eq(&(self as *const _))); - t.set(ptr::null()); + unsafe { assert!((&*t.get()).is_none()) }; }); } } @@ -700,18 +704,45 @@ impl WorkerThread { /// anywhere on the current thread. #[inline] pub(super) fn current() -> *const WorkerThread { - WORKER_THREAD_STATE.get() + WORKER_THREAD_STATE.with(|t| unsafe { (&*t.get()).as_ref().map_or(ptr::null(), |r| r) }) } /// Sets `self` as the worker-thread index for the current thread. /// This is done during worker-thread startup. - unsafe fn set_current(thread: *const WorkerThread) { - WORKER_THREAD_STATE.with(|t| { - assert!(t.get().is_null()); - t.set(thread); + fn set_current(thread: WorkerThread) { + WORKER_THREAD_STATE.with(|t| unsafe { + let t = t.get(); + assert!((&*t).is_none()); + t.write(Some(thread)); }); } + /// Sets `self` as the worker-thread index for the current thread. + /// This is done during worker-thread startup. + fn with_current(thread: WorkerThread, f: F) -> R + where + F: FnOnce(&WorkerThread) -> R, + { + WORKER_THREAD_STATE.with(|t| { + struct Guard(*mut Option); + + impl Drop for Guard { + fn drop(&mut self) { + let t = unsafe { (&mut *self.0).take() }; + drop(t); + } + } + + unsafe { + let t = t.get(); + assert!((&*t).is_none()); + t.write(Some(thread)); + let _g = Guard(t); + f((&*t).as_ref().unwrap_unchecked()) + } + }) + } + /// Returns the registry that owns this worker thread. #[inline] pub(super) fn registry(&self) -> &Arc { @@ -910,14 +941,12 @@ impl WorkerThread { // //////////////////////////////////////////////////////////////////////// unsafe fn main_loop(thread: ThreadBuilder) { - unsafe { - let worker_thread = &WorkerThread::from(thread); - WorkerThread::set_current(worker_thread); + WorkerThread::with_current(WorkerThread::from(thread), |worker_thread| { let registry = &*worker_thread.registry; let index = worker_thread.index; // let registry know we are ready to do work - Latch::set(®istry.thread_infos[index].primed); + unsafe { Latch::set(®istry.thread_infos[index].primed) }; // Worker threads should not panic. If they do, just abort, as the // internal state of the thread pool is corrupted. Note that if @@ -929,7 +958,7 @@ unsafe fn main_loop(thread: ThreadBuilder) { registry.catch_unwind(|| handler(index)); } - worker_thread.wait_until_out_of_work(); + unsafe { worker_thread.wait_until_out_of_work() }; // Normal termination, do not abort. mem::forget(abort_guard); @@ -939,7 +968,7 @@ unsafe fn main_loop(thread: ThreadBuilder) { registry.catch_unwind(|| handler(index)); // We're already exiting the thread, there's nothing else to do. } - } + }); } /// If already in a worker-thread, just execute `op`. Otherwise,