util/worker_pool.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file worker_pool.h | ||
| 3 | /// @brief A small FIFO thread pool for background jobs. | ||
| 4 | #pragma once | ||
| 5 | |||
| 6 | #include "ernie/util/profiler.h" | ||
| 7 | |||
| 8 | #include <algorithm> | ||
| 9 | #include <condition_variable> | ||
| 10 | #include <deque> | ||
| 11 | #include <functional> | ||
| 12 | #include <future> | ||
| 13 | #include <memory> | ||
| 14 | #include <mutex> | ||
| 15 | #include <string> | ||
| 16 | #include <thread> | ||
| 17 | #include <type_traits> | ||
| 18 | #include <utility> | ||
| 19 | #include <vector> | ||
| 20 | |||
| 21 | namespace ernie { | ||
| 22 | |||
| 23 | /// @brief A few threads running jobs in submission order. | ||
| 24 | /// On destruction, running jobs finish and queued ones are dropped | ||
| 25 | /// (their futures report `std::future_error`). | ||
| 26 | class WorkerPool { | ||
| 27 | public: | ||
| 28 | /// @brief Starts the threads. | ||
| 29 | /// @param threads The thread count; at least one is started. | ||
| 30 | /// @param name The profiler name prefix ("Worker 1", ...). | ||
| 31 | 5 | explicit WorkerPool(unsigned threads, const std::string& name = "Worker") { | |
| 32 | 16 | for (unsigned i = 0; i < std::max(threads, 1u); ++i) { | |
| 33 | 18 | threads_.emplace_back([this, thread = name + " " + std::to_string(i + 1)] { | |
| 34 | 9 | profiler::SetThreadName(thread); | |
| 35 | #if ERNIE_TRACY | ||
| 36 | tracy::SetThreadName(profiler::Intern(thread)); | ||
| 37 | #endif | ||
| 38 | 9 | Run(); | |
| 39 | 9 | }); | |
| 40 | } | ||
| 41 | 5 | } | |
| 42 | |||
| 43 | 5 | ~WorkerPool() { | |
| 44 | 5 | { | |
| 45 | 5 | const std::lock_guard lock(mutex_); | |
| 46 | 5 | stopping_ = true; | |
| 47 | 5 | jobs_.clear(); | |
| 48 | 5 | } | |
| 49 | 5 | wake_.notify_all(); | |
| 50 | 14 | for (std::thread& thread : threads_) { | |
| 51 | 9 | thread.join(); | |
| 52 | } | ||
| 53 | 5 | } | |
| 54 | |||
| 55 | WorkerPool(const WorkerPool&) = delete; | ||
| 56 | WorkerPool& operator=(const WorkerPool&) = delete; | ||
| 57 | |||
| 58 | /// @brief Queues a job. | ||
| 59 | /// @param job A callable taking no arguments. | ||
| 60 | /// @return A future for its result or exception. | ||
| 61 | 125 | template <class Job> std::future<std::invoke_result_t<Job>> Submit(Job job) { | |
| 62 | using Result = std::invoke_result_t<Job>; | ||
| 63 | 125 | auto task = std::make_shared<std::packaged_task<Result()>>(std::move(job)); | |
| 64 | 125 | std::future<Result> result = task->get_future(); | |
| 65 | { | ||
| 66 | 125 | const std::lock_guard lock(mutex_); | |
| 67 | 250 | jobs_.emplace_back([task] { (*task)(); }); | |
| 68 | ✗ | } | |
| 69 | 125 | wake_.notify_one(); | |
| 70 | 125 | return result; | |
| 71 | 125 | } | |
| 72 | |||
| 73 | private: | ||
| 74 | 9 | void Run() { | |
| 75 | 255 | for (;;) { | |
| 76 | 132 | std::function<void()> job; | |
| 77 | 132 | { | |
| 78 | 132 | std::unique_lock lock(mutex_); | |
| 79 | 273 | wake_.wait(lock, [this] { return stopping_ || !jobs_.empty(); }); | |
| 80 | 132 | if (stopping_) { | |
| 81 | 9 | return; | |
| 82 | } | ||
| 83 | 123 | job = std::move(jobs_.front()); | |
| 84 | 123 | jobs_.pop_front(); | |
| 85 | 9 | } | |
| 86 | 123 | job(); | |
| 87 | 132 | } | |
| 88 | } | ||
| 89 | |||
| 90 | std::mutex mutex_; | ||
| 91 | std::condition_variable wake_; | ||
| 92 | std::deque<std::function<void()>> jobs_; | ||
| 93 | bool stopping_ = false; | ||
| 94 | std::vector<std::thread> threads_; | ||
| 95 | }; | ||
| 96 | |||
| 97 | } // namespace ernie | ||
| 98 |