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, BackupId}; use notifier::Notifier; use sender::Sender; use shutdown::ShutdownTrigger; 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; /// Thread worker /// /// This is passed to the [`around_worker`] callback set on [`Builder`]. This /// callback is only expected to call [`run`] on it. /// /// [`Builder`]: struct.Builder.html /// [`around_worker`]: struct.Builder.html#method.around_worker /// [`run`]: struct.Worker.html#method.run #[derive(Debug)] pub struct Worker { // Shared scheduler data pub(crate) pool: Arc, // 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, // Set when the worker should finalize on drop should_finalize: Cell, // Completes the shutdown process when the `ThreadPool` and all `Worker`s get dropped. trigger: Arc, // Keep the value on the current thread. _p: PhantomData>, } /// 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>>, /// Tracks the blocking capacity allocation state. can_block: Cell, } /// Identifies 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) 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 new( id: WorkerId, backup_id: BackupId, pool: Arc, trigger: Arc, ) -> Worker { Worker { pool, id, backup_id, current_task: CurrentTask::new(), is_blocking: Cell::new(false), should_finalize: Cell::new(false), trigger, _p: PhantomData, } } pub(crate) fn is_blocking(&self) -> bool { self.is_blocking.get() } /// 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 _); let pool = self.pool.clone(); let mut sender = Sender { pool }; // 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) = self.pool.config.around_worker { callback.call(self, enter); } else { self.run(); } }); }); // 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) -> 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 })) } }) } /// 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.pool.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.pool.spawn_thread(self.id.clone(), &self.pool); // 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 /// different thread pool instances share worker identifier values. pub fn id(&self) -> &WorkerId { &self.id } /// Run the worker /// /// This function blocks until the worker is shutting down. pub fn run(&self) { const MAX_SPINS: usize = 3; const LIGHT_SLEEP_INTERVAL: usize = 32; // Get the notifier. let notify = Arc::new(Notifier { pool: self.pool.clone(), }); let mut first = true; let mut spin_cnt = 0; let mut tick = 0; while self.check_run_state(first) { first = false; // Poll inbound until empty, transferring all tasks to the internal // queue. let consistent = self.drain_inbound(); // Run the next available task if self.try_run_task(¬ify) { if self.is_blocking.get() { // Exit out of the run state return; } if tick % LIGHT_SLEEP_INTERVAL == 0 { self.sleep_light(); } tick = tick.wrapping_add(1); spin_cnt = 0; // As long as there is work, keep looping. continue; } if !consistent { spin_cnt = 0; continue; } spin_cnt += 1; // Yield the thread several times before it actually goes to sleep. if spin_cnt <= MAX_SPINS { thread::yield_now(); continue; } tick = 0; spin_cnt = 0; // Starting to get sleeeeepy if !self.sleep() { return; } // 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.pool.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) -> bool { if self.try_run_owned_task(notify) { return true; } self.try_steal_task(notify) } /// Checks the worker's current state, updating it as needed. /// /// Returns `true` if the worker should run. #[inline] 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 { let pool_state: pool::State = self.pool.state.load(Acquire).into(); if pool_state.is_terminated() { return false; } let mut next = state; match state.lifecycle() { Running => break, Notified | Signaled => { // transition back to running next.set_lifecycle(Running); } Shutdown | Sleeping => { // The worker should never be in these states when calling // this function. panic!("unexpected worker state; lifecycle={:?}", state.lifecycle()); } } let actual = self.entry().state.compare_and_swap( state.into(), next.into(), AcqRel).into(); if actual == state { break; } state = actual; } // `first` is set to true the first time this function is called after // the thread has started. // // This check is to handle the scenario where a worker gets signaled // while it is already happily running. The `is_signaled` state is // intended to wake up a worker that has been previously sleeping in // effect increasing the number of active workers. If this is the first // time `check_run_state` is called, then being in a signalled state is // normal and the thread was started to handle it. However, if this is // **not** the first time the fn was called, then the number of active // workers has not been increased by the signal, so `signal_work` has to // be called again to try to wake up another worker. // // For example, if the thread pool is configured to allow 4 workers. // Worker 1 is processing tasks from its `deque`. Worker 2 receives its // first task. Worker 2 will pick a random worker to signal. It does // this by popping off the sleep stack, but there is no guarantee that // workers on the sleep stack are actually sleeping. It is possible that // Worker 1 gets signaled. // // Without this check, in the above case, no additional workers will get // started, which results in the thread pool permanently being at 2 // workers even though it should reach 4. if !first && state.is_signaled() { trace!("Worker::check_run_state; delegate signal"); // This worker is not ready to be signaled, so delegate the signal // to another worker. self.pool.signal_work(&self.pool); } true } /// Runs the next task on this worker's queue. /// /// Returns `true` if work was found. fn try_run_owned_task(&self, notify: &Arc) -> bool { use deque::Pop; // Poll the internal queue for a task to run match self.entry().pop_task() { Pop::Data(task) => { self.run_task(task, notify); true } Pop::Empty => false, Pop::Retry => true, } } /// Tries to steal a task from another worker. /// /// Returns `true` if work was found fn try_steal_task(&self, notify: &Arc) -> bool { use deque::Steal; debug_assert!(!self.is_blocking.get()); let len = self.pool.workers.len(); let mut idx = self.pool.rand_usize() % len; let mut found_work = false; let start = idx; loop { if idx < len { match self.pool.workers[idx].steal_tasks(self.entry()) { Steal::Data(task) => { trace!("stole task"); self.run_task(task, notify); trace!("try_steal_task -- signal_work; self={}; from={}", self.id.0, idx); // Signal other workers that work is available // // TODO: Should this be called here or before // `run_task`? self.pool.signal_work(&self.pool); return true; } Steal::Empty => {} Steal::Retry => found_work = true, } idx += 1; } else { idx = 0; } if idx == start { break; } } found_work } fn run_task(&self, task: Arc, notify: &Arc) { use task::Run::*; let run = self.run_task2(&task, notify); // 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 => { 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.pool.submit_external(task, &self.pool); } else { self.entry().push_internal(task); } } Complete => { let mut state: pool::State = self.pool.state.load(Acquire).into(); loop { let mut next = state; next.dec_num_futures(); let actual = self.pool.state.compare_and_swap( state.into(), next.into(), AcqRel).into(); if actual == state { trace!("task complete; state={:?}", next); if state.num_futures() == 1 { // If the thread pool has been flagged as shutdown, // start terminating workers. This involves waking // up any sleeping worker so that they can notice // the shutdown state. if next.is_terminated() { self.pool.terminate_sleeping_workers(); } } // The worker's run loop will detect the shutdown state // next iteration. return; } state = actual; } } } } /// 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, notify: &Arc) -> 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.pool.notify_blocking_task(&self.worker.pool); } 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) } /// Drains all tasks on the extern queue and pushes them onto the internal /// queue. /// /// Returns `true` if the operation was able to complete in a consistent /// state. #[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 { // TODO: Why is this called on every iteration? Would it // not be better to only signal when work was found // after waking up? trace!("found work while draining; signal_work"); self.pool.signal_work(&self.pool); } return true; } Inconsistent => { if found_work { trace!("found work while draining; signal_work"); self.pool.signal_work(&self.pool); } return false; } Data(task) => { found_work = true; self.entry().push_internal(task); } } } } /// Put the worker to sleep /// /// Returns `true` if woken up due to new work arriving. fn sleep(&self) -> bool { use self::Lifecycle::*; // 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(); // The first part of the sleep process is to transition the worker state // to "pushed". Now, it may be that the worker is already pushed on the // sleeper stack, in which case, we don't push again. loop { let mut next = state; match state.lifecycle() { Running => { // Try setting the pushed state next.set_pushed(); // Transition the worker state to sleeping next.set_lifecycle(Sleeping); } Notified | Signaled => { // No need to sleep, transition back to running and move on. next.set_lifecycle(Running); } Shutdown | Sleeping => { // The worker cannot transition to sleep when already in a // sleeping state. panic!("unexpected worker state; actual={:?}", state.lifecycle()); } } let actual = self.entry().state.compare_and_swap( state.into(), next.into(), AcqRel).into(); if actual == state { if state.is_notified() { // The previous state was notified, so we don't need to // sleep. return true; } if !state.is_pushed() { debug_assert!(next.is_pushed()); trace!(" sleeping -- push to stack; idx={}", self.id.0); // We obtained permission to push the worker into the // sleeper queue. if let Err(_) = self.pool.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 // shutting down. return true; } } break; } state = actual; } trace!(" -> starting to sleep; idx={}", self.id.0); // 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. loop { unsafe { (*self.entry().park.get()) .park() .unwrap(); } 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 => { // Still sleeping. Park again. } 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 state to // `Sleeping`. No other thread can concurrently // transition to `Shutdown` or `Running`. unreachable!(); } } } } /// This doesn't actually put the thread to sleep. It calls /// `park.park_timeout` with a duration of 0. This allows the park /// implementation to perform any work that might be done on an interval. fn sleep_light(&self) { unsafe { (*self.entry().park.get()) .park_timeout(Duration::from_millis(0)) .unwrap(); } } fn entry(&self) -> &Entry { debug_assert!(!self.is_blocking.get()); &self.pool.workers[self.id.0] } } impl Drop for Worker { fn drop(&mut self) { 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 // queue is drained in the next step. self.drain_inbound(); // Drain the work queue self.entry().drain_tasks(); // TODO: Drain the work queue... } } } // ===== 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 { 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, 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) } /// Returns this identifier represented as an integer. /// /// Worker identifiers in a single thread pool are guaranteed to correspond to integers in the /// range `0..pool_size`. pub fn to_usize(&self) -> usize { self.0 } }