[ThreadPool] Rollback recent changes until I figure out the breakage.

llvm-svn: 288018

3 files changed, 36 insertions, 12 deletions

diff --git a/llvm/include/llvm/Support/ThreadPool.h b/llvm/include/llvm/Support/ThreadPool.h
index da1834dbea9..665cec2465b 100644
--- a/llvm/include/llvm/Support/ThreadPool.h
+++ b/llvm/include/llvm/Support/ThreadPool.h
@@ -72,27 +72,32 @@ public:
   /// Blocking destructor: the pool will wait for all the threads to complete.
   ~ThreadPool();
 
-  /// Asynchronous submission of a task to the pool.
+  /// 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 void async(Function &&F, Args &&... ArgList) {
+  inline std::shared_future<VoidTy> async(Function &&F, Args &&... ArgList) {
     auto Task =
         std::bind(std::forward<Function>(F), std::forward<Args>(ArgList)...);
 #ifndef _MSC_VER
-    asyncImpl(std::move(Task));
+    return asyncImpl(std::move(Task));
 #else
     // This lambda has to be marked mutable because MSVC 2013's std::bind call
     // operator isn't const qualified.
-    asyncImpl([Task](VoidTy) mutable { Task(); });
+    return asyncImpl([Task](VoidTy) mutable -> VoidTy {
+      Task();
+      return VoidTy();
+    });
 #endif
   }
 
-  /// Asynchronous submission of a task to the pool.
+  /// 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 void async(Function &&F) {
+  inline std::shared_future<VoidTy> async(Function &&F) {
 #ifndef _MSC_VER
-    asyncImpl(std::forward<Function>(F));
+    return asyncImpl(std::forward<Function>(F));
 #else
-    asyncImpl([F] (VoidTy) { F(); });
+    return asyncImpl([F] (VoidTy) -> VoidTy { F(); return VoidTy(); });
 #endif
   }
 
@@ -101,8 +106,9 @@ public:
   void wait();
 
 private:
-  /// Asynchronous submission of a task to the pool.
-  void asyncImpl(TaskTy F);
+  /// 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<VoidTy> asyncImpl(TaskTy F);
 
   /// Threads in flight
   std::vector<llvm::thread> Threads;
diff --git a/llvm/lib/Support/ThreadPool.cpp b/llvm/lib/Support/ThreadPool.cpp
index 022f32a1378..db03a4d6240 100644
--- a/llvm/lib/Support/ThreadPool.cpp
+++ b/llvm/lib/Support/ThreadPool.cpp
@@ -82,7 +82,7 @@ void ThreadPool::wait() {
                            [&] { return !ActiveThreads && Tasks.empty(); });
 }
 
-void ThreadPool::asyncImpl(TaskTy Task) {
+std::shared_future<ThreadPool::VoidTy> 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();
@@ -96,6 +96,7 @@ void ThreadPool::asyncImpl(TaskTy Task) {
     Tasks.push(std::move(PackagedTask));
   }
   QueueCondition.notify_one();
+  return Future.share();
 }
 
 // The destructor joins all threads, waiting for completion.
@@ -135,7 +136,7 @@ void ThreadPool::wait() {
   }
 }
 
-void ThreadPool::asyncImpl(TaskTy Task) {
+std::shared_future<ThreadPool::VoidTy> ThreadPool::asyncImpl(TaskTy Task) {
 #ifndef _MSC_VER
   // Get a Future with launch::deferred execution using std::async
   auto Future = std::async(std::launch::deferred, std::move(Task)).share();
@@ -147,6 +148,7 @@ void ThreadPool::asyncImpl(TaskTy Task) {
   PackagedTaskTy PackagedTask([Future](bool) -> bool { Future.get(); return false; });
 #endif
   Tasks.push(std::move(PackagedTask));
+  return Future;
 }
 
 ThreadPool::~ThreadPool() {
diff --git a/llvm/unittests/Support/ThreadPool.cpp b/llvm/unittests/Support/ThreadPool.cpp
index bb972f7daa9..8e03aacfb1e 100644
--- a/llvm/unittests/Support/ThreadPool.cpp
+++ b/llvm/unittests/Support/ThreadPool.cpp
@@ -131,6 +131,22 @@ TEST_F(ThreadPoolTest, Async) {
   ASSERT_EQ(2, i.load());
 }
 
+TEST_F(ThreadPoolTest, GetFuture) {
+  CHECK_UNSUPPORTED();
+  ThreadPool Pool{2};
+  std::atomic_int i{0};
+  Pool.async([this, &i] {
+    waitForMainThread();
+    ++i;
+  });
+  // 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, PoolDestruction) {
   CHECK_UNSUPPORTED();
   // Test that we are waiting on destruction