Make ThreadPool support running tasks that return values

[zturner]

Summary

previously ThreadPool could only execute tasks that did something, not tasks that returned values. We can improve this with a little bit of type erasure. The idea is that instead of queueing up packaged_task<void()>'s, we queue up unique_ptr<TaskBase>, where TaskBase is a base class that has an pure virtual execute() method. A templated derived class stores an appropriately typed package_task<T()>, which is run with execute() is called.

[davide]

Looks fine to me. Do you plan to use It in a subsequent patch?

[labath]

looks good

llvm/include/llvm/Support/ThreadPool.h

30

you could include cassert so it fits in nicely with the other c++ headers.

55–56

Are these still needed/used ?

[zturner]

In D48115#1130685, @davide wrote:

Looks fine to me. Do you plan to use It in a subsequent patch?

Yes, maybe not immediately, but it's one more step along the way of things I find I'm needing.

llvm/include/llvm/Support/ThreadPool.h

 #include <atomic>
 #include <condition_variable>
 #include <functional>
 #include <memory>
 #include <mutex>
 #include <queue>
 #include <utility>
 
+#include <assert.h>
+
 namespace llvm {
 
 /// A ThreadPool for asynchronous parallel execution on a defined number of
 /// threads.
 ///
 /// The pool keeps a vector of threads alive, waiting on a condition variable
 /// for some work to become available.
 class ThreadPool {
+  struct TaskBase {
+    virtual ~TaskBase() {}
+    virtual void execute() = 0;
+  };
+
+  template <typename ReturnType> struct TypedTask : public TaskBase {
+    explicit TypedTask(std::packaged_task<ReturnType()> Task)
+        : Task(std::move(Task)) {}
+
+    void execute() override { Task(); }
+
+    std::packaged_task<ReturnType()> Task;
+  };
+
 public:
   using TaskTy = std::function<void()>;
   using PackagedTaskTy = std::packaged_task<void()>;
 
   /// Construct a pool with the number of threads found by
   /// hardware_concurrency().
   ThreadPool();
 
   /// Construct a pool of \p ThreadCount threads
   ThreadPool(unsigned ThreadCount);
 
   /// Blocking destructor: the pool will wait for all the threads to complete.
   ~ThreadPool();
 
   /// Asynchronous submission of a task to the pool. The returned future can be
   /// used to wait for the task to finish and is *non-blocking* on destruction.
   template <typename Function, typename... Args>
-  inline std::shared_future<void> async(Function &&F, Args &&... ArgList) {
+  inline std::shared_future<typename std::result_of<Function(Args...)>::type>
+  async(Function &&F, Args &&... ArgList) {
     auto Task =
         std::bind(std::forward<Function>(F), std::forward<Args>(ArgList)...);
     return asyncImpl(std::move(Task));
   }
 
   /// Asynchronous submission of a task to the pool. The returned future can be
   /// used to wait for the task to finish and is *non-blocking* on destruction.
   template <typename Function>
-  inline std::shared_future<void> async(Function &&F) {
+  inline std::shared_future<typename std::result_of<Function()>::type>
+  async(Function &&F) {
     return asyncImpl(std::forward<Function>(F));
   }
 
   /// Blocking wait for all the threads to complete and the queue to be empty.
   /// It is an error to try to add new tasks while blocking on this call.
   void wait();
 
 private:
   /// Asynchronous submission of a task to the pool. The returned future can be
   /// used to wait for the task to finish and is *non-blocking* on destruction.
-  std::shared_future<void> asyncImpl(TaskTy F);
+  template <typename TaskTy>
+  std::shared_future<typename std::result_of<TaskTy()>::type>
+  asyncImpl(TaskTy &&Task) {
+    typedef decltype(Task()) ResultTy;
+
+    /// Wrap the Task in a packaged_task to return a future object.
+    std::packaged_task<ResultTy()> PackagedTask(std::move(Task));
+    auto Future = PackagedTask.get_future();
+    std::unique_ptr<TaskBase> TB =
+        llvm::make_unique<TypedTask<ResultTy>>(std::move(PackagedTask));
+
+    {
+      // Lock the queue and push the new task
+      std::unique_lock<std::mutex> LockGuard(QueueLock);
+
+      // Don't allow enqueueing after disabling the pool
+      assert(EnableFlag && "Queuing a thread during ThreadPool destruction");
+
+      Tasks.push(std::move(TB));
+    }
+    QueueCondition.notify_one();
+    return Future.share();
+  }
 
   /// Threads in flight
   std::vector<llvm::thread> Threads;
 
   /// Tasks waiting for execution in the pool.
-  std::queue<PackagedTaskTy> Tasks;
+  std::queue<std::unique_ptr<TaskBase>> Tasks;
 
   /// Locking and signaling for accessing the Tasks queue.
   std::mutex QueueLock;
   std::condition_variable QueueCondition;
 
   /// Locking and signaling for job completion
   std::mutex CompletionLock;
   std::condition_variable CompletionCondition;

llvm/lib/Support/ThreadPool.cpp

 ThreadPool::ThreadPool(unsigned ThreadCount)
     : ActiveThreads(0), EnableFlag(true) {
   // Create ThreadCount threads that will loop forever, wait on QueueCondition
   // for tasks to be queued or the Pool to be destroyed.
   Threads.reserve(ThreadCount);
   for (unsigned ThreadID = 0; ThreadID < ThreadCount; ++ThreadID) {
     Threads.emplace_back([&] {
       while (true) {
-        PackagedTaskTy Task;
+        std::unique_ptr<TaskBase> Task;
         {
           std::unique_lock<std::mutex> LockGuard(QueueLock);
           // Wait for tasks to be pushed in the queue
           QueueCondition.wait(LockGuard,
                               [&] { return !EnableFlag || !Tasks.empty(); });
           // Exit condition
           if (!EnableFlag && Tasks.empty())
             return;
           // Yeah, we have a task, grab it and release the lock on the queue
 
           // We first need to signal that we are active before popping the queue
           // in order for wait() to properly detect that even if the queue is
           // empty, there is still a task in flight.
           {
             std::unique_lock<std::mutex> LockGuard(CompletionLock);
             ++ActiveThreads;
           }
           Task = std::move(Tasks.front());
           Tasks.pop();
         }
         // Run the task we just grabbed
-        Task();
+        Task->execute();
 
         {
           // Adjust `ActiveThreads`, in case someone waits on ThreadPool::wait()
           std::unique_lock<std::mutex> LockGuard(CompletionLock);
           --ActiveThreads;
         }
 
         // Notify task completion, in case someone waits on ThreadPool::wait()
         CompletionCondition.notify_all();
       }
     });
   }
 }
 
 void ThreadPool::wait() {
   // Wait for all threads to complete and the queue to be empty
   std::unique_lock<std::mutex> LockGuard(CompletionLock);
   // The order of the checks for ActiveThreads and Tasks.empty() matters because
   // any active threads might be modifying the Tasks queue, and this would be a
   // race.
   CompletionCondition.wait(LockGuard,
                            [&] { return !ActiveThreads && Tasks.empty(); });
 }
 
-std::shared_future<void> ThreadPool::asyncImpl(TaskTy Task) {
-  /// Wrap the Task in a packaged_task to return a future object.
-  PackagedTaskTy PackagedTask(std::move(Task));
-  auto Future = PackagedTask.get_future();
-  {
-    // Lock the queue and push the new task
-    std::unique_lock<std::mutex> LockGuard(QueueLock);
-
-    // Don't allow enqueueing after disabling the pool
-    assert(EnableFlag && "Queuing a thread during ThreadPool destruction");
-
-    Tasks.push(std::move(PackagedTask));
-  }
-  QueueCondition.notify_one();
-  return Future.share();
-}
-
 // The destructor joins all threads, waiting for completion.
 ThreadPool::~ThreadPool() {
   {
     std::unique_lock<std::mutex> LockGuard(QueueLock);
     EnableFlag = false;
   }
   QueueCondition.notify_all();
   for (auto &Worker : Threads)

llvm/unittests/Support/ThreadPool.cpp

TEST_F(ThreadPoolTest, GetFuture) {
   // Force the future using get()
   Pool.async([&i] { ++i; }).get();
   ASSERT_NE(2, i.load());
   setMainThreadReady();
   Pool.wait();
   ASSERT_EQ(2, i.load());
 }
 
+TEST_F(ThreadPoolTest, TaskWithResult) {
+  CHECK_UNSUPPORTED();
+  // By making only 1 thread in the pool the two tasks are serialized with
+  // respect to each other, which means that the second one must return 2.
+  ThreadPool Pool{1};
+  std::atomic_int i{0};
+  Pool.async([this, &i] {
+    waitForMainThread();
+    ++i;
+  });
+  // Force the future using get()
+  std::shared_future<int> Future = Pool.async([&i] { return ++i; });
+  ASSERT_EQ(0, i.load());
+  setMainThreadReady();
+  int Result = Future.get();
+  ASSERT_EQ(2, i.load());
+  ASSERT_EQ(2, Result);
+}
+
 TEST_F(ThreadPoolTest, PoolDestruction) {
   CHECK_UNSUPPORTED();
   // Test that we are waiting on destruction
   std::atomic_int checked_in{0};
   {
     ThreadPool Pool;
     for (size_t i = 0; i < 5; ++i) {
       Pool.async([this, &checked_in] {