diff --git a/internal/platform/implementation/windows/thread_pool.cc b/internal/platform/implementation/windows/thread_pool.cc index 8a50d397..dfbcb644 100644 --- a/internal/platform/implementation/windows/thread_pool.cc +++ b/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::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(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 diff --git a/internal/platform/implementation/windows/thread_pool.h b/internal/platform/implementation/windows/thread_pool.h index ffeb166e..9e9207ea 100644 --- a/internal/platform/implementation/windows/thread_pool.h +++ b/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 shutdown_latch_ = nullptr; + friend VOID CALLBACK WorkCallback(PTP_CALLBACK_INSTANCE instance, PVOID parameter, PTP_WORK work); }; diff --git a/internal/platform/implementation/windows/thread_pool_test.cc b/internal/platform/implementation/windows/thread_pool_test.cc index 9045f49d..11038eae 100644 --- a/internal/platform/implementation/windows/thread_pool_test.cc +++ b/internal/platform/implementation/windows/thread_pool_test.cc @@ -14,11 +14,11 @@ #include "internal/platform/implementation/windows/thread_pool.h" +#include #include #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