diff --git a/tokio-threadpool/src/builder.rs b/tokio-threadpool/src/builder.rs index ca701f2c1..05af18204 100644 --- a/tokio-threadpool/src/builder.rs +++ b/tokio-threadpool/src/builder.rs @@ -1,7 +1,7 @@ use callback::Callback; use config::{Config, MAX_WORKERS}; use park::{BoxPark, BoxedPark, DefaultPark}; -use sender::Sender; +use shutdown::ShutdownTrigger; use pool::{Pool, MAX_BACKUP}; use thread_pool::ThreadPool; use worker::{self, Worker, WorkerId}; @@ -396,31 +396,38 @@ impl Builder { /// # } /// ``` pub fn build(&self) -> ThreadPool { - let mut workers = vec![]; - trace!("build; num-workers={}", self.pool_size); - for i in 0..self.pool_size { - let id = WorkerId::new(i); - let park = (self.new_park)(&id); - let unpark = park.unpark(); + // Create the worker entry list + let workers: Arc<[worker::Entry]> = { + let mut workers = vec![]; - workers.push(worker::Entry::new(park, unpark)); - } + for i in 0..self.pool_size { + let id = WorkerId::new(i); + let park = (self.new_park)(&id); + let unpark = park.unpark(); + + workers.push(worker::Entry::new(park, unpark)); + } + + workers.into() + }; + + // Create a trigger that will clean up resources on shutdown. + // + // The `Pool` contains a weak reference to it, while `Worker`s and the `ThreadPool` contain + // strong references. + let trigger = Arc::new(ShutdownTrigger::new(workers.clone())); // Create the pool - let pool = Arc::new( - Pool::new( - workers.into_boxed_slice(), - self.max_blocking, - self.config.clone())); + let pool = Arc::new(Pool::new( + workers, + Arc::downgrade(&trigger), + self.max_blocking, + self.config.clone(), + )); - // Wrap with `Sender` - let sender = Some(Sender { - pool - }); - - ThreadPool { sender } + ThreadPool::new2(pool, trigger) } } diff --git a/tokio-threadpool/src/lib.rs b/tokio-threadpool/src/lib.rs index 604214131..ec8d3a368 100644 --- a/tokio-threadpool/src/lib.rs +++ b/tokio-threadpool/src/lib.rs @@ -149,7 +149,6 @@ mod notifier; mod pool; mod sender; mod shutdown; -mod shutdown_task; mod task; mod thread_pool; mod worker; diff --git a/tokio-threadpool/src/pool/mod.rs b/tokio-threadpool/src/pool/mod.rs index c6dce5473..52dbf8051 100644 --- a/tokio-threadpool/src/pool/mod.rs +++ b/tokio-threadpool/src/pool/mod.rs @@ -14,18 +14,17 @@ use self::backup::Handoff; use self::backup_stack::BackupStack; use config::Config; -use shutdown_task::ShutdownTask; +use shutdown::ShutdownTrigger; use task::{Task, Blocking}; use worker::{self, Worker, WorkerId}; use futures::Poll; -use futures::task::AtomicTask; use std::cell::Cell; use std::num::Wrapping; use std::sync::atomic::Ordering::{Acquire, AcqRel}; use std::sync::atomic::AtomicUsize; -use std::sync::Arc; +use std::sync::{Arc, Weak}; use std::thread; use crossbeam_utils::CachePadded; @@ -57,8 +56,14 @@ pub(crate) struct Pool { // 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]>, + // The number of workers will *usually* be small. + pub workers: Arc<[worker::Entry]>, + + // Completes the shutdown process when the `ThreadPool` and all `Worker`s get dropped. + // + // When spawning a new `Worker`, this weak reference is upgraded and handed out to the new + // thread. + pub trigger: Weak, // Backup thread state // @@ -74,18 +79,18 @@ pub(crate) struct Pool { // are pending blocking capacity. blocking: Blocking, - // Task notified when the worker shuts down - pub shutdown_task: ShutdownTask, - // Configuration pub config: Config, } -const TERMINATED: usize = 1; - impl Pool { /// Create a new `Pool` - pub fn new(workers: Box<[worker::Entry]>, max_blocking: usize, config: Config) -> Pool { + pub fn new( + workers: Arc<[worker::Entry]>, + trigger: Weak, + max_blocking: usize, + config: Config, + ) -> Pool { let pool_size = workers.len(); let total_size = max_blocking + pool_size; @@ -112,12 +117,10 @@ impl Pool { sleep_stack: CachePadded::new(worker::Stack::new()), num_workers: AtomicUsize::new(0), workers, + trigger, backup, backup_stack, blocking, - shutdown_task: ShutdownTask { - task: AtomicTask::new(), - }, config, }; @@ -191,10 +194,6 @@ 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. @@ -205,12 +204,6 @@ 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. @@ -226,40 +219,6 @@ impl Pool { 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(); - } - } - - /// 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); - - loop { - if curr & TERMINATED == TERMINATED { - return Err(()); - } - - 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) -> Poll<(), ::BlockingError> { @@ -385,10 +344,14 @@ impl Pool { return; } - if self.thread_started().is_err() { + let trigger = match self.trigger.upgrade() { // The pool is shutting down. - return; - } + None => { + // The pool is shutting down. + return; + } + Some(t) => t, + }; let mut th = thread::Builder::new(); @@ -416,7 +379,7 @@ impl Pool { debug_assert!(pool.backup[backup_id.0].is_running()); // TODO: Avoid always cloning - let worker = Worker::new(worker_id, backup_id, pool.clone()); + let worker = Worker::new(worker_id, backup_id, pool.clone(), trigger.clone()); // Run the worker. If the worker transitioned to a "blocking" // state, then `is_blocking` will be true. @@ -463,8 +426,6 @@ impl Pool { if let Some(ref f) = pool.config.before_stop { f(); } - - pool.thread_stopped(); }); if let Err(e) = res { diff --git a/tokio-threadpool/src/shutdown.rs b/tokio-threadpool/src/shutdown.rs index 957276079..8167d0148 100644 --- a/tokio-threadpool/src/shutdown.rs +++ b/tokio-threadpool/src/shutdown.rs @@ -1,7 +1,9 @@ -use pool::Pool; -use sender::Sender; +use worker; use futures::{Future, Poll, Async}; +use futures::task::AtomicTask; + +use std::sync::{Arc, Mutex}; /// Future that resolves when the thread pool is shutdown. /// @@ -16,12 +18,25 @@ use futures::{Future, Poll, Async}; /// [`shutdown_now`]: struct.ThreadPool.html#method.shutdown_now #[derive(Debug)] pub struct Shutdown { - pub(crate) sender: Sender, + inner: Arc>, +} + +/// Shared state between `Shutdown` and `ShutdownTrigger`. +/// +/// This is used for notifying the `Shutdown` future when `ShutdownTrigger` gets dropped. +#[derive(Debug)] +struct Inner { + /// The task to notify when the threadpool completes the shutdown process. + task: AtomicTask, + /// `true` if the threadpool has been shut down. + completed: bool, } impl Shutdown { - fn pool(&self) -> &Pool { - &*self.sender.pool + pub(crate) fn new(trigger: &ShutdownTrigger) -> Shutdown { + Shutdown { + inner: trigger.inner.clone(), + } } } @@ -30,14 +45,44 @@ impl Future for Shutdown { type Error = (); fn poll(&mut self) -> Poll<(), ()> { - use futures::task; + let inner = self.inner.lock().unwrap(); - self.pool().shutdown_task.task.register_task(task::current()); - - if !self.pool().is_shutdown() { - return Ok(Async::NotReady); + if !inner.completed { + inner.task.register(); + Ok(Async::NotReady) + } else { + Ok(().into()) } - - Ok(().into()) + } +} + +/// When dropped, cleans up threadpool's resources and completes the shutdown process. +#[derive(Debug)] +pub(crate) struct ShutdownTrigger { + inner: Arc>, + workers: Arc<[worker::Entry]>, +} + +unsafe impl Send for ShutdownTrigger {} +unsafe impl Sync for ShutdownTrigger {} + +impl ShutdownTrigger { + pub(crate) fn new(workers: Arc<[worker::Entry]>) -> ShutdownTrigger { + ShutdownTrigger { + inner: Arc::new(Mutex::new(Inner { + task: AtomicTask::new(), + completed: false, + })), + workers, + } + } +} + +impl Drop for ShutdownTrigger { + fn drop(&mut self) { + // Notify the task interested in shutdown. + let mut inner = self.inner.lock().unwrap(); + inner.completed = true; + inner.task.notify(); } } diff --git a/tokio-threadpool/src/shutdown_task.rs b/tokio-threadpool/src/shutdown_task.rs deleted file mode 100644 index 3dcbb4284..000000000 --- a/tokio-threadpool/src/shutdown_task.rs +++ /dev/null @@ -1,12 +0,0 @@ -use futures::task::AtomicTask; - -#[derive(Debug)] -pub(crate) struct ShutdownTask { - pub task: AtomicTask, -} - -impl ShutdownTask { - pub fn notify(&self) { - self.task.notify(); - } -} diff --git a/tokio-threadpool/src/thread_pool.rs b/tokio-threadpool/src/thread_pool.rs index c9e772921..960ffb780 100644 --- a/tokio-threadpool/src/thread_pool.rs +++ b/tokio-threadpool/src/thread_pool.rs @@ -1,11 +1,13 @@ use builder::Builder; use pool::Pool; use sender::Sender; -use shutdown::Shutdown; +use shutdown::{Shutdown, ShutdownTrigger}; use futures::{Future, Poll}; use futures::sync::oneshot; +use std::sync::Arc; + /// Work-stealing based thread pool for executing futures. /// /// If a `ThreadPool` instance is dropped without explicitly being shutdown, @@ -15,7 +17,13 @@ use futures::sync::oneshot; /// Create `ThreadPool` instances using `Builder`. #[derive(Debug)] pub struct ThreadPool { - pub(crate) sender: Option, + inner: Option, +} + +#[derive(Debug)] +struct Inner { + sender: Sender, + trigger: Arc, } impl ThreadPool { @@ -28,6 +36,18 @@ impl ThreadPool { Builder::new().build() } + pub(crate) fn new2( + pool: Arc, + trigger: Arc, + ) -> ThreadPool { + ThreadPool { + inner: Some(Inner { + sender: Sender { pool }, + trigger, + }), + } + } + /// Spawn a future onto the thread pool. /// /// This function takes ownership of the future and randomly assigns it to a @@ -111,12 +131,12 @@ impl ThreadPool { /// The handle is used to spawn futures onto the thread pool. It also /// implements the `Executor` trait. pub fn sender(&self) -> &Sender { - self.sender.as_ref().unwrap() + &self.inner.as_ref().unwrap().sender } /// Return a mutable reference to the sender handle pub fn sender_mut(&mut self) -> &mut Sender { - self.sender.as_mut().unwrap() + &mut self.inner.as_mut().unwrap().sender } /// Shutdown the pool once it becomes idle. @@ -130,8 +150,9 @@ impl ThreadPool { /// shutdown. The returned future completes once all worker threads have /// completed the shutdown process. pub fn shutdown_on_idle(mut self) -> Shutdown { - self.pool().shutdown(false, false); - Shutdown { sender: self.sender.take().unwrap() } + let inner = self.inner.take().unwrap(); + inner.sender.pool.shutdown(false, false); + Shutdown::new(&inner.trigger) } /// Shutdown the pool @@ -143,8 +164,9 @@ impl ThreadPool { /// worker threads are signaled and will shutdown. The returned future /// completes once all worker threads have completed the shutdown process. pub fn shutdown(mut self) -> Shutdown { - self.pool().shutdown(true, false); - Shutdown { sender: self.sender.take().unwrap() } + let inner = self.inner.take().unwrap(); + inner.sender.pool.shutdown(true, false); + Shutdown::new(&inner.trigger) } /// Shutdown the pool immediately @@ -156,20 +178,23 @@ impl ThreadPool { /// worker threads are signaled and will shutdown. The returned future /// completes once all worker threads have completed the shutdown process. pub fn shutdown_now(mut self) -> Shutdown { - self.pool().shutdown(true, true); - Shutdown { sender: self.sender.take().unwrap() } - } - - fn pool(&self) -> &Pool { - &*self.sender.as_ref().unwrap().pool + let inner = self.inner.take().unwrap(); + inner.sender.pool.shutdown(true, true); + Shutdown::new(&inner.trigger) } } impl Drop for ThreadPool { fn drop(&mut self) { - if let Some(sender) = self.sender.take() { - sender.pool.shutdown(true, true); - let shutdown = Shutdown { sender }; + if let Some(inner) = self.inner.take() { + // Begin the shutdown process. + inner.sender.pool.shutdown(true, true); + let shutdown = Shutdown::new(&inner.trigger); + + // Drop `inner` in order to drop its shutdown trigger. + drop(inner); + + // Wait until all worker threads terminate and the threadpool's resources clean up. let _ = shutdown.wait(); } } diff --git a/tokio-threadpool/src/worker/mod.rs b/tokio-threadpool/src/worker/mod.rs index 9999d8639..cf259f767 100644 --- a/tokio-threadpool/src/worker/mod.rs +++ b/tokio-threadpool/src/worker/mod.rs @@ -14,6 +14,7 @@ pub(crate) use self::state::{ use pool::{self, Pool, BackupId}; use notifier::Notifier; use sender::Sender; +use shutdown::ShutdownTrigger; use task::{self, Task, CanBlock}; use tokio_executor; @@ -60,6 +61,9 @@ pub struct Worker { // 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>, } @@ -86,7 +90,12 @@ pub struct WorkerId(pub(crate) usize); 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) -> Worker { + pub(crate) fn new( + id: WorkerId, + backup_id: BackupId, + pool: Arc, + trigger: Arc, + ) -> Worker { Worker { pool, id, @@ -94,6 +103,7 @@ impl Worker { current_task: CurrentTask::new(), is_blocking: Cell::new(false), should_finalize: Cell::new(false), + trigger, _p: PhantomData, } }