use crossbeam_deque::{Deque, Steal, Stealer};
#[cfg(rayon_unstable)]
use internal::task::Task;
#[cfg(rayon_unstable)]
use job::Job;
use job::{JobRef, StackJob};
use latch::{CountLatch, Latch, LatchProbe, LockLatch, SpinLatch, TickleLatch};
use log::Event::*;
use sleep::Sleep;
use std::any::Any;
use std::cell::Cell;
use std::collections::hash_map::DefaultHasher;
use std::hash::Hasher;
use std::mem;
use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
use std::sync::{Arc, Mutex, Once, ONCE_INIT};
use std::thread;
use std::usize;
use unwind;
use util::leak;
use {ErrorKind, ExitHandler, PanicHandler, DeadlockHandler, StartHandler, AcquireThreadHandler,
MainHandler, ThreadPoolBuildError, ThreadPoolBuilder, ReleaseThreadHandler};
pub struct Registry {
thread_infos: Vec<ThreadInfo>,
state: Mutex<RegistryState>,
sleep: Sleep,
job_uninjector: Stealer<JobRef>,
panic_handler: Option<Box<PanicHandler>>,
pub(crate) deadlock_handler: Option<Box<DeadlockHandler>>,
start_handler: Option<Box<StartHandler>>,
exit_handler: Option<Box<ExitHandler>>,
main_handler: Option<Box<MainHandler>>,
pub(crate) acquire_thread_handler: Option<Box<AcquireThreadHandler>>,
pub(crate) release_thread_handler: Option<Box<ReleaseThreadHandler>>,
terminate_latch: CountLatch,
}
struct RegistryState {
job_injector: Deque<JobRef>,
}
static mut THE_REGISTRY: Option<&'static Arc<Registry>> = None;
static THE_REGISTRY_SET: Once = ONCE_INIT;
fn global_registry() -> &'static Arc<Registry> {
THE_REGISTRY_SET.call_once(|| unsafe { init_registry(ThreadPoolBuilder::new()).unwrap() });
unsafe { THE_REGISTRY.expect("The global thread pool has not been initialized.") }
}
pub fn init_global_registry(
builder: ThreadPoolBuilder,
) -> Result<&'static Registry, ThreadPoolBuildError> {
let mut called = false;
let mut init_result = Ok(());;
THE_REGISTRY_SET.call_once(|| unsafe {
init_result = init_registry(builder);
called = true;
});
if called {
init_result.map(|()| &**global_registry())
} else {
Err(ThreadPoolBuildError::new(
ErrorKind::GlobalPoolAlreadyInitialized,
))
}
}
unsafe fn init_registry(builder: ThreadPoolBuilder) -> Result<(), ThreadPoolBuildError> {
Registry::new(builder).map(|registry| THE_REGISTRY = Some(leak(registry)))
}
struct Terminator<'a>(&'a Arc<Registry>);
impl<'a> Drop for Terminator<'a> {
fn drop(&mut self) {
self.0.terminate()
}
}
impl Registry {
pub fn new(mut builder: ThreadPoolBuilder) -> Result<Arc<Registry>, ThreadPoolBuildError> {
let n_threads = builder.get_num_threads();
let breadth_first = builder.get_breadth_first();
let inj_worker = Deque::new();
let inj_stealer = inj_worker.stealer();
let workers: Vec<_> = (0..n_threads).map(|_| Deque::new()).collect();
let stealers: Vec<_> = workers.iter().map(|d| d.stealer()).collect();
let registry = Arc::new(Registry {
thread_infos: stealers.into_iter().map(|s| ThreadInfo::new(s)).collect(),
state: Mutex::new(RegistryState::new(inj_worker)),
sleep: Sleep::new(n_threads),
job_uninjector: inj_stealer,
terminate_latch: CountLatch::new(),
panic_handler: builder.take_panic_handler(),
deadlock_handler: builder.take_deadlock_handler(),
start_handler: builder.take_start_handler(),
main_handler: builder.take_main_handler(),
exit_handler: builder.take_exit_handler(),
acquire_thread_handler: builder.take_acquire_thread_handler(),
release_thread_handler: builder.take_release_thread_handler(),
});
let t1000 = Terminator(®istry);
for (index, worker) in workers.into_iter().enumerate() {
let registry = registry.clone();
let mut b = thread::Builder::new();
if let Some(name) = builder.get_thread_name(index) {
b = b.name(name);
}
if let Some(stack_size) = builder.get_stack_size() {
b = b.stack_size(stack_size);
}
if let Err(e) =
b.spawn(move || unsafe { main_loop(worker, registry, index, breadth_first) })
{
return Err(ThreadPoolBuildError::new(ErrorKind::IOError(e)));
}
}
mem::forget(t1000);
Ok(registry.clone())
}
#[cfg(rayon_unstable)]
pub fn global() -> Arc<Registry> {
global_registry().clone()
}
pub fn current() -> Arc<Registry> {
unsafe {
let worker_thread = WorkerThread::current();
if worker_thread.is_null() {
global_registry().clone()
} else {
(*worker_thread).registry.clone()
}
}
}
pub fn current_num_threads() -> usize {
unsafe {
let worker_thread = WorkerThread::current();
if worker_thread.is_null() {
global_registry().num_threads()
} else {
(*worker_thread).registry.num_threads()
}
}
}
pub fn id(&self) -> RegistryId {
RegistryId {
addr: self as *const Self as usize,
}
}
pub fn num_threads(&self) -> usize {
self.thread_infos.len()
}
pub fn handle_panic(&self, err: Box<Any + Send>) {
match self.panic_handler {
Some(ref handler) => {
let abort_guard = unwind::AbortIfPanic;
handler(err);
mem::forget(abort_guard);
}
None => {
let _ = unwind::AbortIfPanic;
}
}
}
pub fn wait_until_primed(&self) {
for info in &self.thread_infos {
info.primed.wait();
}
}
pub(crate) fn wait_until_stopped(&self) {
self.release_thread();
for info in &self.thread_infos {
info.stopped.wait();
}
self.acquire_thread();
}
pub(crate) fn acquire_thread(&self) {
if let Some(ref acquire_thread_handler) = self.acquire_thread_handler {
acquire_thread_handler();
}
}
pub(crate) fn release_thread(&self) {
if let Some(ref release_thread_handler) = self.release_thread_handler {
release_thread_handler();
}
}
pub fn inject_or_push(&self, job_ref: JobRef) {
let worker_thread = WorkerThread::current();
unsafe {
if !worker_thread.is_null() && (*worker_thread).registry().id() == self.id() {
(*worker_thread).push(job_ref);
} else {
self.inject(&[job_ref]);
}
}
}
#[cfg(rayon_unstable)]
pub unsafe fn submit_task<T>(&self, task: Arc<T>)
where
T: Task,
{
let task_job = TaskJob::new(task);
let task_job_ref = TaskJob::into_job_ref(task_job);
return self.inject_or_push(task_job_ref);
struct TaskJob<T: Task> {
_data: T,
}
impl<T: Task> TaskJob<T> {
fn new(arc: Arc<T>) -> Arc<Self> {
unsafe { mem::transmute(arc) }
}
pub fn into_task(this: Arc<TaskJob<T>>) -> Arc<T> {
unsafe { mem::transmute(this) }
}
unsafe fn into_job_ref(this: Arc<Self>) -> JobRef {
let this: *const Self = mem::transmute(this);
JobRef::new(this)
}
}
impl<T: Task> Job for TaskJob<T> {
unsafe fn execute(this: *const Self) {
let this: Arc<Self> = mem::transmute(this);
let task: Arc<T> = TaskJob::into_task(this);
Task::execute(task);
}
}
}
pub fn inject(&self, injected_jobs: &[JobRef]) {
log!(InjectJobs {
count: injected_jobs.len()
});
{
let state = self.state.lock().unwrap();
assert!(
!self.terminate_latch.probe(),
"inject() sees state.terminate as true"
);
for &job_ref in injected_jobs {
state.job_injector.push(job_ref);
}
}
self.sleep.tickle(usize::MAX);
}
fn pop_injected_job(&self, worker_index: usize) -> Option<JobRef> {
loop {
match self.job_uninjector.steal() {
Steal::Empty => return None,
Steal::Data(d) => {
log!(UninjectedWork {
worker: worker_index
});
return Some(d);
}
Steal::Retry => {}
}
}
}
pub fn in_worker<OP, R>(&self, op: OP) -> R
where
OP: FnOnce(&WorkerThread, bool) -> R + Send,
R: Send,
{
unsafe {
let worker_thread = WorkerThread::current();
if worker_thread.is_null() {
self.in_worker_cold(op)
} else if (*worker_thread).registry().id() != self.id() {
self.in_worker_cross(&*worker_thread, op)
} else {
op(&*worker_thread, false)
}
}
}
#[cold]
unsafe fn in_worker_cold<OP, R>(&self, op: OP) -> R
where
OP: FnOnce(&WorkerThread, bool) -> R + Send,
R: Send,
{
debug_assert!(WorkerThread::current().is_null());
let job = StackJob::new(
0,
|injected| {
let worker_thread = WorkerThread::current();
assert!(injected && !worker_thread.is_null());
op(&*worker_thread, true)
},
LockLatch::new(),
);
self.inject(&[job.as_job_ref()]);
self.release_thread();
job.latch.wait();
self.acquire_thread();
job.into_result()
}
#[cold]
unsafe fn in_worker_cross<OP, R>(&self, current_thread: &WorkerThread, op: OP) -> R
where
OP: FnOnce(&WorkerThread, bool) -> R + Send,
R: Send,
{
debug_assert!(current_thread.registry().id() != self.id());
let latch = TickleLatch::new(SpinLatch::new(), ¤t_thread.registry().sleep);
let job = StackJob::new(
0,
|injected| {
let worker_thread = WorkerThread::current();
assert!(injected && !worker_thread.is_null());
op(&*worker_thread, true)
},
latch,
);
self.inject(&[job.as_job_ref()]);
current_thread.wait_until(&job.latch);
job.into_result()
}
pub fn increment_terminate_count(&self) {
self.terminate_latch.increment();
}
pub fn terminate(&self) {
self.terminate_latch.set();
self.sleep.tickle(usize::MAX);
}
}
#[inline]
pub fn mark_blocked() {
let worker_thread = WorkerThread::current();
assert!(!worker_thread.is_null());
unsafe {
let registry = &(*worker_thread).registry;
registry.sleep.mark_blocked(®istry.deadlock_handler)
}
}
#[inline]
pub fn mark_unblocked(registry: &Registry) {
registry.sleep.mark_unblocked()
}
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct RegistryId {
addr: usize,
}
impl RegistryState {
pub fn new(job_injector: Deque<JobRef>) -> RegistryState {
RegistryState {
job_injector: job_injector,
}
}
}
struct ThreadInfo {
primed: LockLatch,
stopped: LockLatch,
stealer: Stealer<JobRef>,
}
impl ThreadInfo {
fn new(stealer: Stealer<JobRef>) -> ThreadInfo {
ThreadInfo {
primed: LockLatch::new(),
stopped: LockLatch::new(),
stealer: stealer,
}
}
}
pub struct WorkerThread {
worker: Deque<JobRef>,
pub(crate) index: usize,
breadth_first: bool,
rng: XorShift64Star,
pub(crate) registry: Arc<Registry>,
}
thread_local! {
static WORKER_THREAD_STATE: Cell<*const WorkerThread> =
Cell::new(0 as *const WorkerThread)
}
impl WorkerThread {
#[inline]
pub fn current() -> *const WorkerThread {
WORKER_THREAD_STATE.with(|t| t.get())
}
unsafe fn set_current(thread: *const WorkerThread) {
WORKER_THREAD_STATE.with(|t| {
assert!(t.get().is_null());
t.set(thread);
});
}
pub fn registry(&self) -> &Arc<Registry> {
&self.registry
}
#[inline]
pub fn index(&self) -> usize {
self.index
}
#[inline]
pub unsafe fn push(&self, job: JobRef) {
self.worker.push(job);
self.registry.sleep.tickle(self.index);
}
#[inline]
pub fn local_deque_is_empty(&self) -> bool {
self.worker.len() == 0
}
#[inline]
pub unsafe fn take_local_job(&self) -> Option<JobRef> {
if !self.breadth_first {
self.worker.pop()
} else {
loop {
match self.worker.steal() {
Steal::Empty => return None,
Steal::Data(d) => return Some(d),
Steal::Retry => {}
}
}
}
}
#[inline]
pub unsafe fn wait_until<L: LatchProbe + ?Sized>(&self, latch: &L) {
log!(WaitUntil { worker: self.index });
if !latch.probe() {
self.wait_until_cold(latch);
}
}
#[cold]
unsafe fn wait_until_cold<L: LatchProbe + ?Sized>(&self, latch: &L) {
let abort_guard = unwind::AbortIfPanic;
let mut yields = 0;
while !latch.probe() {
if let Some(job) = self
.take_local_job()
.or_else(|| self.steal())
.or_else(|| self.registry.pop_injected_job(self.index))
{
yields = self.registry.sleep.work_found(self.index, yields);
self.execute(job);
} else {
yields = self.registry.sleep.no_work_found(
self.index,
yields,
&self.registry
);
}
}
self.registry.sleep.work_found(self.index, yields);
log!(LatchSet { worker: self.index });
mem::forget(abort_guard);
}
pub unsafe fn execute(&self, job: JobRef) {
job.execute();
self.registry.sleep.tickle(self.index);
}
unsafe fn steal(&self) -> Option<JobRef> {
debug_assert!(self.worker.pop().is_none());
let num_threads = self.registry.thread_infos.len();
if num_threads <= 1 {
return None;
}
let start = self.rng.next_usize(num_threads);
(start..num_threads)
.chain(0..start)
.filter(|&i| i != self.index)
.filter_map(|victim_index| {
let victim = &self.registry.thread_infos[victim_index];
loop {
match victim.stealer.steal() {
Steal::Empty => return None,
Steal::Data(d) => {
log!(StoleWork {
worker: self.index,
victim: victim_index
});
return Some(d);
}
Steal::Retry => {}
}
}
})
.next()
}
}
unsafe fn main_loop(
worker: Deque<JobRef>,
registry: Arc<Registry>,
index: usize,
breadth_first: bool,
) {
let worker_thread = WorkerThread {
worker: worker,
breadth_first: breadth_first,
index: index,
rng: XorShift64Star::new(),
registry: registry.clone(),
};
WorkerThread::set_current(&worker_thread);
registry.thread_infos[index].primed.set();
let abort_guard = unwind::AbortIfPanic;
if let Some(ref handler) = registry.start_handler {
let registry = registry.clone();
match unwind::halt_unwinding(|| handler(index)) {
Ok(()) => {}
Err(err) => {
registry.handle_panic(err);
}
}
}
let mut work = || {
worker_thread.wait_until(®istry.terminate_latch);
};
registry.acquire_thread();
if let Some(ref handler) = registry.main_handler {
match unwind::halt_unwinding(|| handler(index, &mut work)) {
Ok(()) => {
}
Err(err) => {
registry.handle_panic(err);
}
}
} else {
work();
}
debug_assert!(worker_thread.take_local_job().is_none());
registry.thread_infos[index].stopped.set();
mem::forget(abort_guard);
if let Some(ref handler) = registry.exit_handler {
let registry = registry.clone();
match unwind::halt_unwinding(|| handler(index)) {
Ok(()) => {}
Err(err) => {
registry.handle_panic(err);
}
}
}
registry.release_thread();
}
pub fn in_worker<OP, R>(op: OP) -> R
where
OP: FnOnce(&WorkerThread, bool) -> R + Send,
R: Send,
{
unsafe {
let owner_thread = WorkerThread::current();
if !owner_thread.is_null() {
op(&*owner_thread, false)
} else {
global_registry().in_worker_cold(op)
}
}
}
struct XorShift64Star {
state: Cell<u64>,
}
impl XorShift64Star {
fn new() -> Self {
let mut seed = 0;
while seed == 0 {
let mut hasher = DefaultHasher::new();
static COUNTER: AtomicUsize = ATOMIC_USIZE_INIT;
hasher.write_usize(COUNTER.fetch_add(1, Ordering::Relaxed));
seed = hasher.finish();
}
XorShift64Star {
state: Cell::new(seed),
}
}
fn next(&self) -> u64 {
let mut x = self.state.get();
debug_assert_ne!(x, 0);
x ^= x >> 12;
x ^= x << 25;
x ^= x >> 27;
self.state.set(x);
x.wrapping_mul(0x2545_f491_4f6c_dd1d)
}
fn next_usize(&self, n: usize) -> usize {
(self.next() % n as u64) as usize
}
}