threadpool: drop incomplete tasks on shutdown (#722)

## Motivation

When the thread pool shuts down, futures that have been polled at least once but not completed yet are simply leaked. We should drop them instead.

## Solution

Multiple changes are introduced:

* Tasks are assigned a home worker the first time they are polled.

* Each worker contains a set of tasks (`Arc<Task>`) it is home to. When a task is assigned a home worker, it is registered in that worker's set of tasks. When the task is completed, it is unregistered from the set.

* When the thread pool shuts down and after all worker threads stop, the remaining tasks in workers' sets are aborted, i.e. they are switched to the `Aborted` state and their `Future`s are dropped.

* The thread pool shutdown process is refactored to make it more robust. We don't  track the number of active threads manually anymore. Instead, there's  `Arc<ShutdownTrigger>` that aborts remaining tasks and completes the `Shutdown` future once it gets destroyed (when all `Worker`s and `ThreadPool` get dropped because they're the only ones to contain strong references to the `ShutdownTrigger`).

Closes #424 
Closes #428

tokio-threadpool/src/worker/entry.rs

@@ -5,11 +5,14 @@ use worker::state::{State, PUSHED_MASK};
 use std::cell::UnsafeCell;
 use std::fmt;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::atomic::Ordering::{AcqRel, Relaxed};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+use std::sync::atomic::Ordering::{Acquire, AcqRel, Relaxed, Release};
+use std::time::Duration;
 
+use crossbeam::queue::SegQueue;
 use crossbeam_utils::CachePadded;
 use deque;
+use slab::Slab;
 
 // TODO: None of the fields should be public
 //
@@ -32,10 +35,20 @@ pub(crate) struct WorkerEntry {
     stealer: deque::Stealer<Arc<Task>>,
 
     // Thread parker
-    pub park: UnsafeCell<BoxPark>,
+    park: UnsafeCell<Option<BoxPark>>,
 
     // Thread unparker
-    pub unpark: BoxUnpark,
+    unpark: UnsafeCell<Option<BoxUnpark>>,
+
+    // Tasks that have been first polled by this worker, but not completed yet.
+    running_tasks: UnsafeCell<Slab<Arc<Task>>>,
+
+    // Tasks that have been first polled by this worker, but completed by another worker.
+    remotely_completed_tasks: SegQueue<Arc<Task>>,
+
+    // Set to `true` when `remotely_completed_tasks` has tasks that need to be removed from
+    // `running_tasks`.
+    needs_drain: AtomicBool,
 }
 
 impl WorkerEntry {
@@ -47,8 +60,11 @@ impl WorkerEntry {
             next_sleeper: UnsafeCell::new(0),
             worker: w,
             stealer: s,
-            park: UnsafeCell::new(park),
-            unpark,
+            park: UnsafeCell::new(Some(park)),
+            unpark: UnsafeCell::new(Some(unpark)),
+            running_tasks: UnsafeCell::new(Slab::new()),
+            remotely_completed_tasks: SegQueue::new(),
+            needs_drain: AtomicBool::new(false),
         }
     }
 
@@ -100,7 +116,7 @@ impl WorkerEntry {
             Sleeping => {
                 // The worker is currently sleeping, the condition variable must
                 // be signaled
-                self.wakeup();
+                self.unpark();
                 true
             }
             Shutdown => false,
@@ -163,7 +179,7 @@ impl WorkerEntry {
         }
 
         // Wakeup the worker
-        self.wakeup();
+        self.unpark();
     }
 
     /// Pop a task
@@ -202,14 +218,94 @@ impl WorkerEntry {
         }
     }
 
+    /// Parks the worker thread.
+    pub fn park(&self) {
+        if let Some(park) = unsafe { (*self.park.get()).as_mut() } {
+            park.park().unwrap();
+        }
+    }
+
+    /// Parks the worker thread for at most `duration`.
+    pub fn park_timeout(&self, duration: Duration) {
+        if let Some(park) = unsafe { (*self.park.get()).as_mut() } {
+            park.park_timeout(duration).unwrap();
+        }
+    }
+
+    /// Unparks the worker thread.
     #[inline]
-    pub fn push_internal(&self, task: Arc<Task>) {
-        self.worker.push(task);
+    pub fn unpark(&self) {
+        if let Some(park) = unsafe { (*self.unpark.get()).as_ref() } {
+            park.unpark();
+        }
+    }
+
+    /// Registers a task in this worker.
+    ///
+    /// Called when the task is being polled for the first time.
+    #[inline]
+    pub fn register_task(&self, task: &Arc<Task>) {
+        let running_tasks = unsafe { &mut *self.running_tasks.get() };
+
+        let key = running_tasks.insert(task.clone());
+        task.reg_index.set(key);
+    }
+
+    /// Unregisters a task from this worker.
+    ///
+    /// Called when the task is completed and was previously registered in this worker.
+    #[inline]
+    pub fn unregister_task(&self, task: Arc<Task>) {
+        let running_tasks = unsafe { &mut *self.running_tasks.get() };
+        running_tasks.remove(task.reg_index.get());
+        self.drain_remotely_completed_tasks();
+    }
+
+    /// Unregisters a task from this worker.
+    ///
+    /// Called when the task is completed by another worker and was previously registered in this
+    /// worker.
+    #[inline]
+    pub fn remotely_complete_task(&self, task: Arc<Task>) {
+        self.remotely_completed_tasks.push(task);
+        self.needs_drain.store(true, Release);
+    }
+
+    /// Drops the remaining incomplete tasks and the parker associated with this worker.
+    ///
+    /// This function is called by the shutdown trigger.
+    pub fn shutdown(&self) {
+        self.drain_remotely_completed_tasks();
+
+        // Abort all incomplete tasks.
+        let running_tasks = unsafe { &mut *self.running_tasks.get() };
+        for (_, task) in running_tasks.iter() {
+            task.abort();
+        }
+        running_tasks.clear();
+
+        // Drop the parker.
+        unsafe {
+            *self.park.get() = None;
+            *self.unpark.get() = None;
+        }
+    }
+
+    /// Drains the `remotely_completed_tasks` queue and removes tasks from `running_tasks`.
+    #[inline]
+    fn drain_remotely_completed_tasks(&self) {
+        if self.needs_drain.compare_and_swap(true, false, Acquire) {
+            let running_tasks = unsafe { &mut *self.running_tasks.get() };
+
+            while let Some(task) = self.remotely_completed_tasks.try_pop() {
+                running_tasks.remove(task.reg_index.get());
+            }
+        }
     }
 
     #[inline]
-    pub fn wakeup(&self) {
-        self.unpark.unpark();
+    pub fn push_internal(&self, task: Arc<Task>) {
+        self.worker.push(task);
     }
 
     #[inline]

tokio-threadpool/src/worker/mod.rs

@@ -451,6 +451,13 @@ impl Worker {
     fn run_task(&self, task: Arc<Task>, notify: &Arc<Notifier>) {
         use task::Run::*;
 
+        // If this is the first time this task is being polled, register it so that we can keep
+        // track of tasks that are in progress.
+        if task.reg_worker.get().is_none() {
+            task.reg_worker.set(Some(self.id.0 as u32));
+            self.entry().register_task(&task);
+        }
+
         let run = self.run_task2(&task, notify);
 
         // TODO: Try to claim back the worker state in case the backup thread
@@ -497,6 +504,16 @@ impl Worker {
                             }
                         }
 
+                        // Find which worker polled this task first.
+                        let worker = task.reg_worker.get().unwrap() as usize;
+
+                        // Unregister the task from the worker it was registered in.
+                        if !self.is_blocking.get() && worker == self.id.0 {
+                            self.entry().unregister_task(task);
+                        } else {
+                            self.pool.workers[worker].remotely_complete_task(task);
+                        }
+
                         // The worker's run loop will detect the shutdown state
                         // next iteration.
                         return;
@@ -672,11 +689,7 @@ impl Worker {
                 }
             }
 
-            unsafe {
-                (*self.entry().park.get())
-                    .park()
-                    .unwrap();
-            }
+            self.entry().park();
 
             trace!("    -> wakeup; idx={}", self.id.0);
         }
@@ -690,11 +703,7 @@ impl Worker {
     fn sleep_light(&self) {
         const STEAL_COUNT: usize = 32;
 
-        unsafe {
-            (*self.entry().park.get())
-                .park_timeout(Duration::from_millis(0))
-                .unwrap();
-        }
+        self.entry().park_timeout(Duration::from_millis(0));
 
         for _ in 0..STEAL_COUNT {
             if let Some(task) = self.pool.queue.pop() {