MLX
 
Loading...
Searching...
No Matches
threadpool.h
Go to the documentation of this file.
1// This code was modified from https://github.com/progschj/ThreadPool
2// The original License is copied below:
3//
4// Copyright (c) 2012 Jakob Progsch, Václav Zeman
5// This software is provided 'as-is', without any express or implied
6// warranty. In no event will the authors be held liable for any damages
7// arising from the use of this software.
8//
9// Permission is granted to anyone to use this software for any purpose,
10// including commercial applications, and to alter it and redistribute it
11// freely, subject to the following restrictions:
12//
13// 1. The origin of this software must not be misrepresented; you must not
14// claim that you wrote the original software. If you use this software
15// in a product, an acknowledgment in the product documentation would be
16// appreciated but is not required.
17//
18// 2. Altered source versions must be plainly marked as such, and must not be
19// misrepresented as being the original software.
20//
21// 3. This notice may not be removed or altered from any source
22// distribution.
23#pragma once
24
25#include <condition_variable>
26#include <functional>
27#include <future>
28#include <memory>
29#include <mutex>
30#include <queue>
31#include <stdexcept>
32#include <thread>
33#include <vector>
34
36 public:
37 ThreadPool(size_t);
38 template <class F, class... Args>
39 auto enqueue(F&& f, Args&&... args)
40 -> std::future<typename std::invoke_result_t<F, Args...>>;
41 void resize(size_t);
43
44 private:
45 void stop_and_wait();
46 void start_threads(size_t);
47
48 std::vector<std::thread> workers;
49 std::queue<std::function<void()>> tasks;
50 std::mutex queue_mutex;
51 std::condition_variable condition;
52 bool stop;
53};
54
55inline ThreadPool::ThreadPool(size_t threads) : stop(false) {
56 start_threads(threads);
57}
58
59template <class F, class... Args>
60auto ThreadPool::enqueue(F&& f, Args&&... args)
61 -> std::future<typename std::invoke_result_t<F, Args...>> {
62 using return_type = typename std::invoke_result_t<F, Args...>;
63
64 auto task = std::make_shared<std::packaged_task<return_type()>>(
65 std::bind(std::forward<F>(f), std::forward<Args>(args)...));
66
67 std::future<return_type> res = task->get_future();
68 {
69 std::unique_lock<std::mutex> lock(queue_mutex);
70
71 if (stop) {
72 throw std::runtime_error(
73 "[ThreadPool::enqueue] Not allowed on stopped ThreadPool");
74 }
75
76 tasks.emplace([task]() { (*task)(); });
77 }
78 condition.notify_one();
79 return res;
80}
81
82inline void ThreadPool::resize(size_t threads) {
83 if (workers.size() == threads) {
84 return;
85 }
86
87 if (workers.size() > threads) {
88 stop_and_wait();
89 }
90 start_threads(threads - workers.size());
91}
92
94 stop_and_wait();
95}
96
97inline void ThreadPool::stop_and_wait() {
98 // Stop the current threads and wait until they finish
99 {
100 std::unique_lock<std::mutex> lock(queue_mutex);
101 stop = true;
102 }
103 condition.notify_all();
104 for (std::thread& worker : workers) {
105 worker.join();
106 }
107
108 // Reset the member variables so that the threadpool is reusable
109 stop = false;
110 workers.clear();
111}
112
113inline void ThreadPool::start_threads(size_t threads) {
114 for (size_t i = 0; i < threads; ++i) {
115 workers.emplace_back([this] {
116 for (;;) {
117 std::function<void()> task;
118
119 {
120 std::unique_lock<std::mutex> lock(this->queue_mutex);
121 this->condition.wait(
122 lock, [this] { return this->stop || !this->tasks.empty(); });
123 if (this->stop && this->tasks.empty())
124 return;
125 task = std::move(this->tasks.front());
126 this->tasks.pop();
127 }
128
129 task();
130 }
131 });
132 }
133}
void resize(size_t)
Definition threadpool.h:82
auto enqueue(F &&f, Args &&... args) -> std::future< typename std::invoke_result_t< F, Args... > >
Definition threadpool.h:60
~ThreadPool()
Definition threadpool.h:93
ThreadPool(size_t)
Definition threadpool.h:55