Wait for all tasks completed in thread pool

PiperOrigin-RevId: 542866176

internal/platform/implementation/windows/thread_pool.cc

@@ -21,6 +21,7 @@
 
 #include "absl/memory/memory.h"
 #include "absl/synchronization/mutex.h"
+#include "internal/platform/count_down_latch.h"
 #include "internal/platform/logging.h"
 #include "internal/platform/runnable.h"
 
@@ -38,9 +39,6 @@ VOID CALLBACK WorkCallback(PTP_CALLBACK_INSTANCE instance, PVOID parameter,
 }
 
 std::unique_ptr<ThreadPool> ThreadPool::Create(int max_pool_size) {
-  NEARBY_LOGS(VERBOSE) << __func__ << ": Create thread pool with maximum size("
-                       << max_pool_size << ").";
-
   PTP_POOL thread_pool = nullptr;
   TP_CALLBACK_ENVIRON thread_pool_environ;
   InitializeThreadpoolEnvironment(&thread_pool_environ);
@@ -93,6 +91,10 @@ ThreadPool::~ThreadPool() {
   NEARBY_LOGS(VERBOSE) << __func__ << ": Thread pool(" << this
                        << ") is released.";
 
+  if (thread_pool_ == nullptr) {
+    return;
+  }
+
   ShutDown();
 }
 
@@ -103,6 +105,11 @@ bool ThreadPool::Run(Runnable task) {
     return false;
   }
 
+  if (shutdown_latch_ != nullptr) {
+    NEARBY_LOGS(WARNING) << __func__ << ": Thread pool is in shutting down.";
+    return false;
+  }
+
   PTP_WORK work;
   tasks_.push(std::move(task));
   NEARBY_LOGS(VERBOSE) << __func__ << ": Scheduled to run task("
@@ -116,6 +123,8 @@ bool ThreadPool::Run(Runnable task) {
     return false;
   }
 
+  ++running_tasks_count_;
+
   //
   // Submit the work to the pool. Because this was a pre-allocated
   // work item (using CreateThreadpoolWork), it is guaranteed to execute.
@@ -125,17 +134,45 @@ bool ThreadPool::Run(Runnable task) {
 }
 
 void ThreadPool::ShutDown() {
-  absl::MutexLock lock(&mutex_);
+  {
+    absl::MutexLock lock(&mutex_);
 
-  NEARBY_LOGS(VERBOSE) << __func__ << ": Shutdown thread pool(" << this << ").";
-  if (thread_pool_ == nullptr) {
-    NEARBY_LOGS(WARNING) << __func__ << ": Shutdown on closed thread pool("
-                         << this << ").";
-    return;
+    if (thread_pool_ == nullptr) {
+      NEARBY_LOGS(WARNING) << __func__ << ": Shutdown on closed thread pool("
+                           << this << ").";
+      return;
+    }
+
+    if (running_tasks_count_ == 0) {
+      CloseThreadpool(thread_pool_);
+      thread_pool_ = nullptr;
+      NEARBY_LOGS(VERBOSE) << __func__ << ": Thread pool(" << this
+                           << ") is shut down.";
+      return;
+    }
+
+    if (shutdown_latch_ != nullptr) {
+      NEARBY_LOGS(VERBOSE) << __func__ << ": Thread pool(" << this
+                           << ") is already in shutting down.";
+      return;
+    }
+
+    NEARBY_LOGS(VERBOSE) << __func__ << ": Thread pool(" << this
+                         << ") is shutting down.";
+
+    shutdown_latch_ = std::make_unique<CountDownLatch>(1);
   }
 
-  CloseThreadpool(thread_pool_);
-  thread_pool_ = nullptr;
+  // Wait for all tasks to complete.
+  shutdown_latch_->Await();
+
+  {
+    absl::MutexLock lock(&mutex_);
+    CloseThreadpool(thread_pool_);
+    thread_pool_ = nullptr;
+    NEARBY_LOGS(VERBOSE) << __func__ << ": Thread pool(" << this
+                         << ") is shut down.";
+  }
 }
 
 void ThreadPool::RunNextTask() {
@@ -153,15 +190,28 @@ void ThreadPool::RunNextTask() {
 
       task = std::move(tasks_.front());
       tasks_.pop();
+
+      if (task == nullptr) {
+        NEARBY_LOGS(WARNING)
+            << __func__ << ": Tried to run task in an empty thread pool.";
+        --running_tasks_count_;
+        if (running_tasks_count_ == 0 && shutdown_latch_ != nullptr) {
+          shutdown_latch_->CountDown();
+        }
+        return;
+      }
     }
   }
-  if (task == nullptr) {
-    NEARBY_LOGS(WARNING) << __func__
-                         << ": Tried to run task in an empty thread pool.";
-    return;
-  }
 
   task();
+
+  {
+    absl::MutexLock lock(&mutex_);
+    --running_tasks_count_;
+    if (running_tasks_count_ == 0 && shutdown_latch_ != nullptr) {
+      shutdown_latch_->CountDown();
+    }
+  }
 }
 
 }  // namespace windows

internal/platform/implementation/windows/thread_pool.h

@@ -23,6 +23,7 @@
 
 #include "absl/base/thread_annotations.h"
 #include "absl/synchronization/mutex.h"
+#include "internal/platform/count_down_latch.h"
 #include "internal/platform/runnable.h"
 
 namespace nearby {
@@ -37,11 +38,9 @@ class ThreadPool {
   // into the thread pool.
   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_);
+  // In Nearby platform, thread pool should make sure all queued tasks completed
+  // in shut down.
+  void ShutDown();
 
  private:
   ThreadPool(PTP_POOL thread_pool, TP_CALLBACK_ENVIRON thread_pool_environ,
@@ -65,6 +64,12 @@ class ThreadPool {
   // The maximum thread count in the thread pool
   int max_pool_size_ ABSL_GUARDED_BY(mutex_) = 0;
 
+  // Current running task count
+  int running_tasks_count_ ABSL_GUARDED_BY(mutex_) = 0;
+
+  // The latch is used to wait for running tasks
+  std::unique_ptr<CountDownLatch> shutdown_latch_ = nullptr;
+
   friend VOID CALLBACK WorkCallback(PTP_CALLBACK_INSTANCE instance,
                                     PVOID parameter, PTP_WORK work);
 };

internal/platform/implementation/windows/thread_pool_test.cc

@@ -14,11 +14,11 @@
 
 #include "internal/platform/implementation/windows/thread_pool.h"
 
+#include <atomic>
 #include <vector>
 
 #include "gtest/gtest.h"
 #include "absl/synchronization/blocking_counter.h"
-#include "absl/synchronization/notification.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 
@@ -66,6 +66,19 @@ TEST(ThreadPool, TasksInMultipleThreadsRunInParallel) {
   pool->ShutDown();
 }
 
+TEST(ThreadPool, ShutdownWaitsForRunningTasks) {
+  auto pool = ThreadPool::Create(1);
+  std::atomic_int value = 0;
+  pool->Run([&]() {
+    absl::SleepFor(absl::Seconds(1));
+    value += 1;
+  });
+
+  pool->ShutDown();
+
+  EXPECT_EQ(value, 1);
+}
+
 }  // namespace
 }  // namespace windows
 }  // namespace nearby