Threadpool blocking (#317)

This patch adds a `blocking` to `tokio-threadpool`. This function serves
as a way to annotate sections of code that will perform blocking
operations. This informs the thread pool that an additional thread needs
to be spawned to replace the current thread, which will no longer be
able to process the work queue.

tokio-threadpool/src/pool/backup.rs

@@ -0,0 +1,304 @@
+use park::DefaultPark;
+use worker::{WorkerId};
+
+use std::cell::UnsafeCell;
+use std::fmt;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering::{self, Acquire, AcqRel, Relaxed};
+
+/// State associated with a thread in the thread pool.
+///
+/// The pool manages a number of threads. Some of those threads are considered
+/// "primary" threads and process the work queue. When a task being run on a
+/// primary thread enters a blocking context, the responsibility of processing
+/// the work queue must be handed off to another thread. This is done by first
+/// checking for idle threads on the backup stack. If one is found, the worker
+/// token (`WorkerId`) is handed off to that running thread. If none are found,
+/// a new thread is spawned.
+///
+/// This state manages the exchange. A thread that is idle, not assigned to a
+/// work queue, sits around for a specified amount of time. When the worker
+/// token is handed off, it is first stored in `handoff`. The backup thread is
+/// then signaled. At this point, the backup thread wakes up from sleep and
+/// reads `handoff`. At that point, it has been promoted to a primary thread and
+/// will begin processing inbound work on the work queue.
+///
+/// The name `Backup` isn't really great for what the type does, but I have not
+/// come up with a better name... Maybe it should just be named `Thread`.
+#[derive(Debug)]
+pub(crate) struct Backup {
+    /// Worker ID that is being handed to this thread.
+    handoff: UnsafeCell<Option<WorkerId>>,
+
+    /// Thread state.
+    ///
+    /// This tracks:
+    ///
+    /// * Is queued flag
+    /// * If the pool is shutting down.
+    /// * If the thread is running
+    state: AtomicUsize,
+
+    /// Next entry in the treiber stack.
+    next_sleeper: UnsafeCell<BackupId>,
+
+    /// Used to put the thread to sleep
+    park: DefaultPark,
+}
+
+#[derive(Debug, Eq, PartialEq, Copy, Clone)]
+pub(crate) struct BackupId(pub(crate) usize);
+
+#[derive(Debug)]
+pub(crate) enum Handoff {
+    Worker(WorkerId),
+    Idle,
+    Terminated,
+}
+
+/// Tracks thread state.
+#[derive(Clone, Copy, Eq, PartialEq)]
+struct State(usize);
+
+/// Set when the worker is pushed onto the scheduler's stack of sleeping
+/// threads.
+///
+/// This flag also serves as a "notification" bit. If another thread is
+/// attempting to hand off a worker to the backup thread, then the pushed bit
+/// will not be set when the thread tries to shutdown.
+pub const PUSHED: usize = 0b001;
+
+/// Set when the thread is running
+pub const RUNNING: usize = 0b010;
+
+/// Set when the thread pool has terminated
+pub const TERMINATED: usize = 0b100;
+
+// ===== impl Backup =====
+
+impl Backup {
+    pub fn new() -> Backup {
+        Backup {
+            handoff: UnsafeCell::new(None),
+            state: AtomicUsize::new(State::new().into()),
+            next_sleeper: UnsafeCell::new(BackupId(0)),
+            park: DefaultPark::new(),
+        }
+    }
+
+    /// Called when the thread is starting
+    pub fn start(&self, worker_id: &WorkerId) {
+        debug_assert!({
+            let state: State = self.state.load(Relaxed).into();
+
+            debug_assert!(!state.is_pushed());
+            debug_assert!(state.is_running());
+            debug_assert!(!state.is_terminated());
+
+            true
+        });
+
+        // The handoff value is equal to `worker_id`
+        debug_assert_eq!(unsafe { (*self.handoff.get()).as_ref()  }, Some(worker_id));
+
+        unsafe { *self.handoff.get() = None; }
+    }
+
+    pub fn is_running(&self) -> bool {
+        let state: State = self.state.load(Relaxed).into();
+        state.is_running()
+    }
+
+    /// Hands off the worker to a thread.
+    ///
+    /// Returns `true` if the thread needs to be spawned.
+    pub fn worker_handoff(&self, worker_id: WorkerId) -> bool {
+        unsafe {
+            // The backup worker should not already have been handoff a worker.
+            debug_assert!((*self.handoff.get()).is_none());
+
+            // Set the handoff
+            *self.handoff.get() = Some(worker_id);
+        }
+
+        // This *probably* can just be `Release`... memory orderings, how do
+        // they work?
+        let prev = State::worker_handoff(&self.state);
+        debug_assert!(prev.is_pushed());
+
+        if prev.is_running() {
+            // Wakeup the backup thread
+            self.park.notify();
+            false
+        } else {
+            true
+        }
+    }
+
+    /// Terminate the worker
+    pub fn signal_stop(&self) {
+        let prev: State = self.state.fetch_xor(TERMINATED | PUSHED, AcqRel).into();
+
+        debug_assert!(!prev.is_terminated());
+        debug_assert!(prev.is_pushed());
+
+        if prev.is_running() {
+            self.park.notify();
+        }
+    }
+
+    /// Release the worker
+    pub fn release(&self) {
+        let prev: State = self.state.fetch_xor(RUNNING, AcqRel).into();
+
+        debug_assert!(prev.is_running());
+    }
+
+    /// Wait for a worker handoff
+    pub fn wait_for_handoff(&self, sleep: bool) -> Handoff {
+        let mut state: State = self.state.load(Acquire).into();
+
+        // Run in a loop since there can be spurious wakeups
+        loop {
+            if !state.is_pushed() {
+                if state.is_terminated() {
+                    return Handoff::Terminated;
+                }
+
+                let worker_id = unsafe {
+                    (*self.handoff.get()).take()
+                        .expect("no worker handoff")
+                };
+
+                return Handoff::Worker(worker_id);
+            }
+
+            if sleep {
+                // TODO: Park with a timeout
+                self.park.park_sync(None);
+
+                // Reload the state
+                state = self.state.load(Acquire).into();
+                debug_assert!(state.is_running());
+            } else {
+                debug_assert!(state.is_running());
+
+                // Transition out of running
+                let mut next = state;
+                next.unset_running();
+
+                let actual = self.state.compare_and_swap(
+                    state.into(),
+                    next.into(),
+                    AcqRel).into();
+
+                if actual == state {
+                    debug_assert!(!next.is_running());
+
+                    return Handoff::Idle;
+                }
+
+                state = actual;
+            }
+        }
+    }
+
+    pub fn is_pushed(&self) -> bool {
+        let state: State = self.state.load(Relaxed).into();
+        state.is_pushed()
+    }
+
+    pub fn set_pushed(&self, ordering: Ordering) {
+        let prev: State = self.state.fetch_or(PUSHED, ordering).into();
+        debug_assert!(!prev.is_pushed());
+    }
+
+    #[inline]
+    pub fn next_sleeper(&self) -> BackupId {
+        unsafe { *self.next_sleeper.get() }
+    }
+
+    #[inline]
+    pub fn set_next_sleeper(&self, val: BackupId) {
+        unsafe { *self.next_sleeper.get() = val; }
+    }
+}
+
+// ===== impl State =====
+
+impl State {
+    /// Returns a new, default, thread `State`
+    pub fn new() -> State {
+        State(0)
+    }
+
+    /// Returns true if the thread entry is pushed in the sleeper stack
+    pub fn is_pushed(&self) -> bool {
+        self.0 & PUSHED == PUSHED
+    }
+
+    pub fn set_pushed(&mut self) {
+        self.0 |= PUSHED;
+    }
+
+    fn unset_pushed(&mut self) {
+        self.0 &= !PUSHED;
+    }
+
+    pub fn is_running(&self) -> bool {
+        self.0 & RUNNING == RUNNING
+    }
+
+    pub fn set_running(&mut self) {
+        self.0 |= RUNNING;
+    }
+
+    pub fn unset_running(&mut self) {
+        self.0 &= !RUNNING;
+    }
+
+    pub fn is_terminated(&self) -> bool {
+        self.0 & TERMINATED == TERMINATED
+    }
+
+    fn worker_handoff(state: &AtomicUsize) -> State {
+        let mut curr: State = state.load(Acquire).into();
+
+        loop {
+            let mut next = curr;
+            next.set_running();
+            next.unset_pushed();
+
+            let actual = state.compare_and_swap(
+                curr.into(), next.into(), AcqRel).into();
+
+            if actual == curr {
+                return curr;
+            }
+
+            curr = actual;
+        }
+    }
+}
+
+impl From<usize> for State {
+    fn from(src: usize) -> State {
+        State(src)
+    }
+}
+
+impl From<State> for usize {
+    fn from(src: State) -> usize {
+        src.0
+    }
+}
+
+impl fmt::Debug for State {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_struct("backup::State")
+            .field("is_pushed", &self.is_pushed())
+            .field("is_running", &self.is_running())
+            .field("is_terminated", &self.is_terminated())
+            .finish()
+    }
+}

