threadpool: refactor pool shutdown (#769)

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)
     }
 }
 

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;

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<ShutdownTrigger>,
 
     // 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<ShutdownTrigger>,
+        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<Task>) -> 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 {

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<Mutex<Inner>>,
+}
+
+/// 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<Mutex<Inner>>,
+    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();
     }
 }

tokio-threadpool/src/shutdown_task.rs

@@ -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();
-    }
-}

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<Sender>,
+    inner: Option<Inner>,
+}
+
+#[derive(Debug)]
+struct Inner {
+    sender: Sender,
+    trigger: Arc<ShutdownTrigger>,
 }
 
 impl ThreadPool {
@@ -28,6 +36,18 @@ impl ThreadPool {
         Builder::new().build()
     }
 
+    pub(crate) fn new2(
+        pool: Arc<Pool>,
+        trigger: Arc<ShutdownTrigger>,
+    ) -> 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();
         }
     }

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<bool>,
 
+    // Completes the shutdown process when the `ThreadPool` and all `Worker`s get dropped.
+    trigger: Arc<ShutdownTrigger>,
+
     // Keep the value on the current thread.
     _p: PhantomData<Rc<()>>,
 }
@@ -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<Pool>) -> Worker {
+    pub(crate) fn new(
+        id: WorkerId,
+        backup_id: BackupId,
+        pool: Arc<Pool>,
+        trigger: Arc<ShutdownTrigger>,
+    ) -> 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,
         }
     }