Rewrite ThreadPool.

internal/platform/implementation/linux/thread_pool.cc

@@ -12,117 +12,106 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "internal/platform/implementation/linux/thread_pool.h"
-
+#include <atomic>
 #include <queue>
 #include <utility>
 
-#include "absl/memory/memory.h"
 #include "absl/synchronization/mutex.h"
+#include "internal/platform/implementation/linux/thread_pool.h"
 #include "internal/platform/logging.h"
 #include "internal/platform/runnable.h"
 
 namespace nearby {
 namespace linux {
-
-std::unique_ptr<ThreadPool> ThreadPool::Create(int max_pool_size) {
-  NEARBY_LOGS(VERBOSE) << __func__ << ": Create thread pool with maximum size("
-                       << max_pool_size << ").";
-
-  if (max_pool_size <= 0) {
-    NEARBY_LOGS(ERROR) << __func__
-                       << ": Maximum pool size must be positive integer value.";
-    return nullptr;
-  }
-
-  std::unique_ptr<std::vector<std::thread>> thread_pool = std::make_unique<std::vector<std::thread>>();
-
-  // Sets thread pool maximum value.
-  thread_pool->resize(max_pool_size);
-
-  return absl::WrapUnique(
-      new ThreadPool(thread_pool, max_pool_size));
+ThreadPool::ThreadPool(size_t max_pool_size)
+    : max_pool_size_(max_pool_size), shut_down_(false) {
+  threads_.reserve(max_pool_size);
 }
 
-ThreadPool::ThreadPool(std::unique_ptr<std::vector<std::thread>> &thread_pool, int max_pool_size)
-    : thread_pool_(std::move(thread_pool)),
-      max_pool_size_(max_pool_size) {
-  NEARBY_LOGS(VERBOSE) << __func__ << ": Thread pool(" << this
-                       << ") is created with size:" << max_pool_size_;
+ThreadPool::~ThreadPool() { ShutDown(); }
 
-  for (int size = 0; size < max_pool_size_; size++) {
-    thread_pool_->at(size) = std::thread([this]() {
-      while (!tasks_.empty()) {
-        RunNextTask();
-        // Possibly don't need but here to prevent 100% usage for loop
-        sleep(300);
+bool ThreadPool::Start() {
+  shut_down_.store(false, std::memory_order_acquire);
+
+  absl::MutexLock l(&mutex_);
+  if (!threads_.empty()) {
+    NEARBY_LOGS(ERROR) << __func__ << "thread pool is already active";
+    return false;
+  }
+
+  auto runner = [&]() {
+    while (true) {
+      if (shut_down_) {
+        return;
       }
-    });
+
+      auto task = NextTask();
+
+      if (task == nullptr) {
+        NEARBY_LOGS(WARNING) << __func__ << ": Tried to run a null task.";
+        continue;
+      }
+      task();
+    }
+  };
+
+  for (size_t i = 0; i < max_pool_size_; i++) {
+    threads_.emplace_back(runner);
   }
+
+  return true;
 }
 
-ThreadPool::~ThreadPool() {
-  NEARBY_LOGS(VERBOSE) << __func__ << ": Thread pool(" << this
-                       << ") is released.";
+bool ThreadPool::Run(Runnable &&task) {
+  if (shut_down_) {
+    NEARBY_LOGS(ERROR) << __func__ << "thread pool has shut down";
+    return false;
+  }
 
-  ShutDown();
-}
-
-bool ThreadPool::Run(Runnable task) {
-  absl::MutexLock lock(&mutex_);
-
-  if (thread_pool_->size() == max_pool_size_) {
+  absl::MutexLock l(&mutex_);
+  if (threads_.empty()) {
+    NEARBY_LOGS(ERROR) << __func__ << "thread pool is not active";
     return false;
   }
 
   tasks_.push(std::move(task));
-  NEARBY_LOGS(VERBOSE) << __func__ << ": Scheduled to run task("
-                       << &tasks_.back() << ").";
   return true;
 }
 
 void ThreadPool::ShutDown() {
-  absl::MutexLock lock(&mutex_);
+  shut_down_.store(true, std::memory_order_acquire);
 
-  if (!thread_pool_->empty()) {
-    NEARBY_LOGS(WARNING) << __func__ << ": Request to shutdown thread pool with " << thread_pool_->size() << " tasks not finished(" << this << ").";
-  }
-
-  NEARBY_LOGS(VERBOSE) << __func__ << ": Shutdown thread pool(" << this << ").";
-  if (thread_pool_ == nullptr) {
-    NEARBY_LOGS(WARNING) << __func__ << ": Shutdown on closed thread pool("
-                         << this << ").";
-    return;
-  }
-
-  thread_pool_.reset();
-}
-
-void ThreadPool::RunNextTask() {
-  Runnable task = nullptr;
+  NEARBY_LOGS(INFO)
+      << __func__ << ": asked to shut down, waiting for active threads to stop";
 
   {
-    absl::MutexLock lock(&mutex_);
-
-    if (!thread_pool_) {
-      return;
-    }
-    if (!tasks_.empty()) {
-      NEARBY_LOGS(VERBOSE) << __func__ << ": Run task(" << &tasks_.front()
-                           << ").";
-
-      task = std::move(tasks_.front());
-      tasks_.pop();
+    absl::ReaderMutexLock l(&mutex_);
+    for (auto &thread : threads_) {
+      thread.join();
     }
   }
-  if (task == nullptr) {
-    NEARBY_LOGS(WARNING) << __func__
-                         << ": Tried to run task in an empty thread pool.";
-    return;
-  }
 
-  task();
+  absl::MutexLock l(&mutex_);
+  threads_.clear();
+  NEARBY_LOGS(INFO) << __func__ << ": shut down thread pool";
 }
 
-}  // namespace linux
-}  // namespace nearby
+Runnable ThreadPool::NextTask() {
+  Runnable task;
+  auto task_available = [&]() {
+    mutex_.AssertReaderHeld();
+    return !tasks_.empty();
+  };
+
+  {
+    absl::MutexLock l(&mutex_, absl::Condition(&task_available));
+
+    task = std::move(tasks_.front());
+    tasks_.pop();
+  }
+
+  return task;
+}
+
+} // namespace linux
+} // namespace nearby

internal/platform/implementation/linux/thread_pool.h

@@ -15,11 +15,12 @@
 #ifndef PLATFORM_IMPL_LINUX_THREAD_POOL_H_
 #define PLATFORM_IMPL_LINUX_THREAD_POOL_H_
 
+#include <atomic>
 #include <memory>
 #include <queue>
+#include <thread>
 #include <utility>
 #include <vector>
-#include <thread>
 
 #include "absl/base/thread_annotations.h"
 #include "absl/synchronization/mutex.h"
@@ -29,44 +30,30 @@ namespace nearby {
 namespace linux {
 
 class ThreadPool {
- public:
-  virtual ~ThreadPool();
-  static std::unique_ptr<ThreadPool> Create(int max_pool_size);
+public:
+  ThreadPool(size_t max_pool_size);
+  ~ThreadPool();
+
+  bool Start() ABSL_LOCKS_EXCLUDED(mutex_);
 
   // Runs a task on thread pool. The result indicates whether the task is put
   // into the thread pool.
-  bool Run(Runnable task) ABSL_LOCKS_EXCLUDED(mutex_);
+  bool Run(Runnable &&task) ABSL_LOCKS_EXCLUDED(mutex_);
 
-  // The thread pool is closed immediately if there is no outstanding work,
-  // I/O, timer, or wait objects that are bound to the pool; otherwise, the
-  // thread pool is released asynchronously after the outstanding objects are
-  // freed.
   void ShutDown() ABSL_LOCKS_EXCLUDED(mutex_);
 
- private:
-  ThreadPool(std::unique_ptr<std::vector<std::thread>> &thread_pool, int max_pool_size);
+private:
+  Runnable NextTask() ABSL_LOCKS_EXCLUDED(mutex_);
 
-  void RunNextTask();
+private:
+  size_t max_pool_size_;
+  std::atomic_bool shut_down_;
 
-  // Starts each task and injects a function that removes the task when it is finished
-  std::thread tasks_runner_;
-
-  // Protects the access to tasks of the thread pool.
-  mutable absl::Mutex mutex_;
-
-  // The task queue of the thread pool. Thread pool will pick up task to run
-  // when it is idle.
+  absl::Mutex mutex_;
+  std::vector<std::thread> threads_ ABSL_GUARDED_BY(mutex_);
   std::queue<Runnable> tasks_ ABSL_GUARDED_BY(mutex_);
-
-  // All the threads run in the pool
-  std::unique_ptr<std::vector<std::thread>> thread_pool_ ABSL_GUARDED_BY(mutex_);
-
-  // The maximum thread count in the thread pool
-  int max_pool_size_ ABSL_GUARDED_BY(mutex_) = 0;
-
 };
+} // namespace linux
+} // namespace nearby
 
-}  // namespace linux
-}  // namespace nearby
-
-#endif  //  PLATFORM_IMPL_LINUX_THREAD_POOL_H_
+#endif //  PLATFORM_IMPL_LINUX_THREAD_POOL_H_