synchronization/internal/thread_pool.h

0001 // Copyright 2017 The Abseil Authors.
0002 //
0003 // Licensed under the Apache License, Version 2.0 (the "License");
0004 // you may not use this file except in compliance with the License.
0005 // You may obtain a copy of the License at
0006 //
0007 //      https://www.apache.org/licenses/LICENSE-2.0
0008 //
0009 // Unless required by applicable law or agreed to in writing, software
0010 // distributed under the License is distributed on an "AS IS" BASIS,
0011 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0012 // See the License for the specific language governing permissions and
0013 // limitations under the License.
0014
0015 #ifndef ABSL_SYNCHRONIZATION_INTERNAL_THREAD_POOL_H_
0016 #define ABSL_SYNCHRONIZATION_INTERNAL_THREAD_POOL_H_
0017
0018 #include <cassert>
0019 #include <cstddef>
0020 #include <functional>
0021 #include <queue>
0022 #include <thread>  // NOLINT(build/c++11)
0023 #include <utility>
0024 #include <vector>
0025
0026 #include "absl/base/thread_annotations.h"
0027 #include "absl/functional/any_invocable.h"
0028 #include "absl/synchronization/mutex.h"
0029
0030 namespace absl {
0031 ABSL_NAMESPACE_BEGIN
0032 namespace synchronization_internal {
0033
0034 // A simple ThreadPool implementation for tests.
0035 class ThreadPool {
0036  public:
0037   explicit ThreadPool(int num_threads) {
0038     threads_.reserve(num_threads);
0039     for (int i = 0; i < num_threads; ++i) {
0040       threads_.push_back(std::thread(&ThreadPool::WorkLoop, this));
0041     }
0042   }
0043
0044   ThreadPool(const ThreadPool &) = delete;
0045   ThreadPool &operator=(const ThreadPool &) = delete;
0046
0047   ~ThreadPool() {
0048     {
0049       absl::MutexLock l(&mu_);
0050       for (size_t i = 0; i < threads_.size(); i++) {
0051         queue_.push(nullptr);  // Shutdown signal.
0052       }
0053     }
0054     for (auto &t : threads_) {
0055       t.join();
0056     }
0057   }
0058
0059   // Schedule a function to be run on a ThreadPool thread immediately.
0060   void Schedule(absl::AnyInvocable<void()> func) {
0061     assert(func != nullptr);
0062     absl::MutexLock l(&mu_);
0063     queue_.push(std::move(func));
0064   }
0065
0066  private:
0067   bool WorkAvailable() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
0068     return !queue_.empty();
0069   }
0070
0071   void WorkLoop() {
0072     while (true) {
0073       absl::AnyInvocable<void()> func;
0074       {
0075         absl::MutexLock l(&mu_);
0076         mu_.Await(absl::Condition(this, &ThreadPool::WorkAvailable));
0077         func = std::move(queue_.front());
0078         queue_.pop();
0079       }
0080       if (func == nullptr) {  // Shutdown signal.
0081         break;
0082       }
0083       func();
0084     }
0085   }
0086
0087   absl::Mutex mu_;
0088   std::queue<absl::AnyInvocable<void()>> queue_ ABSL_GUARDED_BY(mu_);
0089   std::vector<std::thread> threads_;
0090 };
0091
0092 }  // namespace synchronization_internal
0093 ABSL_NAMESPACE_END
0094 }  // namespace absl
0095
0096 #endif  // ABSL_SYNCHRONIZATION_INTERNAL_THREAD_POOL_H_