diff --git a/libs/util/include/psemek/util/threadpool.hpp b/libs/util/include/psemek/util/threadpool.hpp index 40f0b22e..cb9cdd33 100644 --- a/libs/util/include/psemek/util/threadpool.hpp +++ b/libs/util/include/psemek/util/threadpool.hpp @@ -7,6 +7,8 @@ #include #include #include +#include +#include namespace psemek::util { @@ -19,6 +21,7 @@ namespace psemek::util threadpool(std::string const & name, std::size_t thread_count) : name(name) + , working_count_{0} { start(thread_count); } @@ -49,12 +52,19 @@ namespace psemek::util void wait() { tasks_queue.wait(); + + std::unique_lock lock{working_count_mutex_}; + working_count_cv_.wait(lock, [this]{ return working_count_ == 0; }); } private: std::string const name; std::vector threads; util::synchronized_queue> tasks_queue; + + std::size_t working_count_; + std::mutex working_count_mutex_; + std::condition_variable working_count_cv_; }; } diff --git a/libs/util/source/threadpool.cpp b/libs/util/source/threadpool.cpp index 9faa2989..cee1a6e0 100644 --- a/libs/util/source/threadpool.cpp +++ b/libs/util/source/threadpool.cpp @@ -21,7 +21,16 @@ namespace psemek::util if (!task) break; + { + std::lock_guard lock{working_count_mutex_}; + ++working_count_; + } task(); + { + std::lock_guard lock{working_count_mutex_}; + --working_count_; + } + working_count_cv_.notify_all(); } }); }