onnxruntime/winml/test/concurrency/ThreadPool.h

#pragma once

#include <vector>
#include <thread>
#include <queue>
#include <mutex>
#include <future>

class ThreadPool {
 private:
  std::condition_variable m_cond_var;
  bool m_destruct_pool;
  std::mutex m_mutex;
  std::vector<std::thread> m_threads;
  std::queue<std::function<void()>> m_work_queue;

 public:
  ThreadPool(unsigned int initial_pool_size);
  ~ThreadPool();
  template <typename F, typename... Args>
  inline auto SubmitWork(F&& f, Args&&... args) -> std::future<decltype(f(args...))> {
    auto func = std::bind(std::forward<F>(f), std::forward<Args>(args)...);
    auto task = std::make_shared<std::packaged_task<decltype(f(args...))()>>(std::forward<decltype(func)>(func));
    {
      std::lock_guard<std::mutex> lock(m_mutex);
      // wrap packed task into a void return function type so that it can be stored in queue
      m_work_queue.push([task]() { (*task)(); });
    }

    m_cond_var.notify_one();  // unblocks one of the waiting threads
    return task->get_future();
  }
};