tokio-threadpool/src/pool/backup_stack.rs

@@ -0,0 +1,185 @@
+use pool::{Backup, BackupId};
+
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering::{Acquire, AcqRel};
+
+#[derive(Debug)]
+pub(crate) struct BackupStack {
+    state: AtomicUsize,
+}
+
+#[derive(Debug, Eq, PartialEq, Clone, Copy)]
+struct State(usize);
+
+pub(crate) const MAX_BACKUP: usize = 1 << 15;
+
+/// Extracts the head of the backup stack from the state
+const STACK_MASK: usize = ((1 << 16) - 1);
+
+/// Used to mark the stack as empty
+pub(crate) const EMPTY: BackupId = BackupId(MAX_BACKUP);
+
+/// Used to mark the stack as terminated
+pub(crate) const TERMINATED: BackupId = BackupId(EMPTY.0 + 1);
+
+/// How many bits the treiber ABA guard is offset by
+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;
+
+// ===== impl BackupStack =====
+
+impl BackupStack {
+    pub fn new() -> BackupStack {
+        let state = AtomicUsize::new(State::new().into());
+        BackupStack { state }
+    }
+
+    /// Push a backup thread onto the stack
+    ///
+    /// # Return
+    ///
+    /// Returns `Ok` on success.
+    ///
+    /// Returns `Err` if the pool has transitioned to the `TERMINATED` state.
+    /// Whene terminated, pushing new entries is no longer permitted.
+    pub fn push(&self, entries: &[Backup], id: BackupId) -> Result<(), ()> {
+        let mut state: State = self.state.load(Acquire).into();
+
+        entries[id.0].set_pushed(AcqRel);
+
+        loop {
+            let mut next = state;
+
+            let head = state.head();
+
+            if head == TERMINATED {
+                // The pool is terminated, cannot push the sleeper.
+                return Err(());
+            }
+
+            entries[id.0].set_next_sleeper(head);
+            next.set_head(id);
+
+            let actual = self.state.compare_and_swap(
+                state.into(), next.into(), AcqRel).into();
+
+            if state == actual {
+                return Ok(());
+            }
+
+            state = actual;
+        }
+    }
+
+    /// Pop a backup thread off the stack.
+    ///
+    /// If `terminate` is set and the stack is empty when this function is
+    /// called, the state of the stack is transitioned to "terminated". At this
+    /// point, no further entries can be pushed onto the stack.
+    ///
+    /// # Return
+    ///
+    /// * Returns the index of the popped worker and the worker's observed
+    ///   state.
+    ///
+    /// * `Ok(None)` if the stack is empty.
+    /// * `Err(_)` is returned if the pool has been shutdown.
+    pub fn pop(&self, entries: &[Backup], terminate: bool) -> Result<Option<BackupId>, ()> {
+        // Figure out the empty value
+        let terminal = match terminate {
+            true => TERMINATED,
+            false => EMPTY,
+        };
+
+        let mut state: State = self.state.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 Ok(None);
+                }
+
+                let actual = self.state.compare_and_swap(
+                    state.into(), next.into(), AcqRel).into();
+
+                if actual != state {
+                    state = actual;
+                    continue;
+                }
+
+                return Ok(None);
+            } else if head == TERMINATED {
+                return Err(());
+            }
+
+            debug_assert!(head.0 < MAX_BACKUP);
+
+            let mut next = state;
+
+            let next_head = entries[head.0].next_sleeper();
+
+            // TERMINATED can never be set as the "next pointer" on a worker.
+            debug_assert!(next_head != TERMINATED);
+
+            if next_head == EMPTY {
+                next.set_head(terminal);
+            } else {
+                next.set_head(next_head);
+            }
+
+            let actual = self.state.compare_and_swap(
+                state.into(), next.into(), AcqRel).into();
+
+            if actual == state {
+                debug_assert!(entries[head.0].is_pushed());
+                return Ok(Some(head));
+            }
+
+            state = actual;
+        }
+    }
+}
+
+// ===== impl State =====
+
+impl State {
+    fn new() -> State {
+        State(EMPTY.0)
+    }
+
+    fn head(&self) -> BackupId {
+        BackupId(self.0 & STACK_MASK)
+    }
+
+    fn set_head(&mut self, val: BackupId) {
+        let val = val.0;
+
+        // The ABA guard protects against the ABA problem w/ treiber stacks
+        let aba_guard = ((self.0 >> ABA_GUARD_SHIFT) + 1) & ABA_GUARD_MASK;
+
+        self.0 = (aba_guard << ABA_GUARD_SHIFT) | val;
+    }
+}
+
+impl From<usize> for State {
+    fn from(src: usize) -> Self {
+        State(src)
+    }
+}
+
+impl From<State> for usize {
+    fn from(src: State) -> Self {
+        src.0
+    }
+}

