flodavid/suyu
1
0
Fork 0
forked from suyu/suyu
Code Pull requests Activity

common/thread_worker: Stop workers on stop_token when waiting

Browse source
This commit is contained in:
ReinUsesLisp 2021-06-29 01:47:13 -03:00
parent da34d37044
commit 0ddbbb64e5
1 changed files with 20 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
src/common/thread_worker.h
View file

@ -7,7 +7,9 @@
#include <atomic> #include <atomic>
#include <functional> #include <functional>
#include <mutex> #include <mutex>
#include <stop_token>
#include <string> #include <string>
#include <thread>
#include <type_traits> #include <type_traits>
#include <vector> #include <vector>
#include <queue> #include <queue>
@ -34,19 +36,19 @@ class StatefulThreadWorker {
public: public:
explicit StatefulThreadWorker(size_t num_workers, std::string name, StateMaker func = {}) explicit StatefulThreadWorker(size_t num_workers, std::string name, StateMaker func = {})
: workers_queued{num_workers}, thread_name{std::move(name)} { : workers_queued{num_workers}, thread_name{std::move(name)} {
const auto lambda = [this, func] { const auto lambda = [this, func](std::stop_token stop_token) {
Common::SetCurrentThreadName(thread_name.c_str()); Common::SetCurrentThreadName(thread_name.c_str());
{ {
std::conditional_t<with_state, StateType, int> state{func()}; std::conditional_t<with_state, StateType, int> state{func()};
while (!stop) { while (!stop_token.stop_requested()) {
Task task; Task task;
{ {
std::unique_lock lock{queue_mutex}; std::unique_lock lock{queue_mutex};
if (requests.empty()) { if (requests.empty()) {
wait_condition.notify_all(); wait_condition.notify_all();
} }
condition.wait(lock, [this] { return stop || !requests.empty(); }); condition.wait(lock, stop_token, [this] { return !requests.empty(); });
if (stop) { if (stop_token.stop_requested()) {
break; break;
} }
task = std::move(requests.front()); task = std::move(requests.front());
@ -63,21 +65,17 @@ public:
++workers_stopped; ++workers_stopped;
wait_condition.notify_all(); wait_condition.notify_all();
}; };
threads.reserve(num_workers);
for (size_t i = 0; i < num_workers; ++i) { for (size_t i = 0; i < num_workers; ++i) {
threads.emplace_back(lambda); threads.emplace_back(lambda);
} }
} }
~StatefulThreadWorker() { StatefulThreadWorker& operator=(const StatefulThreadWorker&) = delete;
{ StatefulThreadWorker(const StatefulThreadWorker&) = delete;
std::unique_lock lock{queue_mutex};
stop = true; StatefulThreadWorker& operator=(StatefulThreadWorker&&) = delete;
} StatefulThreadWorker(StatefulThreadWorker&&) = delete;
condition.notify_all();
for (std::thread& thread : threads) {
thread.join();
}
}
void QueueWork(Task work) { void QueueWork(Task work) {
{ {
@ -88,7 +86,12 @@ public:
condition.notify_one(); condition.notify_one();
} }
void WaitForRequests() { void WaitForRequests(std::stop_token stop_token = {}) {
std::stop_callback callback(stop_token, [this] {
for (auto& thread : threads) {
thread.request_stop();
}
});
std::unique_lock lock{queue_mutex}; std::unique_lock lock{queue_mutex};
wait_condition.wait(lock, [this] { wait_condition.wait(lock, [this] {
return workers_stopped >= workers_queued || work_done >= work_scheduled; return workers_stopped >= workers_queued || work_done >= work_scheduled;
@ -96,17 +99,16 @@ public:
} }
private: private:
std::vector<std::thread> threads;
std::queue<Task> requests; std::queue<Task> requests;
std::mutex queue_mutex; std::mutex queue_mutex;
std::condition_variable condition; std::condition_variable_any condition;
std::condition_variable wait_condition; std::condition_variable wait_condition;
std::atomic_bool stop{};
std::atomic<size_t> work_scheduled{}; std::atomic<size_t> work_scheduled{};
std::atomic<size_t> work_done{}; std::atomic<size_t> work_done{};
std::atomic<size_t> workers_stopped{}; std::atomic<size_t> workers_stopped{};
std::atomic<size_t> workers_queued{}; std::atomic<size_t> workers_queued{};
std::string thread_name; std::string thread_name;
std::vector<std::jthread> threads;
}; };
using ThreadWorker = StatefulThreadWorker<>; using ThreadWorker = StatefulThreadWorker<>;