onnxruntime/winml/test/concurrency/ThreadPool.cpp

#include "testPch.h"
#include "ThreadPool.h"
#include <ctime>

ThreadPool::ThreadPool(unsigned int initial_pool_size) : m_destruct_pool(false), m_threads() {
  for (unsigned int i = 0; i < initial_pool_size; i++) {
    m_threads.emplace_back([this]() {
      while (true) {
        std::unique_lock<std::mutex> lock(m_mutex);
        // thread listening for event and acquire lock if event triggered
        m_cond_var.wait(lock, [this] { return m_destruct_pool || !m_work_queue.empty(); });
        if (!m_work_queue.empty()) {
          auto work = m_work_queue.front();
          m_work_queue.pop();
          lock.unlock();
          work();
        } else {
          // Work queue is empty but lock acquired
          // This means we are destructing the pool
          break;
        }
      }
    });
  }
}

ThreadPool::~ThreadPool() {
  m_destruct_pool = true;
  m_cond_var.notify_all();  // notify destruction to threads
  for (auto& thread : m_threads) {
    thread.join();
  }
}