common/thread_worker: Stop workers on stop_token when waiting
This commit is contained in:
parent
da34d37044
commit
0ddbbb64e5
|
@ -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<>;
|
||||||
|
|
Reference in New Issue