From 7b9e0428421149786a818063e83c6c73a8280ef9 Mon Sep 17 00:00:00 2001
From: lisyarus <lisyarus@gmail.com>
Date: Wed, 30 Sep 2020 16:20:48 +0300
Subject: [PATCH] Fix threadpool::wait - wait for finished tasks as well, not
 just empty task queue

---
 libs/util/include/psemek/util/threadpool.hpp | 10 ++++++++++
 libs/util/source/threadpool.cpp              |  9 +++++++++
 2 files changed, 19 insertions(+)

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 <future>
 #include <vector>
 #include <string>
+#include <mutex>
+#include <condition_variable>
 
 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<util::thread> threads;
 		util::synchronized_queue<movable_function<void()>> 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();
 				}
 			});
 		}