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/worker/mod.rs

@@ -1,27 +1,30 @@
 mod entry;
+mod stack;
 mod state;
 
 pub(crate) use self::entry::{
     WorkerEntry as Entry,
 };
+pub(crate) use self::stack::Stack;
 pub(crate) use self::state::{
     State,
     Lifecycle,
 };
 
-use pool::{self, Pool};
+use pool::{self, Pool, BackupId};
 use notifier::Notifier;
 use sender::Sender;
-use task::Task;
+use task::{self, Task, CanBlock};
 
 use tokio_executor;
 
+use futures::{Poll, Async};
+
 use std::cell::Cell;
 use std::marker::PhantomData;
 use std::rc::Rc;
 use std::sync::atomic::Ordering::{AcqRel, Acquire};
 use std::sync::Arc;
-use std::thread;
 use std::time::{Duration, Instant};
 
 /// Thread worker
@@ -36,6 +39,19 @@ pub struct Worker {
     // WorkerEntry index
     pub(crate) id: WorkerId,
 
+    // Backup thread ID assigned to processing this worker.
+    backup_id: BackupId,
+
+    // Set to the task that is currently being polled by the worker. This is
+    // needed so that `blocking` blocks are able to interact with this task.
+    //
+    // This has to be a raw pointer to make it compile, but great care is taken
+    // when this is set.
+    current_task: CurrentTask,
+
+    // Set when the thread is in blocking mode.
+    is_blocking: Cell<bool>,
+
     // Set when the worker should finalize on drop
     should_finalize: Cell<bool>,
 
@@ -43,65 +59,72 @@ pub struct Worker {
     _p: PhantomData<Rc<()>>,
 }
 
+/// Tracks the state related to the currently running task.
+#[derive(Debug)]
+struct CurrentTask {
+    /// This has to be a raw pointer to make it compile, but great care is taken
+    /// when this is set.
+    task: Cell<Option<*const Arc<Task>>>,
+
+    /// Tracks the blocking capacity allocation state.
+    can_block: Cell<CanBlock>,
+}
+
 /// Identifiers a thread pool worker.
 ///
 /// This identifier is unique scoped by the thread pool. It is possible that
 /// different thread pool instances share worker identifier values.
 #[derive(Debug, Clone, Hash, Eq, PartialEq)]
-pub struct WorkerId {
-    pub(crate) idx: usize,
-}
+pub struct WorkerId(pub(crate) usize);
 
 // Pointer to the current worker info
 thread_local!(static CURRENT_WORKER: Cell<*const Worker> = Cell::new(0 as *const _));
 
 impl Worker {
-    pub(crate) fn spawn(id: WorkerId, inner: &Arc<Pool>) {
-        trace!("spawning new worker thread; id={}", id.idx);
-
-        let mut th = thread::Builder::new();
-
-        if let Some(ref prefix) = inner.config.name_prefix {
-            th = th.name(format!("{}{}", prefix, id.idx));
+    pub(crate) fn new(id: WorkerId, backup_id: BackupId, inner: Arc<Pool>) -> Worker {
+        Worker {
+            inner,
+            id,
+            backup_id,
+            current_task: CurrentTask::new(),
+            is_blocking: Cell::new(false),
+            should_finalize: Cell::new(false),
+            _p: PhantomData,
         }
+    }
 
-        if let Some(stack) = inner.config.stack_size {
-            th = th.stack_size(stack);
-        }
+    pub(crate) fn is_blocking(&self) -> bool {
+        self.is_blocking.get()
+    }
 
-        let inner = inner.clone();
+    /// Run the worker
+    ///
+    /// Returns `true` if the thread should keep running as a `backup` thread.
+    pub(crate) fn do_run(&self) -> bool {
+        // Create another worker... It's ok, this is just a new type around
+        // `Pool` that is expected to stay on the current thread.
+        CURRENT_WORKER.with(|c| {
+            c.set(self as *const _);
 
-        th.spawn(move || {
-            let worker = Worker {
-                inner,
-                id,
-                should_finalize: Cell::new(false),
-                _p: PhantomData,
-            };
+            let inner = self.inner.clone();
+            let mut sender = Sender { inner };
 
-            // Make sure the ref to the worker does not move
-            let wref = &worker;
+            // Enter an execution context
+            let mut enter = tokio_executor::enter().unwrap();
 
-            // Create another worker... It's ok, this is just a new type around
-            // `Pool` that is expected to stay on the current thread.
-            CURRENT_WORKER.with(|c| {
-                c.set(wref as *const _);
-
-                let inner = wref.inner.clone();
-                let mut sender = Sender { inner };
-
-                // Enter an execution context
-                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();
-                    }
-                });
+            tokio_executor::with_default(&mut sender, &mut enter, |enter| {
+                if let Some(ref callback) = self.inner.config.around_worker {
+                    callback.call(self, enter);
+                } else {
+                    self.run();
+                }
             });
-        }).unwrap();
+        });
+
+        // Can't be in blocking mode and finalization mode
+        debug_assert!(!self.is_blocking.get() || !self.should_finalize.get());
+
+        self.is_blocking.get()
     }
 
     pub(crate) fn with_current<F: FnOnce(Option<&Worker>) -> R, R>(f: F) -> R {
@@ -116,6 +139,67 @@ impl Worker {
         })
     }
 
+    /// Transition the current worker to a blocking worker
+    pub(crate) fn transition_to_blocking(&self) -> Poll<(), ::BlockingError> {
+        use self::CanBlock::*;
+
+        // If we get this far, then `current_task` has been set.
+        let task_ref = self.current_task.get_ref();
+
+        // First step is to acquire blocking capacity for the task.
+        match self.current_task.can_block() {
+            // Capacity to block has already been allocated to this task.
+            Allocated => {}
+
+            // The task has already requested capacity to block, but there is
+            // none yet available.
+            NoCapacity => return Ok(Async::NotReady),
+
+            // The task has yet to ask for capacity
+            CanRequest => {
+                // Atomically attempt to acquire blocking capacity, and if none
+                // is available, register the task to be notified once capacity
+                // becomes available.
+                match self.inner.poll_blocking_capacity(task_ref)? {
+                    Async::Ready(()) => {
+                        self.current_task.set_can_block(Allocated);
+                    }
+                    Async::NotReady => {
+                        self.current_task.set_can_block(NoCapacity);
+                        return Ok(Async::NotReady);
+                    }
+                }
+            }
+        }
+
+        // The task has been allocated blocking capacity. At this point, this is
+        // when the current thread transitions from a worker to a backup thread.
+        // To do so requires handing over the worker to another backup thread.
+
+        if self.is_blocking.get() {
+            // The thread is already in blocking mode, so there is nothing else
+            // to do. Return `Ready` and allow the caller to block the thread.
+            return Ok(().into());
+        }
+
+        trace!("transition to blocking state");
+
+        // Transitioning to blocking requires handing over the worker state to
+        // another thread so that the work queue can continue to be processed.
+
+        self.inner.spawn_thread(self.id.clone(), &self.inner);
+
+        // Track that the thread has now fully entered the blocking state.
+        self.is_blocking.set(true);
+
+        Ok(().into())
+    }
+
+    /// Transition from blocking
+    pub(crate) fn transition_from_blocking(&self) {
+        // TODO: Attempt to take ownership of the worker again.
+    }
+
     /// Returns a reference to the worker's identifier.
     ///
     /// This identifier is unique scoped by the thread pool. It is possible that
@@ -149,6 +233,11 @@ impl Worker {
 
             // Run the next available task
             if self.try_run_task(&notify, &mut sender) {
+                if self.is_blocking.get() {
+                    // Exit out of the run state
+                    return;
+                }
+
                 if tick % LIGHT_SLEEP_INTERVAL == 0 {
                     self.sleep_light();
                 }
@@ -160,17 +249,6 @@ impl Worker {
                 continue;
             }
 
-            // No work in this worker's queue, it is time to try stealing.
-            if self.try_steal_task(&notify, &mut sender) {
-                if tick % LIGHT_SLEEP_INTERVAL == 0 {
-                    self.sleep_light();
-                }
-
-                tick = tick.wrapping_add(1);
-                spin_cnt = 0;
-                continue;
-            }
-
             if !consistent {
                 spin_cnt = 0;
                 continue;
@@ -190,9 +268,31 @@ impl Worker {
             // If there still isn't any work to do, shutdown the worker?
         }
 
+        // The pool is terminating. However, transitioning the pool state to
+        // terminated is the very first step of the finalization process. Other
+        // threads may not see this state and try to spawn a new thread. To
+        // ensure consistency, before the current thread shuts down, it must
+        // return the backup token to the stack.
+        //
+        // The returned result is ignored because `Err` represents the pool
+        // shutting down. We are currently aware of this fact.
+        let _ = self.inner.release_backup(self.backup_id);
+
         self.should_finalize.set(true);
     }
 
+    /// Try to run a task
+    ///
+    /// Returns `true` if work was found.
+    #[inline]
+    fn try_run_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
+        if self.try_run_owned_task(notify, sender) {
+            return true;
+        }
+
+        self.try_steal_task(notify, sender)
+    }
+
     /// Checks the worker's current state, updating it as needed.
     ///
     /// Returns `true` if the worker should run.
@@ -200,6 +300,8 @@ impl Worker {
     fn check_run_state(&self, first: bool) -> bool {
         use self::Lifecycle::*;
 
+        debug_assert!(!self.is_blocking.get());
+
         let mut state: State = self.entry().state.load(Acquire).into();
 
         loop {
@@ -249,8 +351,7 @@ impl Worker {
     /// Runs the next task on this worker's queue.
     ///
     /// Returns `true` if work was found.
-    #[inline]
-    fn try_run_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
+    fn try_run_owned_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
         use deque::Steal::*;
 
         // Poll the internal queue for a task to run
@@ -267,10 +368,11 @@ impl Worker {
     /// Tries to steal a task from another worker.
     ///
     /// Returns `true` if work was found
-    #[inline]
     fn try_steal_task(&self, notify: &Arc<Notifier>, sender: &mut Sender) -> bool {
         use deque::Steal::*;
 
+        debug_assert!(!self.is_blocking.get());
+
         let len = self.inner.workers.len();
         let mut idx = self.inner.rand_usize() % len;
         let mut found_work = false;
@@ -285,9 +387,12 @@ impl Worker {
                         self.run_task(task, notify, sender);
 
                         trace!("try_steal_task -- signal_work; self={}; from={}",
-                               self.id.idx, idx);
+                               self.id.0, idx);
 
                         // Signal other workers that work is available
+                        //
+                        // TODO: Should this be called here or before
+                        // `run_task`?
                         self.inner.signal_work(&self.inner);
 
                         return true;
@@ -312,10 +417,28 @@ impl Worker {
     fn run_task(&self, task: Arc<Task>, notify: &Arc<Notifier>, sender: &mut Sender) {
         use task::Run::*;
 
-        match task.run(notify, sender) {
+        let run = self.run_task2(&task, notify, sender);
+
+        // TODO: Try to claim back the worker state in case the backup thread
+        // did not start up fast enough. This is a performance optimization.
+
+        match run {
             Idle => {}
             Schedule => {
-                self.entry().push_internal(task);
+                if self.is_blocking.get() {
+                    // The future has been notified while it was running.
+                    // However, the future also entered a blocking section,
+                    // which released the worker state from this thread.
+                    //
+                    // This means that scheduling the future must be done from
+                    // a point of view external to the worker set.
+                    //
+                    // We have to call `submit_external` instead of `submit`
+                    // here because `self` is still set as the current worker.
+                    self.inner.submit_external(task, &self.inner);
+                } else {
+                    self.entry().push_internal(task);
+                }
             }
             Complete => {
                 let mut state: pool::State = self.inner.state.load(Acquire).into();
@@ -351,6 +474,52 @@ impl Worker {
         }
     }
 
+    /// Actually run the task. This is where `Worker::current_task` is set.
+    ///
+    /// Great care is needed to ensure that `current_task` is unset in this
+    /// function.
+    fn run_task2(&self,
+                 task: &Arc<Task>,
+                 notify: &Arc<Notifier>,
+                 sender: &mut Sender)
+        -> task::Run
+    {
+        struct Guard<'a> {
+            worker: &'a Worker,
+            allocated_at_run: bool
+        }
+
+        impl<'a> Drop for Guard<'a> {
+            fn drop(&mut self) {
+                // A task is allocated at run when it was explicitly notified
+                // that the task has capacity to block. When this happens, that
+                // capacity is automatically allocated to the notified task.
+                // This capacity is "use it or lose it", so if the thread is not
+                // transitioned to blocking in this call, then another task has
+                // to be notified.
+                if self.allocated_at_run && !self.worker.is_blocking.get() {
+                    self.worker.inner.notify_blocking_task(&self.worker.inner);
+                }
+
+                self.worker.current_task.clear();
+            }
+        }
+
+        let can_block = task.consume_blocking_allocation();
+
+        // Set `current_task`
+        self.current_task.set(task, can_block);
+
+        // Create the guard, this ensures that `current_task` is unset when the
+        // function returns, even if the return is caused by a panic.
+        let _g = Guard {
+            worker: self,
+            allocated_at_run: can_block == CanBlock::Allocated
+        };
+
+        task.run(notify, sender)
+    }
+
     /// Drains all tasks on the extern queue and pushes them onto the internal
     /// queue.
     ///
@@ -396,7 +565,11 @@ impl Worker {
     fn sleep(&self) -> bool {
         use self::Lifecycle::*;
 
-        trace!("Worker::sleep; worker={:?}", self);
+        // Putting a worker to sleep is a multipart operation. This is, in part,
+        // due to the fact that a worker can be notified without it being popped
+        // from the sleep stack. Extra care is needed to deal with this.
+
+        trace!("Worker::sleep; worker={:?}", self.id);
 
         let mut state: State = self.entry().state.load(Acquire).into();
 
@@ -439,12 +612,12 @@ impl Worker {
                 if !state.is_pushed() {
                     debug_assert!(next.is_pushed());
 
-                    trace!("  sleeping -- push to stack; idx={}", self.id.idx);
+                    trace!("  sleeping -- push to stack; idx={}", self.id.0);
 
                     // We obtained permission to push the worker into the
                     // sleeper queue.
-                    if let Err(_) = self.inner.push_sleeper(self.id.idx) {
-                        trace!("  sleeping -- push to stack failed; idx={}", self.id.idx);
+                    if let Err(_) = self.inner.push_sleeper(self.id.0) {
+                        trace!("  sleeping -- push to stack failed; idx={}", self.id.0);
                         // The push failed due to the pool being terminated.
                         //
                         // This is true because the "work" being woken up for is
@@ -459,7 +632,7 @@ impl Worker {
             state = actual;
         }
 
-        trace!("    -> starting to sleep; idx={}", self.id.idx);
+        trace!("    -> starting to sleep; idx={}", self.id.0);
 
         let sleep_until = self.inner.config.keep_alive
             .map(|dur| Instant::now() + dur);
@@ -467,6 +640,7 @@ impl Worker {
         // The state has been transitioned to sleeping, we can now wait by
         // calling the parker. This is done in a loop as condvars can wakeup
         // spuriously.
+        'sleep:
         loop {
             let mut drop_thread = false;
 
@@ -495,14 +669,91 @@ impl Worker {
                 }
             }
 
-            trace!("    -> wakeup; idx={}", self.id.idx);
+            trace!("    -> wakeup; idx={}", self.id.0);
 
             // Reload the state
             state = self.entry().state.load(Acquire).into();
 
+            // If the worker has been notified, transition back to running.
+            match state.lifecycle() {
+                Sleeping => {
+                    if !drop_thread {
+                        // This goes back to the outer loop.
+                        continue 'sleep;
+                    }
+                }
+                Notified | Signaled => {
+                    // Transition back to running
+                    loop {
+                        let mut next = state;
+                        next.set_lifecycle(Running);
+
+                        let actual = self.entry().state.compare_and_swap(
+                            state.into(), next.into(), AcqRel).into();
+
+                        if actual == state {
+                            return true;
+                        }
+
+                        state = actual;
+                    }
+                }
+                Shutdown | Running => {
+                    // To get here, the block above transitioned the tate to
+                    // `Sleeping`. No other thread can concurrently
+                    // transition to `Shutdown` or `Running`.
+                    unreachable!();
+                }
+            }
+
+            // The thread has reached the maximum permitted sleep duration.
+            // It is now going to begin to shutdown.
+            //
+            // Doing this requires first releasing the thread to the backup
+            // stack. Because the moment the worker state is transitioned to
+            // `Shutdown`, other threads **expect** the thread's backup
+            // entry to be available on the backup stack.
+            //
+            // However, it is possible that the worker is notified between
+            // us pushing the backup entry onto the backup stack and
+            // transitioning the worker to `Shutdown`. If this happens, the
+            // current thread lost the token to run the backup entry and has
+            // to shutdown no matter what.
+            //
+            // To deal with this, the worker is transitioned to another
+            // thread. This is a pretty rare condition.
+            //
+            // If pushing on the backup stack fails, then the pool is being
+            // terminated and the thread should just shutdown
+            let backup_push_err = self.inner.release_backup(self.backup_id).is_err();
+
+            if backup_push_err {
+                debug_assert!({
+                    let state: State = self.entry().state.load(Acquire).into();
+                    state.lifecycle() != Sleeping
+                });
+
+                self.should_finalize.set(true);
+
+                return true;
+            }
+
             loop {
-                match state.lifecycle() {
-                    Sleeping => {}
+                let mut next = state;
+                next.set_lifecycle(Shutdown);
+
+                let actual: State = self.entry().state.compare_and_swap(
+                    state.into(), next.into(), AcqRel).into();
+
+                if actual == state {
+                    // Transitioned to a shutdown state
+                    return false;
+                }
+
+                match actual.lifecycle() {
+                    Sleeping => {
+                        state = actual;
+                    }
                     Notified | Signaled => {
                         // Transition back to running
                         loop {
@@ -513,7 +764,8 @@ impl Worker {
                                 state.into(), next.into(), AcqRel).into();
 
                             if actual == state {
-                                return true;
+                                self.inner.spawn_thread(self.id.clone(), &self.inner);
+                                return false;
                             }
 
                             state = actual;
@@ -526,27 +778,7 @@ impl Worker {
                         unreachable!();
                     }
                 }
-
-                if !drop_thread {
-                    // This goees back to the outer loop.
-                    break;
-                }
-
-                let mut next = state;
-                next.set_lifecycle(Shutdown);
-
-                let actual = self.entry().state.compare_and_swap(
-                    state.into(), next.into(), AcqRel).into();
-
-                if actual == state {
-                    // Transitioned to a shutdown state
-                    return false;
-                }
-
-                state = actual;
             }
-
-            // The worker hasn't been notified, go back to sleep
         }
     }
 
@@ -562,13 +794,14 @@ impl Worker {
     }
 
     fn entry(&self) -> &Entry {
-        &self.inner.workers[self.id.idx]
+        debug_assert!(!self.is_blocking.get());
+        &self.inner.workers[self.id.0]
     }
 }
 
 impl Drop for Worker {
     fn drop(&mut self) {
-        trace!("shutting down thread; idx={}", self.id.idx);
+        trace!("shutting down thread; idx={}", self.id.0);
 
         if self.should_finalize.get() {
             // Get all inbound work and push it onto the work queue. The work
@@ -579,13 +812,51 @@ impl Drop for Worker {
             self.entry().drain_tasks();
 
             // TODO: Drain the work queue...
-            self.inner.worker_terminated();
         }
     }
 }
 
-impl WorkerId {
-    pub(crate) fn new(idx: usize) -> WorkerId {
-        WorkerId { idx }
+// ===== impl CurrentTask =====
+
+impl CurrentTask {
+    /// Returns a default `CurrentTask` representing no task.
+    fn new() -> CurrentTask {
+        CurrentTask {
+            task: Cell::new(None),
+            can_block: Cell::new(CanBlock::CanRequest),
+        }
+    }
+
+    /// Returns a reference to the task.
+    fn get_ref(&self) -> &Arc<Task> {
+        unsafe { &*self.task.get().unwrap() }
+    }
+
+    fn can_block(&self) -> CanBlock {
+        self.can_block.get()
+    }
+
+    fn set_can_block(&self, can_block: CanBlock) {
+        self.can_block.set(can_block);
+    }
+
+    fn set(&self, task: &Arc<Task>, can_block: CanBlock) {
+        self.task.set(Some(task as *const _));
+        self.can_block.set(can_block);
+    }
+
+    /// Reset the `CurrentTask` to null state.
+    fn clear(&self) {
+        self.task.set(None);
+        self.can_block.set(CanBlock::CanRequest);
+    }
+}
+
+// ===== impl WorkerId =====
+
+impl WorkerId {
+    /// Returns a `WorkerId` representing the worker entry at index `idx`.
+    pub(crate) fn new(idx: usize) -> WorkerId {
+        WorkerId(idx)
     }
 }