thread_pool.h 2.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. // Copyright 2017 The Abseil Authors.
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // https://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. #ifndef ABSL_SYNCHRONIZATION_INTERNAL_THREAD_POOL_H_
  15. #define ABSL_SYNCHRONIZATION_INTERNAL_THREAD_POOL_H_
  16. #include <cassert>
  17. #include <cstddef>
  18. #include <functional>
  19. #include <queue>
  20. #include <thread> // NOLINT(build/c++11)
  21. #include <vector>
  22. #include "absl/base/thread_annotations.h"
  23. #include "absl/synchronization/mutex.h"
  24. namespace absl {
  25. ABSL_NAMESPACE_BEGIN
  26. namespace synchronization_internal {
  27. // A simple ThreadPool implementation for tests.
  28. class ThreadPool {
  29. public:
  30. explicit ThreadPool(int num_threads) {
  31. for (int i = 0; i < num_threads; ++i) {
  32. threads_.push_back(std::thread(&ThreadPool::WorkLoop, this));
  33. }
  34. }
  35. ThreadPool(const ThreadPool &) = delete;
  36. ThreadPool &operator=(const ThreadPool &) = delete;
  37. ~ThreadPool() {
  38. {
  39. absl::MutexLock l(&mu_);
  40. for (size_t i = 0; i < threads_.size(); i++) {
  41. queue_.push(nullptr); // Shutdown signal.
  42. }
  43. }
  44. for (auto &t : threads_) {
  45. t.join();
  46. }
  47. }
  48. // Schedule a function to be run on a ThreadPool thread immediately.
  49. void Schedule(std::function<void()> func) {
  50. assert(func != nullptr);
  51. absl::MutexLock l(&mu_);
  52. queue_.push(std::move(func));
  53. }
  54. private:
  55. bool WorkAvailable() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
  56. return !queue_.empty();
  57. }
  58. void WorkLoop() {
  59. while (true) {
  60. std::function<void()> func;
  61. {
  62. absl::MutexLock l(&mu_);
  63. mu_.Await(absl::Condition(this, &ThreadPool::WorkAvailable));
  64. func = std::move(queue_.front());
  65. queue_.pop();
  66. }
  67. if (func == nullptr) { // Shutdown signal.
  68. break;
  69. }
  70. func();
  71. }
  72. }
  73. absl::Mutex mu_;
  74. std::queue<std::function<void()>> queue_ ABSL_GUARDED_BY(mu_);
  75. std::vector<std::thread> threads_;
  76. };
  77. } // namespace synchronization_internal
  78. ABSL_NAMESPACE_END
  79. } // namespace absl
  80. #endif // ABSL_SYNCHRONIZATION_INTERNAL_THREAD_POOL_H_