add wait_for_all to the ProactiveThreadpool

gguo4 · gguo4 · commit 2b4b40471df8 · 2018-09-03T17:28:20.000-05:00
diff --git a/taskflow/threadpool/proactive_threadpool.hpp b/taskflow/threadpool/proactive_threadpool.hpp
@@ -95,6 +95,11 @@ class ProactiveThreadpool {
             if(_task_queue.empty()){
               w.ready = false;
               _workers.push_back(&w);
+
+              if(_workers.size() == num_workers()){
+                _complete.notify_one();
+              }
+
               w.cv.wait(lock, [&w](){ return w.ready; });              
               t = std::move(w.task);
             }
@@ -215,6 +220,18 @@ class ProactiveThreadpool {
   
   }
 
+  void wait_for_all(){
+
+    if(is_worker()){
+      throw std::runtime_error("Worker thread cannot wait for all");
+    }
+
+    std::unique_lock<std::mutex> lock(_mutex);
+    _complete.wait(lock, [this](){ return _workers.size() == num_workers(); }); 
+
+  }
+
+
   private:
     
     template <typename T>
@@ -238,6 +255,8 @@ class ProactiveThreadpool {
     mutable std::mutex _mutex;
 
     std::condition_variable _empty;
+    std::condition_variable _complete;
+
     std::deque<UnitTask> _task_queue;
     std::vector<std::thread> _threads;
     std::vector<Worker*> _workers; 
diff --git a/unittest/threadpool.cpp b/unittest/threadpool.cpp
@@ -54,7 +54,38 @@ void test_threadpool_async(ThreadpoolType& tp, const size_t task_num){
   for(size_t i=0; i<int_future.size(); i++){
     CHECK(int_future[i].get() == int_result[i]);
   } 
- 
+}
+
+template <typename ThreadpoolType>
+void test_threadpool_wait_for_all(ThreadpoolType& tp){
+
+  const size_t worker_num = tp.num_workers();
+  const size_t task_num = 20;
+  std::atomic<size_t> counter{0};
+  
+  for(size_t i=0; i<task_num; i++){
+    tp.silent_async([&counter](){ 
+      std::this_thread::sleep_for(std::chrono::milliseconds(200));
+      counter++; 
+    });
+  }
+  CHECK(counter < task_num);
+  //std::cout << "counter: " << counter << std::endl;
+
+  tp.shutdown();
+  tp.spawn(worker_num);
+
+  counter = 0;
+  for(size_t i=0; i<task_num; i++){
+    tp.silent_async([&counter](){ 
+      std::this_thread::sleep_for(std::chrono::milliseconds(200));
+      counter++; 
+    });
+  }
+  tp.wait_for_all();
+  //std::cout << "counter:" << counter << std::endl;
+  CHECK(counter == task_num);
+
 }
 
 // --------------------------------------------------------
@@ -76,5 +107,10 @@ TEST_CASE("Threadpool.ProactiveThreadpool" * doctest::timeout(5)) {
     test_threadpool_silent_async(tp, task_num);
   }
 
+  SUBCASE("WaitForAll"){
+    tf::ProactiveThreadpool tp(4);
+    test_threadpool_wait_for_all(tp);
+  }
+
 }
 
