From 76f0d5873b921011296549c3a73b5cac05a72367 Mon Sep 17 00:00:00 2001
From: antonio2368 <antonio2368@users.noreply.github.com>
Date: Thu, 19 Nov 2020 16:46:01 +0100
Subject: [PATCH] Use unfinished tasks num instead of idle thread num (#45)

---
 src/utils/thread_pool.cpp        | 8 +++++---
 src/utils/thread_pool.hpp        | 4 ++--
 tests/unit/utils_thread_pool.cpp | 6 +++---
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/utils/thread_pool.cpp b/src/utils/thread_pool.cpp
index 3c3177556..4155c663e 100644
--- a/src/utils/thread_pool.cpp
+++ b/src/utils/thread_pool.cpp
@@ -11,6 +11,7 @@ ThreadPool::ThreadPool(const size_t pool_size) {
 void ThreadPool::AddTask(std::function<void()> new_task) {
   task_queue_.WithLock([&](auto &queue) {
     queue.emplace(std::make_unique<TaskSignature>(std::move(new_task)));
+    unfinished_tasks_num_.fetch_add(1);
   });
   queue_cv_.notify_one();
 }
@@ -55,22 +56,23 @@ void ThreadPool::ThreadLoop() {
         return;
       }
       (*task)();
+      unfinished_tasks_num_.fetch_sub(1);
       task = PopTask();
     }
 
     std::unique_lock guard(pool_lock_);
-    idle_thread_num_.fetch_add(1);
     queue_cv_.wait(guard, [&] {
       task = PopTask();
       return task || terminate_pool_.load();
     });
-    idle_thread_num_.fetch_sub(1);
     if (terminate_pool_.load()) {
       return;
     }
   }
 }
 
-size_t ThreadPool::IdleThreadNum() const { return idle_thread_num_.load(); }
+size_t ThreadPool::UnfinishedTasksNum() const {
+  return unfinished_tasks_num_.load();
+}
 
 }  // namespace utils
diff --git a/src/utils/thread_pool.hpp b/src/utils/thread_pool.hpp
index 408f55615..bf9617761 100644
--- a/src/utils/thread_pool.hpp
+++ b/src/utils/thread_pool.hpp
@@ -29,7 +29,7 @@ class ThreadPool {
   ThreadPool &operator=(const ThreadPool &) = delete;
   ThreadPool &operator=(ThreadPool &&) = delete;
 
-  size_t IdleThreadNum() const;
+  size_t UnfinishedTasksNum() const;
 
  private:
   std::unique_ptr<TaskSignature> PopTask();
@@ -38,7 +38,7 @@ class ThreadPool {
 
   std::vector<std::thread> thread_pool_;
 
-  std::atomic<size_t> idle_thread_num_{0};
+  std::atomic<size_t> unfinished_tasks_num_{0};
   std::atomic<bool> terminate_pool_{false};
   std::atomic<bool> stopped_{false};
   utils::Synchronized<std::queue<std::unique_ptr<TaskSignature>>,
diff --git a/tests/unit/utils_thread_pool.cpp b/tests/unit/utils_thread_pool.cpp
index f24dbc7ed..76bf81c37 100644
--- a/tests/unit/utils_thread_pool.cpp
+++ b/tests/unit/utils_thread_pool.cpp
@@ -9,8 +9,8 @@
 using namespace std::chrono_literals;
 
 TEST(ThreadPool, Basic) {
-  constexpr size_t adder_count = 100'000;
-  constexpr std::array<size_t, 5> pool_sizes{1, 2, 4, 8, 100};
+  constexpr size_t adder_count =
+      500'000; constexpr std::array<size_t, 5> pool_sizes{1, 2, 4, 8, 100};
 
   for (const auto pool_size : pool_sizes) {
     utils::ThreadPool pool{pool_size};
@@ -20,7 +20,7 @@ TEST(ThreadPool, Basic) {
       pool.AddTask([&] { count.fetch_add(1); });
     }
 
-    while (pool.IdleThreadNum() != pool_size) {
+    while (pool.UnfinishedTasksNum() != 0) {
       std::this_thread::sleep_for(std::chrono::milliseconds(10));
     }