#![doc(html_root_url = "https://docs.rs/tokio-threadpool/0.1.1")]
#![deny(warnings, missing_docs, missing_debug_implementations)]
extern crate tokio_executor;
extern crate futures;
extern crate crossbeam_deque as deque;
extern crate num_cpus;
extern crate rand;
#[macro_use]
extern crate log;
#[cfg(feature = "unstable-futures")]
extern crate futures2;
mod task;
use tokio_executor::{Enter, SpawnError};
use task::Task;
use futures::{future, Future, Poll, Async};
use futures::executor::Notify;
use futures::task::AtomicTask;
use rand::{Rng, SeedableRng, XorShiftRng};
use std::{fmt, mem, thread, usize};
use std::cell::{Cell, UnsafeCell};
use std::marker::PhantomData;
use std::rc::Rc;
use std::sync::{Arc, Weak, Mutex, Condvar};
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{AcqRel, Acquire, Release, Relaxed};
use std::time::{Instant, Duration};
#[derive(Debug)]
struct ShutdownTask {
task1: AtomicTask,
#[cfg(feature = "unstable-futures")]
task2: futures2::task::AtomicWaker,
}
#[derive(Debug)]
pub struct ThreadPool {
inner: Option<Sender>,
}
#[derive(Debug)]
pub struct Sender {
inner: Arc<Inner>,
}
#[derive(Debug)]
pub struct Shutdown {
inner: Sender,
}
#[derive(Debug)]
pub struct Builder {
config: Config,
pool_size: usize,
}
#[derive(Debug, Clone)]
struct Config {
keep_alive: Option<Duration>,
name_prefix: Option<String>,
stack_size: Option<usize>,
around_worker: Option<Callback>,
}
#[derive(Debug)]
struct Inner {
state: AtomicUsize,
sleep_stack: AtomicUsize,
num_workers: AtomicUsize,
next_thread_id: AtomicUsize,
workers: Box<[WorkerEntry]>,
shutdown_task: ShutdownTask,
config: Config,
}
#[derive(Clone)]
struct Callback {
f: Arc<Fn(&Worker, &mut Enter) + Send + Sync>,
}
#[derive(Debug)]
struct Notifier {
inner: Weak<Inner>,
}
#[cfg(feature = "unstable-futures")]
struct Futures2Wake {
notifier: Arc<Notifier>,
id: usize,
}
#[derive(Eq, PartialEq, Clone, Copy)]
struct State(usize);
const SHUTDOWN_ON_IDLE: usize = 1;
const SHUTDOWN_NOW: usize = 2;
const LIFECYCLE_MASK: usize = 0b11;
const NUM_FUTURES_MASK: usize = !LIFECYCLE_MASK;
const NUM_FUTURES_OFFSET: usize = 2;
const MAX_FUTURES: usize = usize::MAX >> NUM_FUTURES_OFFSET;
#[derive(Eq, PartialEq, Clone, Copy)]
struct SleepStack(usize);
const STACK_MASK: usize = ((1 << 16) - 1);
const MAX_WORKERS: usize = 1 << 15;
const EMPTY: usize = MAX_WORKERS;
const TERMINATED: usize = EMPTY + 1;
const ABA_GUARD_SHIFT: usize = 16;
#[cfg(target_pointer_width = "64")]
const ABA_GUARD_MASK: usize = (1 << (64 - ABA_GUARD_SHIFT)) - 1;
#[cfg(target_pointer_width = "32")]
const ABA_GUARD_MASK: usize = (1 << (32 - ABA_GUARD_SHIFT)) - 1;
#[derive(Debug)]
struct WorkerEntry {
state: AtomicUsize,
next_sleeper: UnsafeCell<usize>,
deque: deque::Deque<Task>,
steal: deque::Stealer<Task>,
park_mutex: Mutex<()>,
park_condvar: Condvar,
inbound: task::Queue,
}
#[derive(Clone, Copy, Eq, PartialEq)]
struct WorkerState(usize);
const PUSHED_MASK: usize = 0b001;
const WORKER_LIFECYCLE_MASK: usize = 0b1110;
const WORKER_LIFECYCLE_SHIFT: usize = 1;
const WORKER_SHUTDOWN: usize = 0;
const WORKER_RUNNING: usize = 1;
const WORKER_SLEEPING: usize = 2;
const WORKER_NOTIFIED: usize = 3;
const WORKER_SIGNALED: usize = 4;
#[derive(Debug)]
pub struct Worker {
inner: Arc<Inner>,
idx: usize,
should_finalize: Cell<bool>,
_p: PhantomData<Rc<()>>,
}
thread_local!(static CURRENT_WORKER: Cell<*const Worker> = Cell::new(0 as *const _));
impl Builder {
pub fn new() -> Builder {
let num_cpus = num_cpus::get();
Builder {
pool_size: num_cpus,
config: Config {
keep_alive: None,
name_prefix: None,
stack_size: None,
around_worker: None,
},
}
}
pub fn pool_size(&mut self, val: usize) -> &mut Self {
assert!(val >= 1, "at least one thread required");
assert!(val <= MAX_WORKERS, "max value is {}", 32768);
self.pool_size = val;
self
}
pub fn keep_alive(&mut self, val: Option<Duration>) -> &mut Self {
self.config.keep_alive = val;
self
}
pub fn name_prefix<S: Into<String>>(&mut self, val: S) -> &mut Self {
self.config.name_prefix = Some(val.into());
self
}
pub fn stack_size(&mut self, val: usize) -> &mut Self {
self.config.stack_size = Some(val);
self
}
pub fn around_worker<F>(&mut self, f: F) -> &mut Self
where F: Fn(&Worker, &mut Enter) + Send + Sync + 'static
{
self.config.around_worker = Some(Callback::new(f));
self
}
pub fn build(&self) -> ThreadPool {
let mut workers = vec![];
trace!("build; num-workers={}", self.pool_size);
for _ in 0..self.pool_size {
workers.push(WorkerEntry::new());
}
let inner = Arc::new(Inner {
state: AtomicUsize::new(State::new().into()),
sleep_stack: AtomicUsize::new(SleepStack::new().into()),
num_workers: AtomicUsize::new(self.pool_size),
next_thread_id: AtomicUsize::new(0),
workers: workers.into_boxed_slice(),
shutdown_task: ShutdownTask {
task1: AtomicTask::new(),
#[cfg(feature = "unstable-futures")]
task2: futures2::task::AtomicWaker::new(),
},
config: self.config.clone(),
});
for i in 0..self.pool_size {
inner.push_sleeper(i).unwrap();
}
let inner = Some(Sender { inner });
ThreadPool { inner }
}
}
impl ThreadPool {
pub fn new() -> ThreadPool {
Builder::new().build()
}
pub fn spawn<F>(&self, future: F)
where F: Future<Item = (), Error = ()> + Send + 'static,
{
self.sender().spawn(future).unwrap();
}
pub fn sender(&self) -> &Sender {
self.inner.as_ref().unwrap()
}
pub fn sender_mut(&mut self) -> &mut Sender {
self.inner.as_mut().unwrap()
}
pub fn shutdown_on_idle(mut self) -> Shutdown {
self.inner().shutdown(false, false);
Shutdown { inner: self.inner.take().unwrap() }
}
pub fn shutdown(mut self) -> Shutdown {
self.inner().shutdown(true, false);
Shutdown { inner: self.inner.take().unwrap() }
}
pub fn shutdown_now(mut self) -> Shutdown {
self.inner().shutdown(true, true);
Shutdown { inner: self.inner.take().unwrap() }
}
fn inner(&self) -> &Inner {
&*self.inner.as_ref().unwrap().inner
}
}
impl Drop for ThreadPool {
fn drop(&mut self) {
if let Some(sender) = self.inner.take() {
sender.inner.shutdown(true, true);
let shutdown = Shutdown { inner: sender };
let _ = shutdown.wait();
}
}
}
impl Sender {
pub fn spawn<F>(&self, future: F) -> Result<(), SpawnError>
where F: Future<Item = (), Error = ()> + Send + 'static,
{
let mut s = self;
tokio_executor::Executor::spawn(&mut s, Box::new(future))
}
fn prepare_for_spawn(&self) -> Result<(), SpawnError> {
let mut state: State = self.inner.state.load(Acquire).into();
loop {
let mut next = state;
if next.num_futures() == MAX_FUTURES {
return Err(SpawnError::at_capacity());
}
if next.lifecycle() == SHUTDOWN_NOW {
return Err(SpawnError::shutdown());
}
next.inc_num_futures();
let actual = self.inner.state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
trace!("execute; count={:?}", next.num_futures());
break;
}
state = actual;
}
Ok(())
}
}
impl tokio_executor::Executor for Sender {
fn status(&self) -> Result<(), tokio_executor::SpawnError> {
let s = self;
tokio_executor::Executor::status(&s)
}
fn spawn(&mut self, future: Box<Future<Item = (), Error = ()> + Send>)
-> Result<(), SpawnError>
{
let mut s = &*self;
tokio_executor::Executor::spawn(&mut s, future)
}
#[cfg(feature = "unstable-futures")]
fn spawn2(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
futures2::executor::Executor::spawn(self, f)
}
}
impl<'a> tokio_executor::Executor for &'a Sender {
fn status(&self) -> Result<(), tokio_executor::SpawnError> {
let state: State = self.inner.state.load(Acquire).into();
if state.num_futures() == MAX_FUTURES {
return Err(SpawnError::at_capacity());
}
if state.lifecycle() == SHUTDOWN_NOW {
return Err(SpawnError::shutdown());
}
Ok(())
}
fn spawn(&mut self, future: Box<Future<Item = (), Error = ()> + Send>)
-> Result<(), SpawnError>
{
self.prepare_for_spawn()?;
let task = Task::new(future);
self.inner.submit(task, &self.inner);
Ok(())
}
#[cfg(feature = "unstable-futures")]
fn spawn2(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
futures2::executor::Executor::spawn(self, f)
}
}
impl<T> future::Executor<T> for Sender
where T: Future<Item = (), Error = ()> + Send + 'static,
{
fn execute(&self, future: T) -> Result<(), future::ExecuteError<T>> {
if let Err(e) = tokio_executor::Executor::status(self) {
let kind = if e.is_at_capacity() {
future::ExecuteErrorKind::NoCapacity
} else {
future::ExecuteErrorKind::Shutdown
};
return Err(future::ExecuteError::new(kind, future));
}
let _ = self.spawn(future);
Ok(())
}
}
#[cfg(feature = "unstable-futures")]
type Task2 = Box<futures2::Future<Item = (), Error = futures2::Never> + Send>;
#[cfg(feature = "unstable-futures")]
impl futures2::executor::Executor for Sender {
fn spawn(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
let mut s = &*self;
futures2::executor::Executor::spawn(&mut s, f)
}
fn status(&self) -> Result<(), futures2::executor::SpawnError> {
let s = &*self;
futures2::executor::Executor::status(&s)
}
}
#[cfg(feature = "unstable-futures")]
impl<'a> futures2::executor::Executor for &'a Sender {
fn spawn(&mut self, f: Task2) -> Result<(), futures2::executor::SpawnError> {
self.prepare_for_spawn()
.map_err(|_| futures2::executor::SpawnError::shutdown())?;
let task = Task::new2(f, |id| into_waker(Arc::new(Futures2Wake::new(id, &self.inner))));
self.inner.submit(task, &self.inner);
Ok(())
}
fn status(&self) -> Result<(), futures2::executor::SpawnError> {
tokio_executor::Executor::status(self)
.map_err(|_| futures2::executor::SpawnError::shutdown())
}
}
impl Clone for Sender {
#[inline]
fn clone(&self) -> Sender {
let inner = self.inner.clone();
Sender { inner }
}
}
impl ShutdownTask {
#[cfg(not(feature = "unstable-futures"))]
fn notify(&self) {
self.task1.notify();
}
#[cfg(feature = "unstable-futures")]
fn notify(&self) {
self.task1.notify();
self.task2.wake();
}
}
impl Shutdown {
fn inner(&self) -> &Inner {
&*self.inner.inner
}
}
impl Future for Shutdown {
type Item = ();
type Error = ();
fn poll(&mut self) -> Poll<(), ()> {
use futures::task;
trace!("Shutdown::poll");
self.inner().shutdown_task.task1.register_task(task::current());
if 0 != self.inner().num_workers.load(Acquire) {
return Ok(Async::NotReady);
}
Ok(().into())
}
}
#[cfg(feature = "unstable-futures")]
impl futures2::Future for Shutdown {
type Item = ();
type Error = ();
fn poll(&mut self, cx: &mut futures2::task::Context) -> futures2::Poll<(), ()> {
trace!("Shutdown::poll");
self.inner().shutdown_task.task2.register(cx.waker());
if 0 != self.inner().num_workers.load(Acquire) {
return Ok(futures2::Async::Pending);
}
Ok(().into())
}
}
impl Inner {
fn shutdown(&self, now: bool, purge_queue: bool) {
let mut state: State = self.state.load(Acquire).into();
trace!("shutdown; state={:?}", state);
debug_assert!(!purge_queue || now);
loop {
let mut next = state;
let num_futures = next.num_futures();
if next.lifecycle() >= SHUTDOWN_NOW {
if !purge_queue || num_futures == 0 {
return;
}
debug_assert!(purge_queue);
next.clear_num_futures();
} else {
next.set_lifecycle(if now || num_futures == 0 {
SHUTDOWN_NOW
} else {
SHUTDOWN_ON_IDLE
});
if purge_queue {
next.clear_num_futures();
}
}
let actual = self.state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if state == actual {
state = next;
break;
}
state = actual;
}
trace!(" -> transitioned to shutdown");
if state.num_futures() != 0 {
return;
}
self.terminate_sleeping_workers();
}
fn terminate_sleeping_workers(&self) {
trace!(" -> shutting down workers");
while let Some((idx, worker_state)) = self.pop_sleeper(WORKER_SIGNALED, TERMINATED) {
trace!(" -> shutdown worker; idx={:?}; state={:?}", idx, worker_state);
self.signal_stop(idx, worker_state);
}
}
fn signal_stop(&self, idx: usize, mut state: WorkerState) {
let worker = &self.workers[idx];
loop {
let mut next = state;
match state.lifecycle() {
WORKER_SHUTDOWN => {
trace!("signal_stop -- WORKER_SHUTDOWN; idx={}", idx);
self.worker_terminated();
return;
}
WORKER_RUNNING | WORKER_SLEEPING => {}
_ => {
trace!("signal_stop -- skipping; idx={}; state={:?}", idx, state);
return;
}
}
next.set_lifecycle(WORKER_SIGNALED);
let actual = worker.state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
break;
}
state = actual;
}
worker.wakeup();
}
fn worker_terminated(&self) {
let prev = self.num_workers.fetch_sub(1, AcqRel);
trace!("worker_terminated; num_workers={}", prev - 1);
if 1 == prev {
trace!("notifying shutdown task");
self.shutdown_task.notify();
}
}
fn submit(&self, task: Task, inner: &Arc<Inner>) {
Worker::with_current(|worker| {
match worker {
Some(worker) => {
let idx = worker.idx;
trace!(" -> submit internal; idx={}", idx);
worker.inner.workers[idx].submit_internal(task);
worker.inner.signal_work(inner);
}
None => {
self.submit_external(task, inner);
}
}
});
}
fn submit_external(&self, task: Task, inner: &Arc<Inner>) {
if let Some((idx, state)) = self.pop_sleeper(WORKER_NOTIFIED, EMPTY) {
trace!("submit to existing worker; idx={}; state={:?}", idx, state);
self.submit_to_external(idx, task, state, inner);
return;
}
let len = self.workers.len();
let idx = self.rand_usize() % len;
trace!(" -> submitting to random; idx={}", idx);
let state: WorkerState = self.workers[idx].state.load(Acquire).into();
self.submit_to_external(idx, task, state, inner);
}
fn submit_to_external(&self,
idx: usize,
task: Task,
state: WorkerState,
inner: &Arc<Inner>)
{
let entry = &self.workers[idx];
if !entry.submit_external(task, state) {
Worker::spawn(idx, inner);
}
}
fn signal_work(&self, inner: &Arc<Inner>) {
if let Some((idx, mut state)) = self.pop_sleeper(WORKER_SIGNALED, EMPTY) {
let entry = &self.workers[idx];
loop {
let mut next = state;
debug_assert!(state.lifecycle() != WORKER_SIGNALED);
next.set_lifecycle(WORKER_SIGNALED);
let actual = entry.state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
break;
}
state = actual;
}
match state.lifecycle() {
WORKER_SLEEPING => {
trace!("signal_work -- wakeup; idx={}", idx);
self.workers[idx].wakeup();
}
WORKER_SHUTDOWN => {
trace!("signal_work -- spawn; idx={}", idx);
Worker::spawn(idx, inner);
}
_ => {}
}
}
}
fn push_sleeper(&self, idx: usize) -> Result<(), ()> {
let mut state: SleepStack = self.sleep_stack.load(Acquire).into();
debug_assert!(WorkerState::from(self.workers[idx].state.load(Relaxed)).is_pushed());
loop {
let mut next = state;
let head = state.head();
if head == TERMINATED {
return Err(());
}
self.workers[idx].set_next_sleeper(head);
next.set_head(idx);
let actual = self.sleep_stack.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if state == actual {
return Ok(());
}
state = actual;
}
}
fn pop_sleeper(&self, max_lifecycle: usize, terminal: usize)
-> Option<(usize, WorkerState)>
{
debug_assert!(terminal == EMPTY || terminal == TERMINATED);
let mut state: SleepStack = self.sleep_stack.load(Acquire).into();
loop {
let head = state.head();
if head == EMPTY {
let mut next = state;
next.set_head(terminal);
if next == state {
debug_assert!(terminal == EMPTY);
return None;
}
let actual = self.sleep_stack.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual != state {
state = actual;
continue;
}
return None;
} else if head == TERMINATED {
return None;
}
debug_assert!(head < MAX_WORKERS);
let mut next = state;
let next_head = self.workers[head].next_sleeper();
debug_assert!(next_head != TERMINATED);
if next_head == EMPTY {
next.set_head(terminal);
} else {
next.set_head(next_head);
}
let actual = self.sleep_stack.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
debug_assert_eq!(1, PUSHED_MASK);
let state: WorkerState = self.workers[head].state
.fetch_sub(PUSHED_MASK, Release).into();
if state.lifecycle() >= max_lifecycle {
continue;
}
return Some((head, state));
}
state = actual;
}
}
fn rand_usize(&self) -> usize {
thread_local!(static THREAD_RNG_KEY: UnsafeCell<Option<XorShiftRng>> = UnsafeCell::new(None));
THREAD_RNG_KEY.with(|t| {
#[cfg(target_pointer_width = "32")]
fn new_rng(thread_id: usize) -> XorShiftRng {
XorShiftRng::from_seed([
thread_id as u32,
0x00000000,
0xa8a7d469,
0x97830e05])
}
#[cfg(target_pointer_width = "64")]
fn new_rng(thread_id: usize) -> XorShiftRng {
XorShiftRng::from_seed([
thread_id as u32,
(thread_id >> 32) as u32,
0xa8a7d469,
0x97830e05])
}
let thread_id = self.next_thread_id.fetch_add(1, Relaxed);
let rng = unsafe { &mut *t.get() };
if rng.is_none() {
*rng = Some(new_rng(thread_id));
}
rng.as_mut().unwrap().next_u32() as usize
})
}
}
impl Notify for Notifier {
fn notify(&self, id: usize) {
trace!("Notifier::notify; id=0x{:x}", id);
let id = id as usize;
let task = unsafe { Task::from_notify_id_ref(&id) };
if !task.schedule() {
trace!(" -> task already scheduled");
return;
}
let task = task.clone();
if let Some(inner) = self.inner.upgrade() {
let _ = inner.submit(task, &inner);
}
}
fn clone_id(&self, id: usize) -> usize {
unsafe {
let handle = Task::from_notify_id_ref(&id);
mem::forget(handle.clone());
}
id
}
fn drop_id(&self, id: usize) {
unsafe {
let _ = Task::from_notify_id(id);
}
}
}
unsafe impl Send for Inner {}
unsafe impl Sync for Inner {}
impl Worker {
fn spawn(idx: usize, inner: &Arc<Inner>) {
trace!("spawning new worker thread; idx={}", idx);
let mut th = thread::Builder::new();
if let Some(ref prefix) = inner.config.name_prefix {
th = th.name(format!("{}{}", prefix, idx));
}
if let Some(stack) = inner.config.stack_size {
th = th.stack_size(stack);
}
let inner = inner.clone();
th.spawn(move || {
let worker = Worker {
inner: inner,
idx: idx,
should_finalize: Cell::new(false),
_p: PhantomData,
};
let wref = &worker;
CURRENT_WORKER.with(|c| {
c.set(wref as *const _);
let inner = wref.inner.clone();
let mut sender = Sender { inner };
let mut enter = tokio_executor::enter().unwrap();
tokio_executor::with_default(&mut sender, &mut enter, |enter| {
if let Some(ref callback) = wref.inner.config.around_worker {
callback.call(wref, enter);
} else {
wref.run();
}
});
});
}).unwrap();
}
fn with_current<F: FnOnce(Option<&Worker>) -> R, R>(f: F) -> R {
CURRENT_WORKER.with(move |c| {
let ptr = c.get();
if ptr.is_null() {
f(None)
} else {
f(Some(unsafe { &*ptr }))
}
})
}
pub fn run(&self) {
let notify = Arc::new(Notifier {
inner: Arc::downgrade(&self.inner),
});
let mut sender = Sender { inner: self.inner.clone() };
let mut first = true;
let mut spin_cnt = 0;
while self.check_run_state(first) {
first = false;
let consistent = self.drain_inbound();
if self.try_run_task(¬ify, &mut sender) {
spin_cnt = 0;
continue;
}
if self.try_steal_task(¬ify, &mut sender) {
spin_cnt = 0;
continue;
}
if !consistent {
spin_cnt = 0;
continue;
}
if spin_cnt < 32 {
spin_cnt += 1;
} else if spin_cnt < 256 {
spin_cnt += 1;
thread::yield_now();
} else {
if !self.sleep() {
return;
}
}
}
self.should_finalize.set(true);
}
#[inline]
fn check_run_state(&self, first: bool) -> bool {
let mut state: WorkerState = self.entry().state.load(Acquire).into();
loop {
let pool_state: State = self.inner.state.load(Acquire).into();
if pool_state.is_terminated() {
return false;
}
let mut next = state;
match state.lifecycle() {
WORKER_RUNNING => break,
WORKER_NOTIFIED | WORKER_SIGNALED => {
next.set_lifecycle(WORKER_RUNNING);
}
lifecycle => panic!("unexpected worker state; lifecycle={}", lifecycle),
}
let actual = self.entry().state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
break;
}
state = actual;
}
if !first && state.is_signaled() {
trace!("Worker::check_run_state; delegate signal");
self.inner.signal_work(&self.inner);
}
true
}
#[inline]
fn try_run_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
use deque::Steal::*;
match self.entry().deque.steal() {
Data(task) => {
self.run_task(task, notify, sender);
true
}
Empty => false,
Retry => true,
}
}
#[inline]
fn try_steal_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
use deque::Steal::*;
let len = self.inner.workers.len();
let mut idx = self.inner.rand_usize() % len;
let mut found_work = false;
let start = idx;
loop {
if idx < len {
match self.inner.workers[idx].steal.steal() {
Data(task) => {
trace!("stole task");
self.run_task(task, notify, sender);
trace!("try_steal_task -- signal_work; self={}; from={}",
self.idx, idx);
self.inner.signal_work(&self.inner);
return true;
}
Empty => {}
Retry => found_work = true,
}
idx += 1;
} else {
idx = 0;
}
if idx == start {
break;
}
}
found_work
}
fn run_task(&self, task: Task, notify: &Arc<Notifier>, sender: &mut Sender) {
use task::Run::*;
match task.run(notify, sender) {
Idle => {}
Schedule => {
self.entry().push_internal(task);
}
Complete => {
let mut state: State = self.inner.state.load(Acquire).into();
loop {
let mut next = state;
next.dec_num_futures();
let actual = self.inner.state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
trace!("task complete; state={:?}", next);
if state.num_futures() == 1 {
if next.is_terminated() {
self.inner.terminate_sleeping_workers();
}
}
return;
}
state = actual;
}
}
}
}
#[inline]
fn drain_inbound(&self) -> bool {
use task::Poll::*;
let mut found_work = false;
loop {
let task = unsafe { self.entry().inbound.poll() };
match task {
Empty => {
if found_work {
trace!("found work while draining; signal_work");
self.inner.signal_work(&self.inner);
}
return true;
}
Inconsistent => {
if found_work {
trace!("found work while draining; signal_work");
self.inner.signal_work(&self.inner);
}
return false;
}
Data(task) => {
found_work = true;
self.entry().push_internal(task);
}
}
}
}
#[inline]
fn sleep(&self) -> bool {
trace!("Worker::sleep; idx={}", self.idx);
let mut state: WorkerState = self.entry().state.load(Acquire).into();
loop {
let mut next = state;
match state.lifecycle() {
WORKER_RUNNING => {
next.set_pushed();
}
WORKER_NOTIFIED | WORKER_SIGNALED => {
next.set_lifecycle(WORKER_RUNNING);
}
actual => panic!("unexpected worker state; {}", actual),
}
let actual = self.entry().state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
if state.is_notified() {
return true;
}
if !state.is_pushed() {
debug_assert!(next.is_pushed());
trace!(" sleeping -- push to stack; idx={}", self.idx);
if let Err(_) = self.inner.push_sleeper(self.idx) {
trace!(" sleeping -- push to stack failed; idx={}", self.idx);
return true;
}
}
break;
}
state = actual;
}
let mut lock = self.entry().park_mutex.lock().unwrap();
loop {
let mut next = state;
match state.lifecycle() {
WORKER_RUNNING => {}
WORKER_NOTIFIED | WORKER_SIGNALED => {
drop(lock);
loop {
let mut next = state;
next.set_lifecycle(WORKER_RUNNING);
let actual = self.entry().state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
return true;
}
state = actual;
}
}
_ => unreachable!(),
}
trace!(" sleeping -- set WORKER_SLEEPING; idx={}", self.idx);
next.set_lifecycle(WORKER_SLEEPING);
let actual = self.entry().state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
break;
}
state = actual;
}
trace!(" -> starting to sleep; idx={}", self.idx);
let sleep_until = self.inner.config.keep_alive
.map(|dur| Instant::now() + dur);
loop {
let mut drop_thread = false;
lock = match sleep_until {
Some(when) => {
let now = Instant::now();
if when >= now {
drop_thread = true;
}
let dur = when - now;
self.entry().park_condvar
.wait_timeout(lock, dur)
.unwrap().0
}
None => {
self.entry().park_condvar.wait(lock).unwrap()
}
};
trace!(" -> wakeup; idx={}", self.idx);
state = self.entry().state.load(Acquire).into();
loop {
match state.lifecycle() {
WORKER_SLEEPING => {}
WORKER_NOTIFIED | WORKER_SIGNALED => {
drop(lock);
loop {
let mut next = state;
next.set_lifecycle(WORKER_RUNNING);
let actual = self.entry().state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
return true;
}
state = actual;
}
}
_ => unreachable!(),
}
if !drop_thread {
break;
}
let mut next = state;
next.set_lifecycle(WORKER_SHUTDOWN);
let actual = self.entry().state.compare_and_swap(
state.into(), next.into(), AcqRel).into();
if actual == state {
return false;
}
state = actual;
}
}
}
fn entry(&self) -> &WorkerEntry {
&self.inner.workers[self.idx]
}
}
impl Drop for Worker {
fn drop(&mut self) {
trace!("shutting down thread; idx={}", self.idx);
if self.should_finalize.get() {
self.drain_inbound();
while let Some(_) = self.entry().deque.pop() {
}
self.inner.worker_terminated();
}
}
}
impl State {
#[inline]
fn new() -> State {
State(0)
}
fn num_futures(&self) -> usize {
self.0 >> NUM_FUTURES_OFFSET
}
fn inc_num_futures(&mut self) {
debug_assert!(self.num_futures() < MAX_FUTURES);
debug_assert!(self.lifecycle() < SHUTDOWN_NOW);
self.0 += 1 << NUM_FUTURES_OFFSET;
}
fn dec_num_futures(&mut self) {
let num_futures = self.num_futures();
if num_futures == 0 {
return;
}
self.0 -= 1 << NUM_FUTURES_OFFSET;
if self.lifecycle() == SHUTDOWN_ON_IDLE && num_futures == 1 {
self.0 = SHUTDOWN_NOW;
}
}
fn clear_num_futures(&mut self) {
self.0 = self.0 & LIFECYCLE_MASK;
}
fn lifecycle(&self) -> usize {
self.0 & LIFECYCLE_MASK
}
fn set_lifecycle(&mut self, val: usize) {
self.0 = (self.0 & NUM_FUTURES_MASK) | val;
}
fn is_terminated(&self) -> bool {
self.lifecycle() == SHUTDOWN_NOW && self.num_futures() == 0
}
}
impl From<usize> for State {
fn from(src: usize) -> Self {
State(src)
}
}
impl From<State> for usize {
fn from(src: State) -> Self {
src.0
}
}
impl fmt::Debug for State {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("State")
.field("lifecycle", &self.lifecycle())
.field("num_futures", &self.num_futures())
.finish()
}
}
impl SleepStack {
#[inline]
fn new() -> SleepStack {
SleepStack(EMPTY)
}
#[inline]
fn head(&self) -> usize {
self.0 & STACK_MASK
}
#[inline]
fn set_head(&mut self, val: usize) {
let aba_guard = ((self.0 >> ABA_GUARD_SHIFT) + 1) & ABA_GUARD_MASK;
self.0 = (aba_guard << ABA_GUARD_SHIFT) | val;
}
}
impl From<usize> for SleepStack {
fn from(src: usize) -> Self {
SleepStack(src)
}
}
impl From<SleepStack> for usize {
fn from(src: SleepStack) -> Self {
src.0
}
}
impl fmt::Debug for SleepStack {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
let head = self.head();
let mut fmt = fmt.debug_struct("SleepStack");
if head < MAX_WORKERS {
fmt.field("head", &head);
} else if head == EMPTY {
fmt.field("head", &"EMPTY");
} else if head == TERMINATED {
fmt.field("head", &"TERMINATED");
}
fmt.finish()
}
}
impl WorkerEntry {
fn new() -> Self {
let w = deque::Deque::new();
let s = w.stealer();
WorkerEntry {
state: AtomicUsize::new(WorkerState::default().into()),
next_sleeper: UnsafeCell::new(0),
deque: w,
steal: s,
inbound: task::Queue::new(),
park_mutex: Mutex::new(()),
park_condvar: Condvar::new(),
}
}
#[inline]
fn submit_internal(&self, task: Task) {
self.push_internal(task);
}
fn submit_external(&self, task: Task, mut state: WorkerState) -> bool {
self.push_external(task);
loop {
let mut next = state;
next.notify();
let actual = self.state.compare_and_swap(
state.into(), next.into(),
AcqRel).into();
if state == actual {
break;
}
state = actual;
}
match state.lifecycle() {
WORKER_SLEEPING => {
self.wakeup();
true
}
WORKER_SHUTDOWN => false,
_ => true,
}
}
#[inline]
fn push_external(&self, task: Task) {
self.inbound.push(task);
}
#[inline]
fn push_internal(&self, task: Task) {
self.deque.push(task);
}
#[inline]
fn wakeup(&self) {
let _lock = self.park_mutex.lock().unwrap();
self.park_condvar.notify_one();
}
#[inline]
fn next_sleeper(&self) -> usize {
unsafe { *self.next_sleeper.get() }
}
#[inline]
fn set_next_sleeper(&self, val: usize) {
unsafe { *self.next_sleeper.get() = val; }
}
}
impl WorkerState {
fn is_pushed(&self) -> bool {
self.0 & PUSHED_MASK == PUSHED_MASK
}
fn set_pushed(&mut self) {
self.0 |= PUSHED_MASK
}
fn is_notified(&self) -> bool {
match self.lifecycle() {
WORKER_NOTIFIED | WORKER_SIGNALED => true,
_ => false,
}
}
fn lifecycle(&self) -> usize {
(self.0 & WORKER_LIFECYCLE_MASK) >> WORKER_LIFECYCLE_SHIFT
}
fn set_lifecycle(&mut self, val: usize) {
self.0 = (self.0 & !WORKER_LIFECYCLE_MASK) |
(val << WORKER_LIFECYCLE_SHIFT)
}
fn is_signaled(&self) -> bool {
self.lifecycle() == WORKER_SIGNALED
}
fn notify(&mut self) {
if self.lifecycle() != WORKER_SIGNALED {
self.set_lifecycle(WORKER_NOTIFIED)
}
}
}
impl Default for WorkerState {
fn default() -> WorkerState {
WorkerState(PUSHED_MASK)
}
}
impl From<usize> for WorkerState {
fn from(src: usize) -> Self {
WorkerState(src)
}
}
impl From<WorkerState> for usize {
fn from(src: WorkerState) -> Self {
src.0
}
}
impl fmt::Debug for WorkerState {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
fmt.debug_struct("WorkerState")
.field("lifecycle", &match self.lifecycle() {
WORKER_SHUTDOWN => "WORKER_SHUTDOWN",
WORKER_RUNNING => "WORKER_RUNNING",
WORKER_SLEEPING => "WORKER_SLEEPING",
WORKER_NOTIFIED => "WORKER_NOTIFIED",
WORKER_SIGNALED => "WORKER_SIGNALED",
_ => unreachable!(),
})
.field("is_pushed", &self.is_pushed())
.finish()
}
}
impl Callback {
fn new<F>(f: F) -> Self
where F: Fn(&Worker, &mut Enter) + Send + Sync + 'static
{
Callback { f: Arc::new(f) }
}
pub fn call(&self, worker: &Worker, enter: &mut Enter) {
(self.f)(worker, enter)
}
}
impl fmt::Debug for Callback {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
write!(fmt, "Fn")
}
}
#[cfg(feature = "unstable-futures")]
impl Futures2Wake {
fn new(id: usize, inner: &Arc<Inner>) -> Futures2Wake {
let notifier = Arc::new(Notifier {
inner: Arc::downgrade(inner),
});
Futures2Wake { id, notifier }
}
}
#[cfg(feature = "unstable-futures")]
impl Drop for Futures2Wake {
fn drop(&mut self) {
self.notifier.drop_id(self.id)
}
}
#[cfg(feature = "unstable-futures")]
struct ArcWrapped(PhantomData<Futures2Wake>);
#[cfg(feature = "unstable-futures")]
unsafe impl futures2::task::UnsafeWake for ArcWrapped {
unsafe fn clone_raw(&self) -> futures2::task::Waker {
let me: *const ArcWrapped = self;
let arc = (*(&me as *const *const ArcWrapped as *const Arc<Futures2Wake>)).clone();
arc.notifier.clone_id(arc.id);
into_waker(arc)
}
unsafe fn drop_raw(&self) {
let mut me: *const ArcWrapped = self;
let me = &mut me as *mut *const ArcWrapped as *mut Arc<Futures2Wake>;
(*me).notifier.drop_id((*me).id);
::std::ptr::drop_in_place(me);
}
unsafe fn wake(&self) {
let me: *const ArcWrapped = self;
let me = &me as *const *const ArcWrapped as *const Arc<Futures2Wake>;
(*me).notifier.notify((*me).id)
}
}
#[cfg(feature = "unstable-futures")]
fn into_waker(rc: Arc<Futures2Wake>) -> futures2::task::Waker {
unsafe {
let ptr = mem::transmute::<Arc<Futures2Wake>, *mut ArcWrapped>(rc);
futures2::task::Waker::new(ptr)
}
}