Make TaskQueue support tasks that return values

[zturner]

Summary

This depends on D48240. After this patch, you will be able to do somethign like

shared_future<int> F = TQ.async([] { return 42; });
shared_future<Error> F2 = TQ.async([] { return someFunctionThatMightFail(); });

[labath]

llvm/include/llvm/Support/TaskQueue.h

115–116

This depends a lot on the results of the discussion with Chandler on the other patch, but I was thinking about this possibly long recursion chain caused by these ParentTask links.

I think think we could avoid the recursion if instead of ParentTask->wait(), we did something like TaskQueue->wait(ThisTask.get()). The TaskQueue function would do something like

while (Tasks.front().get() != TaskToWaitFor) Tasks.front()->execute(); // Details about who pops the task TBD
TaskToWaitFor->execute();

The existing wait function would be a special case of this one for TaskToWaitFor = Tasks.back().get(). I think this would also allow you to avoid having a shared_ptrs in the task queue (just plain unique_ptr should suffice).

llvm/include/llvm/Support/TaskQueue.h

 #define LLVM_SUPPORT_TASK_QUEUE_H
 
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/thread.h"
 
 #include <future>
 
 #include <atomic>
+#include <cassert>
 #include <condition_variable>
 #include <functional>
 #include <memory>
 #include <mutex>
 #include <queue>
 #include <utility>
 
 namespace llvm {
 /// TaskQueue executes work on exactly 1 thread in a first in first out
 /// fashion.  All work is guaranteed to execute in the order it is received.
 class TaskQueue {
 public:
   using TaskFunctionTy = std::function<void()>;
 
 #if LLVM_ENABLE_THREADS
-  using TaskTy = std::packaged_task<void()>;
+  struct TaskBase {
+    virtual ~TaskBase() {}
+    virtual void execute() = 0;
+  };
+  template <typename T> struct TaskImpl : public TaskBase {
+    explicit TaskImpl(std::packaged_task<T()> Task)
+        : PackagedTask(std::move(Task)) {}
+
+    void execute() override { PackagedTask(); }
+    std::packaged_task<T()> PackagedTask;
+  };
+  using QueueItemTy = std::unique_ptr<TaskBase>;
 #else
-  using TaskTy = std::shared_future<void>;
+  struct TaskBase {
+    virtual ~TaskBase() {}
+    virtual void wait() = 0;
+  };
+  template <typename T> struct TaskImpl : public TaskBase {
+    TaskImpl(std::shared_future<T> Future) : Future(Future) {}
+
+    void wait() override { Future.wait(); }
+
+    std::shared_future<T> Future;
+  };
+  using QueueItemTy = std::shared_ptr<TaskBase>;
 #endif
 
   /// Construct a task queue with no work.
   TaskQueue();
 
   /// Blocking destructor: the queue will wait for all work to complete.
   ~TaskQueue();
 
   /// Asynchronous submission of a task to the queue. The returned future can be
   /// used to wait for the task (and all previous tasks that have not yet
   /// completed) to finish.
   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 queue. The returned future can be
   /// used to wait for the task (and all previous tasks that have not yet
   /// completed) to finish.
   template <typename Function>
-  inline std::shared_future<void> async(Function &&F) {
-    return asyncImpl(std::forward<Function>(F));
+  inline std::shared_future<typename std::result_of<Function()>::type>
+  async(Function &&F) {
+    using ResultTy = typename std::result_of<Function()>::type;
+
+#if LLVM_ENABLE_THREADS
+    /// Wrap the Task in a packaged_task to return a future object.
+    std::packaged_task<ResultTy()> PackagedTask(std::move(F));
+    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 queue
+      assert(EnableFlag && "Queuing a thread during TaskQueue destruction");
+
+      Tasks.push(
+          llvm::make_unique<TaskImpl<ResultTy>>(std::move(PackagedTask)));
+    }
+    QueueCondition.notify_one();
+    return Future.share();
+#else
+    std::shared_ptr<TaskBase> ParentTask;
+    if (!Tasks.empty())
+      ParentTask = Tasks.back();
+
+    // Create a wrapper for this task which first waits on the parent task (if
+    // there is one) and then executes this task.
+    auto WrappedTask = [this, ParentTask, F] {
+      if (ParentTask)
+        ParentTask->wait();
+
+      // If we're executing "this" task, then we must have necessarily already
+      // executed all tasks that come before.  That means this task is the first
+      // task in the list, and we can pop it to effectively "drain" the queue.
+      // That way someone could queue up F1, F2, F3, wait out of order on F2,
+      // then call TaskQueue::wait() and the front of the queue would correctly
+      // contain the task for F3.
+      Tasks.pop();
+      return F();
+    };
+
+    // The only way to have a future which invokes the function when you call
+    // F.wait() is by calling std::async with launch::deferred.  So we need this
+    // to support the case where the user waits on the returned future.
+    auto Future =
+        std::async(std::launch::deferred, std::move(WrappedTask)).share();
+
+    auto ThisTask = std::make_shared<TaskImpl<ResultTy>>(Future);
+    Tasks.push(ThisTask);
+    return Future;
+#endif
   }
 
   /// Blocking wait for all work to complete.
   void wait();
 
 private:
   void WorkThreadFunc();
 
-  /// Asynchronous submission of a task to the queue. The returned future can be
-  /// used to wait for the task (and all previous tasks that have not yet
-  /// completed) to finish.
-  std::shared_future<void> asyncImpl(TaskFunctionTy F);
-
   /// Tasks waiting for execution in the queue.
-  std::queue<TaskTy> Tasks;
+  std::queue<QueueItemTy> Tasks;
 
 #if LLVM_ENABLE_THREADS
   /// Signal for the destruction of the pool, asking thread to exit.
   /// This needs to come before WorkThread so that it is initialized first in
   /// the constructor.
   bool EnableFlag = true;
 
   /// Flag indicating whether the queue is in the middle of processing a job.

llvm/lib/Support/TaskQueue.cpp

 using namespace llvm;
 
 #if LLVM_ENABLE_THREADS
 
 TaskQueue::TaskQueue() : WorkThread([this] { WorkThreadFunc(); }) {}
 
 void TaskQueue::WorkThreadFunc() {
   while (true) {
-    TaskTy TheTask;
+    std::unique_ptr<TaskBase> TheTask;
     {
       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;
 
       {
         std::unique_lock<std::mutex> LockGuard(CompletionLock);
         Active = true;
       }
       // Yeah, we have a task, grab it and release the lock on the queue
       TheTask = std::move(Tasks.front());
       Tasks.pop();
     }
 
-    // Run the task we just grabbed and update the ids.
-    TheTask();
+    // Run the task we just grabbed which should make all waiting futures become
+    // ready.
+    TheTask->execute();
 
     {
       std::unique_lock<std::mutex> LockGuard(CompletionLock);
       Active = false;
     }
 
     // Notify task completion, in case someone waits on ThreadPool::wait()
     CompletionCondition.notify_all();
   }
 }
 
 void TaskQueue::wait() {
   // Wait for 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 !Active && Tasks.empty(); });
 }
 
-std::shared_future<void> TaskQueue::asyncImpl(TaskFunctionTy TheTask) {
-  /// Wrap the Task in a packaged_task to return a future object.
-  TaskTy PackagedTask(std::move(TheTask));
-  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 the work thread, waiting for completion.
 TaskQueue::~TaskQueue() {
   {
     std::unique_lock<std::mutex> LockGuard(QueueLock);
     EnableFlag = false;
   }
   QueueCondition.notify_one();
   WorkThread.join();
 }
 
 #else // LLVM_ENABLE_THREADS Disabled
 
 TaskQueue::TaskQueue() {}
 
 void TaskQueue::wait() {
   // Sequential implementation running the tasks
   while (!Tasks.empty()) {
-    auto Task = Tasks.front();
-    Task.wait();
-  }
-}
-
-std::shared_future<void> TaskQueue::asyncImpl(TaskFunctionTy Task) {
-  std::shared_future<void> ParentFuture;
-  if (!Tasks.empty())
-    ParentFuture = Tasks.back();
-
-  auto Future = std::async(std::launch::deferred, [this, ParentFuture, Task] {
-                  if (ParentFuture.valid())
-                    ParentFuture.wait();
+    auto Task = std::move(Tasks.front());
 
-                  Task();
-                  Tasks.pop();
-                }).share();
-
-  Tasks.push(Future);
-  return Future;
+    // This will pop from the queue as part of the internal wrapped task.
+    Task->wait();
+  }
 }
 
 TaskQueue::~TaskQueue() { wait(); }
 
 #endif

llvm/unittests/Support/TaskQueueTest.cpp

 //========- unittests/Support/TaskQueue.cpp - TaskQueue.h tests ------========//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/TaskQueue.h"
 
+#include "llvm/Support/Error.h"
+
 #include "gtest/gtest.h"
 
 using namespace llvm;
 
 class TaskQueueTest : public testing::Test {
 protected:
   TaskQueueTest() {}
 

 #endif
 
   TQ.wait();
   ASSERT_EQ(1, X);
   ASSERT_EQ(1, Y);
   ASSERT_EQ(1, Z);
   ASSERT_EQ(1, A);
 }
+
+TEST_F(TaskQueueTest, TasksWithReturnValues) {
+  TaskQueue TQ;
+  int I = 0;
+  std::string S;
+  char C = 0;
+  std::shared_future<int> F1 = TQ.async([&]() { return I = 42; });
+  std::shared_future<std::string> F2 = TQ.async([&]() { return S = "Test"; });
+  std::shared_future<char> F3 = TQ.async([&]() { return C = '\n'; });
+
+  F3.wait();
+
+  // First check all the local variables, to make sure that an out-of-order call
+  // to get() also retrieved all other values.
+  EXPECT_EQ('\n', C);
+  EXPECT_EQ("Test", S);
+  EXPECT_EQ(42, I);
+
+  // Then check all the returned values.
+  EXPECT_EQ('\n', F3.get());
+  EXPECT_EQ("Test", F2.get());
+  EXPECT_EQ(42, F1.get());
+}