Shutdown the thread pool on drop. (#190)

Currently, if a thread pool instance is dropped without being shutdown,
the workers will run indefinitely. This is not ideal as it leaks the
threadpool.

This patch forces the thread pool to shutdown on drop.

Closes #151

tokio-threadpool/src/lib.rs

@@ -34,10 +34,14 @@ use std::time::{Instant, Duration};
 
 /// Work-stealing based thread pool for executing futures.
 ///
+/// If a `ThreadPool` instance is dropped without explicitly being shutdown,
+/// `shutdown_now` is called implicitly, forcing all tasks that have not yet
+/// completed to be dropped.
+///
 /// Create `ThreadPool` instances using `Builder`.
 #[derive(Debug)]
 pub struct ThreadPool {
-    inner: Sender,
+    inner: Option<Sender>,
 }
 
 /// Submit futures to the associated thread pool for execution.
@@ -72,7 +76,7 @@ pub struct Sender {
 /// [`shutdown_now`]: struct.ThreadPool.html#method.shutdown_now
 #[derive(Debug)]
 pub struct Shutdown {
-    inner: ThreadPool,
+    inner: Sender,
 }
 
 /// Builds a thread pool with custom configuration values.
@@ -536,7 +540,7 @@ impl Builder {
             inner.push_sleeper(i).unwrap();
         }
 
-        let inner = Sender { inner };
+        let inner = Some(Sender { inner });
 
         ThreadPool { inner }
     }
@@ -596,12 +600,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.inner
+        self.inner.as_ref().unwrap()
     }
 
     /// Return a mutable reference to the sender handle
     pub fn sender_mut(&mut self) -> &mut Sender {
-        &mut self.inner
+        self.inner.as_mut().unwrap()
     }
 
     /// Shutdown the pool once it becomes idle.
@@ -614,9 +618,9 @@ impl ThreadPool {
     /// handle will result in an error. All worker threads are signaled and will
     /// shutdown. The returned future completes once all worker threads have
     /// completed the shutdown process.
-    pub fn shutdown_on_idle(self) -> Shutdown {
+    pub fn shutdown_on_idle(mut self) -> Shutdown {
         self.inner().shutdown(false, false);
-        Shutdown { inner: self }
+        Shutdown { inner: self.inner.take().unwrap() }
     }
 
     /// Shutdown the pool
@@ -627,9 +631,9 @@ impl ThreadPool {
     /// Calling `spawn` on any outstanding handle will result in an error. All
     /// worker threads are signaled and will shutdown. The returned future
     /// completes once all worker threads have completed the shutdown process.
-    pub fn shutdown(self) -> Shutdown {
+    pub fn shutdown(mut self) -> Shutdown {
         self.inner().shutdown(true, false);
-        Shutdown { inner: self }
+        Shutdown { inner: self.inner.take().unwrap() }
     }
 
     /// Shutdown the pool immediately
@@ -640,13 +644,23 @@ impl ThreadPool {
     /// Calling `spawn` on any outstanding handle will result in an error. All
     /// worker threads are signaled and will shutdown. The returned future
     /// completes once all worker threads have completed the shutdown process.
-    pub fn shutdown_now(self) -> Shutdown {
+    pub fn shutdown_now(mut self) -> Shutdown {
         self.inner().shutdown(true, true);
-        Shutdown { inner: self }
+        Shutdown { inner: self.inner.take().unwrap() }
     }
 
     fn inner(&self) -> &Inner {
-        &*self.inner.inner
+        &*self.inner.as_ref().unwrap().inner
+    }
+}
+
+impl Drop for ThreadPool {
+    fn drop(&mut self) {
+        if let Some(sender) = self.inner.take() {
+            sender.inner.shutdown(true, true);
+            let shutdown = Shutdown { inner: sender };
+            let _ = shutdown.wait();
+        }
     }
 }
 
@@ -824,7 +838,7 @@ impl Clone for Sender {
 
 impl Shutdown {
     fn inner(&self) -> &Inner {
-        self.inner.inner()
+        &*self.inner.inner
     }
 }
 

tokio-threadpool/tests/threadpool.rs

@@ -135,6 +135,63 @@ fn force_shutdown_drops_futures() {
     }
 }
 
+#[test]
+fn drop_threadpool_drops_futures() {
+    let _ = ::env_logger::init();
+
+    for _ in 0..1_000 {
+        let num_inc = Arc::new(AtomicUsize::new(0));
+        let num_dec = Arc::new(AtomicUsize::new(0));
+        let num_drop = Arc::new(AtomicUsize::new(0));
+
+        struct Never(Arc<AtomicUsize>);
+
+        impl Future for Never {
+            type Item = ();
+            type Error = ();
+
+            fn poll(&mut self) -> Poll<(), ()> {
+                Ok(Async::NotReady)
+            }
+        }
+
+        impl Drop for Never {
+            fn drop(&mut self) {
+                self.0.fetch_add(1, Relaxed);
+            }
+        }
+
+        let a = num_inc.clone();
+        let b = num_dec.clone();
+
+        let mut pool = Builder::new()
+            .around_worker(move |w, _| {
+                a.fetch_add(1, Relaxed);
+                w.run();
+                b.fetch_add(1, Relaxed);
+            })
+            .build();
+        let mut tx = pool.sender().clone();
+
+        tx.spawn(Never(num_drop.clone())).unwrap();
+
+        // Wait for the pool to shutdown
+        drop(pool);
+
+        // Assert that only a single thread was spawned.
+        let a = num_inc.load(Relaxed);
+        assert!(a >= 1);
+
+        // Assert that all threads shutdown
+        let b = num_dec.load(Relaxed);
+        assert_eq!(a, b);
+
+        // Assert that the future was dropped
+        let c = num_drop.load(Relaxed);
+        assert_eq!(c, 1);
+    }
+}
+
 #[test]
 fn thread_shutdown_timeout() {
     use std::sync::Mutex;