diff --git a/tokio-threadpool/src/lib.rs b/tokio-threadpool/src/lib.rs index 13dfb5d2c..c13952273 100644 --- a/tokio-threadpool/src/lib.rs +++ b/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, } /// 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 } } diff --git a/tokio-threadpool/tests/threadpool.rs b/tokio-threadpool/tests/threadpool.rs index d8b97997f..4d86316e1 100644 --- a/tokio-threadpool/tests/threadpool.rs +++ b/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); + + 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;