common/thread_worker: Add wait for requests method

2 files changed, 11 insertions, 0 deletions

diff --git a/src/common/thread_worker.cpp b/src/common/thread_worker.cpp
index 8f9bf447a..745918c7e 100644
--- a/src/common/thread_worker.cpp
+++ b/src/common/thread_worker.cpp
@@ -29,6 +29,10 @@ ThreadWorker::ThreadWorker(std::size_t num_workers, const std::string& name) {
                     }
                     task = std::move(requests.front());
                     requests.pop();
+
+                    if (requests.empty()) {
+                        wait_condition.notify_one();
+                    }
                 }
 
                 task();
@@ -55,4 +59,9 @@ void ThreadWorker::QueueWork(std::function<void()>&& work) {
     condition.notify_one();
 }
 
+void ThreadWorker::WaitForRequests() {
+    std::unique_lock lock{queue_mutex};
+    wait_condition.wait(lock, [this] { return stop || requests.empty(); });
+}
+
 } // namespace Common
diff --git a/src/common/thread_worker.h b/src/common/thread_worker.h
index f1859971f..7a6756eb5 100644
--- a/src/common/thread_worker.h
+++ b/src/common/thread_worker.h
@@ -18,12 +18,14 @@ public:
     explicit ThreadWorker(std::size_t num_workers, const std::string& name);
     ~ThreadWorker();
     void QueueWork(std::function<void()>&& work);
+    void WaitForRequests();
 
 private:
     std::vector<std::thread> threads;
     std::queue<std::function<void()>> requests;
     std::mutex queue_mutex;
     std::condition_variable condition;
+    std::condition_variable wait_condition;
     std::atomic_bool stop{};
 };