tokio-threadpool/src/pool/mod.rs

@@ -1,24 +1,31 @@
+mod backup;
+mod backup_stack;
 mod state;
-mod stack;
 
+pub(crate) use self::backup::{Backup, BackupId};
+pub(crate) use self::backup_stack::MAX_BACKUP;
 pub(crate) use self::state::{
     State,
     Lifecycle,
     MAX_FUTURES,
 };
-use self::stack::SleepStack;
+
+use self::backup::Handoff;
+use self::backup_stack::BackupStack;
 
 use config::Config;
 use shutdown_task::ShutdownTask;
-use task::Task;
+use task::{Task, Blocking};
 use worker::{self, Worker, WorkerId};
 
+use futures::Poll;
 use futures::task::AtomicTask;
 
 use std::cell::UnsafeCell;
 use std::sync::atomic::Ordering::{Acquire, AcqRel, Relaxed};
 use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
+use std::thread;
 
 use rand::{Rng, SeedableRng, XorShiftRng};
 
@@ -29,9 +36,9 @@ pub(crate) struct Pool {
     pub state: AtomicUsize,
 
     // Stack tracking sleeping workers.
-    sleep_stack: SleepStack,
+    sleep_stack: worker::Stack,
 
-    // Number of workers who haven't reached the final state of shutdown
+    // Number of workers that haven't reached the final state of shutdown
     //
     // This is only used to know when to single `shutdown_task` once the
     // shutdown process has completed.
@@ -40,11 +47,28 @@ pub(crate) struct Pool {
     // Used to generate a thread local RNG seed
     pub next_thread_id: AtomicUsize,
 
-    // Storage for workers
+    // Worker state
     //
-    // This will *usually* be a small number
+    // A worker is a thread that is processing the work queue and polling
+    // futures.
+    //
+    // This will *usually* be a small number.
     pub workers: Box<[worker::Entry]>,
 
+    // Backup thread state
+    //
+    // In order to efficiently support `blocking`, a pool of backup threads is
+    // needed. These backup threads are ready to take over a worker if the
+    // future being processed requires blocking.
+    backup: Box<[Backup]>,
+
+    // Stack of sleeping backup threads
+    pub backup_stack: BackupStack,
+
+    // State regarding coordinating blocking sections and tracking tasks that
+    // are pending blocking capacity.
+    blocking: Blocking,
+
     // Task notified when the worker shuts down
     pub shutdown_task: ShutdownTask,
 
@@ -52,17 +76,41 @@ pub(crate) struct Pool {
     pub config: Config,
 }
 
+const TERMINATED: usize = 1;
+
 impl Pool {
     /// Create a new `Pool`
-    pub fn new(workers: Box<[worker::Entry]>, config: Config) -> Pool {
+    pub fn new(workers: Box<[worker::Entry]>, max_blocking: usize, config: Config) -> Pool {
         let pool_size = workers.len();
+        let total_size = max_blocking + pool_size;
+
+        // Create the set of backup entries
+        //
+        // This is `backup + pool_size` because the core thread pool running the
+        // workers is spawned from backup as well.
+        let backup = (0..total_size).map(|_| {
+            Backup::new()
+        }).collect::<Vec<_>>().into_boxed_slice();
+
+        let backup_stack = BackupStack::new();
+
+        for i in (0..backup.len()).rev() {
+            backup_stack.push(&backup, BackupId(i))
+                .unwrap();
+        }
+
+        // Initialize the blocking state
+        let blocking = Blocking::new(max_blocking);
 
         let ret = Pool {
             state: AtomicUsize::new(State::new().into()),
-            sleep_stack: SleepStack::new(),
-            num_workers: AtomicUsize::new(pool_size),
+            sleep_stack: worker::Stack::new(),
+            num_workers: AtomicUsize::new(0),
             next_thread_id: AtomicUsize::new(0),
             workers,
+            backup,
+            backup_stack,
+            blocking,
             shutdown_task: ShutdownTask {
                 task1: AtomicTask::new(),
                 #[cfg(feature = "unstable-futures")]
@@ -141,6 +189,10 @@ impl Pool {
         self.terminate_sleeping_workers();
     }
 
+    pub fn is_shutdown(&self) -> bool {
+        self.num_workers.load(Acquire) == TERMINATED
+    }
+
     /// Called by `Worker` as it tries to enter a sleeping state. Before it
     /// sleeps, it must push itself onto the sleep stack. This enables other
     /// threads to see it when signaling work.
@@ -151,32 +203,67 @@ impl Pool {
     pub fn terminate_sleeping_workers(&self) {
         use worker::Lifecycle::Signaled;
 
+        // First, set the TERMINATED flag on `num_workers`. This signals that
+        // whichever thread transitions the count to zero must notify the
+        // shutdown task.
+        let prev = self.num_workers.fetch_or(TERMINATED, AcqRel);
+        let notify = prev == 0;
+
         trace!("  -> shutting down workers");
         // Wakeup all sleeping workers. They will wake up, see the state
         // transition, and terminate.
         while let Some((idx, worker_state)) = self.sleep_stack.pop(&self.workers, Signaled, true) {
-            trace!("  -> shutdown worker; idx={:?}; state={:?}", idx, worker_state);
+            self.workers[idx].signal_stop(worker_state);
+        }
 
-            if self.workers[idx].signal_stop(worker_state).is_err() {
-                // The worker is already in the shutdown state, immediately
-                // track that it has terminated as the worker will never work
-                // again.
-                self.worker_terminated();
-            }
+        // Now terminate any backup threads
+        //
+        // The call to `pop` must be successful because shutting down the pool
+        // is coordinated and at this point, this is the only thread that will
+        // attempt to transition the backup stack to "terminated".
+        while let Ok(Some(backup_id)) = self.backup_stack.pop(&self.backup, true) {
+            self.backup[backup_id.0].signal_stop();
+        }
+
+        if notify {
+            self.shutdown_task.notify();
         }
     }
 
-    pub fn worker_terminated(&self) {
-        let prev = self.num_workers.fetch_sub(1, AcqRel);
+    /// Track that a worker thread has started
+    ///
+    /// If `Err` is returned, then the thread is not permitted to started.
+    fn thread_started(&self) -> Result<(), ()> {
+        let mut curr = self.num_workers.load(Acquire);
 
-        trace!("worker_terminated; num_workers={}", prev - 1);
+        loop {
+            if curr & TERMINATED == TERMINATED {
+                return Err(());
+            }
 
-        if 1 == prev {
-            trace!("notifying shutdown task");
+            let actual = self.num_workers.compare_and_swap(
+                curr, curr + 2, AcqRel);
+
+            if curr == actual {
+                return Ok(());
+            }
+
+            curr = actual;
+        }
+    }
+
+    fn thread_stopped(&self) {
+        let prev = self.num_workers.fetch_sub(2, AcqRel);
+
+        if prev == TERMINATED | 2 {
             self.shutdown_task.notify();
         }
     }
 
+    pub fn poll_blocking_capacity(&self, task: &Arc<Task>) -> Poll<(), ::BlockingError> {
+        self.blocking.poll_blocking_capacity(task)
+    }
+
     /// Submit a task to the scheduler.
     ///
     /// Called from either inside or outside of the scheduler. If currently on
@@ -184,15 +271,19 @@ impl Pool {
     pub fn submit(&self, task: Arc<Task>, inner: &Arc<Pool>) {
         Worker::with_current(|worker| {
             match worker {
-                Some(worker) => {
-                    let idx = worker.id.idx;
+                // If the worker is in blocking mode, then even though the
+                // thread-local variable is set, the current thread does not
+                // have ownership of that worker entry. This is because the
+                // worker entry has already been handed off to another thread.
+                Some(worker) if !worker.is_blocking() => {
+                    let idx = worker.id.0;
 
                     trace!("    -> submit internal; idx={}", idx);
 
                     worker.inner.workers[idx].submit_internal(task);
                     worker.inner.signal_work(inner);
                 }
-                None => {
+                _ => {
                     self.submit_external(task, inner);
                 }
             }
@@ -203,7 +294,7 @@ impl Pool {
     ///
     /// Called from outside of the scheduler, this function is how new tasks
     /// enter the system.
-    fn submit_external(&self, task: Arc<Task>, inner: &Arc<Pool>) {
+    pub fn submit_external(&self, task: Arc<Task>, inner: &Arc<Pool>) {
         use worker::Lifecycle::Notified;
 
         // First try to get a handle to a sleeping worker. This ensures that
@@ -234,12 +325,130 @@ impl Pool {
         let entry = &self.workers[idx];
 
         if !entry.submit_external(task, state) {
-            self.spawn_worker(idx, inner);
+            self.spawn_thread(WorkerId::new(idx), inner);
         }
     }
 
-    fn spawn_worker(&self, idx: usize, inner: &Arc<Pool>) {
-        Worker::spawn(WorkerId::new(idx), inner);
+    pub fn release_backup(&self, backup_id: BackupId) -> Result<(), ()> {
+        // First update the state, this cannot fail because the caller must have
+        // exclusive access to the backup token.
+        self.backup[backup_id.0].release();
+
+        // Push the backup entry back on the stack
+        self.backup_stack.push(&self.backup, backup_id)
+    }
+
+    pub fn notify_blocking_task(&self, pool: &Arc<Pool>) {
+        self.blocking.notify_task(&pool);
+    }
+
+    /// Provision a thread to run a worker
+    pub fn spawn_thread(&self, id: WorkerId, inner: &Arc<Pool>) {
+        let backup_id = match self.backup_stack.pop(&self.backup, false) {
+            Ok(Some(backup_id)) => backup_id,
+            Ok(None) => panic!("no thread available"),
+            Err(_) => {
+                debug!("failed to spawn worker thread due to the thread pool shutting down");
+                return;
+            }
+        };
+
+        let need_spawn = self.backup[backup_id.0]
+            .worker_handoff(id.clone());
+
+        if !need_spawn {
+            return;
+        }
+
+        if self.thread_started().is_err() {
+            // The pool is shutting down.
+            return;
+        }
+
+        let mut th = thread::Builder::new();
+
+        if let Some(ref prefix) = inner.config.name_prefix {
+            th = th.name(format!("{}{}", prefix, backup_id.0));
+        }
+
+        if let Some(stack) = inner.config.stack_size {
+            th = th.stack_size(stack);
+        }
+
+        let inner = inner.clone();
+
+        let res = th.spawn(move || {
+            if let Some(ref f) = inner.config.after_start {
+                f();
+            }
+
+            let mut worker_id = id;
+
+            inner.backup[backup_id.0].start(&worker_id);
+
+            loop {
+                // The backup token should be in the running state.
+                debug_assert!(inner.backup[backup_id.0].is_running());
+
+                // TODO: Avoid always cloning
+                let worker = Worker::new(worker_id, backup_id, inner.clone());
+
+                // Run the worker. If the worker transitioned to a "blocking"
+                // state, then `is_blocking` will be true.
+                if !worker.do_run() {
+                    // The worker shutdown, so exit the thread.
+                    break;
+                }
+
+                // Push the thread back onto the backup stack. This makes it
+                // available for future handoffs.
+                //
+                // This **must** happen before notifying the task.
+                let res = inner.backup_stack
+                    .push(&inner.backup, backup_id);
+
+                if res.is_err() {
+                    // The pool is being shutdown.
+                    break;
+                }
+
+                // The task switched the current thread to blocking mode.
+                // Now that the blocking task completed, any tasks
+                inner.notify_blocking_task(&inner);
+
+                debug_assert!(inner.backup[backup_id.0].is_running());
+
+                // Wait for a handoff
+                let handoff = inner.backup[backup_id.0]
+                    .wait_for_handoff(true);
+
+                match handoff {
+                    Handoff::Worker(id) => {
+                        debug_assert!(inner.backup[backup_id.0].is_running());
+                        worker_id = id;
+                    }
+                    Handoff::Idle => {
+                        // Worker is idle
+                        break;
+                    }
+                    Handoff::Terminated => {
+                        // TODO: When wait_for_handoff supports blocking with a
+                        // timeout, this will have to be smarter
+                        break;
+                    }
+                }
+            }
+
+            if let Some(ref f) = inner.config.before_stop {
+                f();
+            }
+
+            inner.thread_stopped();
+        });
+
+        if let Err(e) = res {
+            warn!("failed to spawn worker thread; err={:?}", e);
+        }
     }
 
     /// If there are any other workers currently relaxing, signal them that work
@@ -277,7 +486,7 @@ impl Pool {
                 }
                 Shutdown => {
                     trace!("signal_work -- spawn; idx={}", idx);
-                    Worker::spawn(WorkerId::new(idx), inner);
+                    self.spawn_thread(WorkerId(idx), inner);
                 }
                 Running | Notified | Signaled => {
                     // The workers are already active. No need to wake them up.
@@ -286,7 +495,6 @@ impl Pool {
         }
     }
 
-
     /// Generates a random number
     ///
     /// Uses a thread-local seeded XorShift.

tokio-threadpool/src/pool/stack.rs

@@ -1,252 +0,0 @@
-use config::MAX_WORKERS;
-use worker;
-
-use std::{fmt, usize};
-use std::sync::atomic::AtomicUsize;
-use std::sync::atomic::Ordering::{Acquire, AcqRel, Relaxed};
-
-/// Lock-free stack of sleeping workers.
-///
-/// This is implemented as a Treiber stack and references to nodes are
-/// `usize` values, indexing the entry in the `[worker::Entry]` array stored by
-/// `Pool`. Each `Entry` instance maintains a `pushed` bit in its state. This
-/// bit tracks if the entry is already pushed onto the stack or not. A single
-/// entry can only be stored on the stack a single time.
-///
-/// By using indexes instead of pointers, that allows a much greater amount of
-/// data to be used for the ABA guard (see correctness section of wikipedia
-/// page).
-///
-/// Treiber stack: https://en.wikipedia.org/wiki/Treiber_Stack
-#[derive(Debug)]
-pub(crate) struct SleepStack {
-    state: AtomicUsize,
-}
-
-/// State related to the stack of sleeping workers.
-///
-/// - Parked head     16 bits
-/// - Sequence        remaining
-///
-/// The parked head value has a couple of special values:
-///
-/// - EMPTY: No sleepers
-/// - TERMINATED: Don't spawn more threads
-#[derive(Eq, PartialEq, Clone, Copy)]
-pub struct State(usize);
-
-/// Extracts the head of the worker stack from the scheduler state
-const STACK_MASK: usize = ((1 << 16) - 1);
-
-/// Used to mark the stack as empty
-pub(crate) const EMPTY: usize = MAX_WORKERS;
-
-/// Used to mark the stack as terminated
-pub(crate) const TERMINATED: usize = EMPTY + 1;
-
-/// How many bits the treiber ABA guard is offset by
-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;
-
-// ===== impl SleepStack =====
-
-impl SleepStack {
-    /// Create a new `SleepStack` representing the empty state.
-    pub fn new() -> SleepStack {
-        let state = AtomicUsize::new(State::new().into());
-        SleepStack { state }
-    }
-
-    /// Push a worker onto the stack
-    ///
-    /// # Return
-    ///
-    /// Returns `Ok` on success.
-    ///
-    /// Returns `Err` if the pool has transitioned to the `TERMINATED` state.
-    /// Whene terminated, pushing new entries is no longer permitted.
-    pub fn push(&self, entries: &[worker::Entry], idx: usize) -> Result<(), ()> {
-        let mut state: State = self.state.load(Acquire).into();
-
-        debug_assert!(worker::State::from(entries[idx].state.load(Relaxed)).is_pushed());
-
-        loop {
-            let mut next = state;
-
-            let head = state.head();
-
-            if head == TERMINATED {
-                // The pool is terminated, cannot push the sleeper.
-                return Err(());
-            }
-
-            entries[idx].set_next_sleeper(head);
-            next.set_head(idx);
-
-            let actual = self.state.compare_and_swap(
-                state.into(), next.into(), AcqRel).into();
-
-            if state == actual {
-                return Ok(());
-            }
-
-            state = actual;
-        }
-    }
-
-
-    /// Pop a worker off the stack.
-    ///
-    /// If `terminate` is set and the stack is empty when this function is
-    /// called, the state of the stack is transitioned to "terminated". At this
-    /// point, no further workers can be pusheed onto the stack.
-    ///
-    /// # Return
-    ///
-    /// Returns the index of the popped worker and the worker's observed state.
-    ///
-    /// `None` if the stack is empty.
-    pub fn pop(&self, entries: &[worker::Entry],
-                  max_lifecycle: worker::Lifecycle,
-                  terminate: bool)
-        -> Option<(usize, worker::State)>
-    {
-        // Figure out the empty value
-        let terminal = match terminate {
-            true => TERMINATED,
-            false => EMPTY,
-        };
-
-        // If terminating, the max lifecycle *must* be `Signaled`, which is the
-        // highest lifecycle. By passing the greatest possible lifecycle value,
-        // no entries are skipped by this function.
-        //
-        // TODO: It would be better to terminate in a separate function that
-        // atomically takes all values and transitions to a terminated state.
-        debug_assert!(!terminate || max_lifecycle == worker::Lifecycle::Signaled);
-
-        let mut state: State = self.state.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.state.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 = entries[head].next_sleeper();
-
-            // TERMINATED can never be set as the "next pointer" on a worker.
-            debug_assert!(next_head != TERMINATED);
-
-            if next_head == EMPTY {
-                next.set_head(terminal);
-            } else {
-                next.set_head(next_head);
-            }
-
-            let actual = self.state.compare_and_swap(
-                state.into(), next.into(), AcqRel).into();
-
-            if actual == state {
-                // Release ordering is needed to ensure that unsetting the
-                // `pushed` flag happens after popping the sleeper from the
-                // stack.
-                //
-                // Acquire ordering is required to acquire any memory associated
-                // with transitioning the worker's lifecycle.
-                let state = entries[head].fetch_unset_pushed(AcqRel);
-
-                if state.lifecycle() >= max_lifecycle {
-                    // If the worker has already been notified, then it is
-                    // warming up to do more work. In this case, try to pop
-                    // another thread that might be in a relaxed state.
-                    continue;
-                }
-
-                return Some((head, state));
-            }
-
-            state = actual;
-        }
-    }
-}
-
-// ===== impl State =====
-
-impl State {
-    #[inline]
-    fn new() -> State {
-        State(EMPTY)
-    }
-
-    #[inline]
-    fn head(&self) -> usize {
-        self.0 & STACK_MASK
-    }
-
-    #[inline]
-    fn set_head(&mut self, val: usize) {
-        // The ABA guard protects against the ABA problem w/ treiber stacks
-        let aba_guard = ((self.0 >> ABA_GUARD_SHIFT) + 1) & ABA_GUARD_MASK;
-
-        self.0 = (aba_guard << ABA_GUARD_SHIFT) | val;
-    }
-}
-
-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 {
-        let head = self.head();
-
-        let mut fmt = fmt.debug_struct("stack::State");
-
-        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()
-    }
-}