bvh/v2/thread_pool.h

0001 #ifndef BVH_V2_THREAD_POOL_H
0002 #define BVH_V2_THREAD_POOL_H
0003
0004 #include <thread>
0005 #include <mutex>
0006 #include <condition_variable>
0007 #include <vector>
0008 #include <queue>
0009 #include <functional>
0010
0011 namespace bvh::v2 {
0012
0013 class ThreadPool {
0014 public:
0015     using Task = std::function<void(size_t)>;
0016
0017     /// Creates a thread pool with the given number of threads (a value of 0 tries to autodetect
0018     /// the number of threads and uses that as a thread count).
0019     ThreadPool(size_t thread_count = 0) { start(thread_count); }
0020
0021     ~ThreadPool() {
0022         wait();
0023         stop();
0024         join();
0025     }
0026
0027     inline void push(Task&& fun);
0028     inline void wait();
0029
0030     size_t get_thread_count() const { return threads_.size(); }
0031
0032 private:
0033     static inline void worker(ThreadPool*, size_t);
0034
0035     inline void start(size_t);
0036     inline void stop();
0037     inline void join();
0038
0039     int busy_count_ = 0;
0040     bool should_stop_ = false;
0041     std::mutex mutex_;
0042     std::vector<std::thread> threads_;
0043     std::condition_variable avail_;
0044     std::condition_variable done_;
0045     std::queue<Task> tasks_;
0046 };
0047
0048 void ThreadPool::push(Task&& task) {
0049     {
0050         std::unique_lock<std::mutex> lock(mutex_);
0051         tasks_.emplace(std::move(task));
0052     }
0053     avail_.notify_one();
0054 }
0055
0056 void ThreadPool::wait() {
0057     std::unique_lock<std::mutex> lock(mutex_);
0058     done_.wait(lock, [this] { return busy_count_ == 0 && tasks_.empty(); });
0059 }
0060
0061 void ThreadPool::worker(ThreadPool* pool, size_t thread_id) {
0062     while (true) {
0063         Task task;
0064         {
0065             std::unique_lock<std::mutex> lock(pool->mutex_);
0066             pool->avail_.wait(lock, [pool] { return pool->should_stop_ || !pool->tasks_.empty(); });
0067             if (pool->should_stop_ && pool->tasks_.empty())
0068                 break;
0069             task = std::move(pool->tasks_.front());
0070             pool->tasks_.pop();
0071             pool->busy_count_++;
0072         }
0073         task(thread_id);
0074         {
0075             std::unique_lock<std::mutex> lock(pool->mutex_);
0076             pool->busy_count_--;
0077         }
0078         pool->done_.notify_one();
0079     }
0080 }
0081
0082 void ThreadPool::start(size_t thread_count) {
0083     if (thread_count == 0)
0084         thread_count = std::max(1u, std::thread::hardware_concurrency());
0085     for (size_t i = 0; i < thread_count; ++i)
0086         threads_.emplace_back(worker, this, i);
0087 }
0088
0089 void ThreadPool::stop() {
0090     {
0091         std::unique_lock<std::mutex> lock(mutex_);
0092         should_stop_ = true;
0093     }
0094     avail_.notify_all();
0095 }
0096
0097 void ThreadPool::join() {
0098     for (auto& thread : threads_)
0099         thread.join();
0100 }
0101
0102 } // namespace bvh::v2
0103
0104 #endif