ThreadsPool.h

Go to the documentation of this file.

 
#pragma once
 
#ifndef __THREADS_POOL_H_
#define __THREADS_POOL_H_
 
#include <condition_variable>
#include <queue>
#include <vector>
 
#include <fstream>
 
#include "WorkerThread.h"
 
namespace Utils {
 
template <class Job>
class ThreadsPool {
  using JobWorkerThread = WorkerThread<ThreadsPool<Job>, Job>;
  friend class WorkerThread<ThreadsPool<Job>, Job>;
 
 public:
  explicit ThreadsPool(int nrThreads = 0) {
    if (nrThreads <= 0) nrThreads = 1;
 
    for (int i = 0; i < nrThreads; ++i)
      Threads.emplace_back(std::make_unique<JobWorkerThread>(this));
  }
 
  ~ThreadsPool() { Stop(); }
 
  void Stop() {
    {
      std::lock_guard lock(Mutex);
 
      for (auto &worker : Threads) worker->SetStopUnlocked();
    }
 
    NotifyAll();
 
    for (auto &worker : Threads) worker->Join();
 
    Threads.clear();
  }
 
  void Start() {
    for (auto &worker : Threads) worker->Start();
  }
 
  void AddRunJob(const std::shared_ptr<Job> &job) {
    {
      std::lock_guard lock(Mutex);
      JobsQueue.push(job);
    }
 
    NotifyOne();
  }
 
  void AddRunJob(std::shared_ptr<Job> &&job) {
    {
      std::lock_guard lock(Mutex);
      JobsQueue.push(std::move(job));
    }
 
    NotifyOne();
  }
 
  void WaitForFinish() {
    do {
      std::unique_lock lock(FinishMutex);
 
      if (FinishCount >= FinishLimit) return;
 
      ConditionFinish.wait(lock, [this] { return FinishCount >= FinishLimit; });
 
      if (FinishCount >= FinishLimit) return;
    } while (true);
  }
 
  void Resize(size_t nrThreads) {
    if (nrThreads <= 0) nrThreads = 1;
 
    size_t oldSize = Threads.size();
 
    if (oldSize == nrThreads) return;
 
    if (oldSize < nrThreads) {
      for (size_t i = oldSize; i < nrThreads; ++i) {
        Threads.emplace_back(std::make_unique<JobWorkerThread>(this));
        Threads.back()->Start();
      }
    } else {
      {
        std::lock_guard lock(Mutex);
        for (size_t i = oldSize; i < nrThreads; ++i)
          Threads[i]->SetStopUnlocked();
      }
 
      NotifyAll();
 
      for (size_t i = oldSize; i < nrThreads; ++i) Threads[i]->Join();
      Threads.resize(nrThreads);
    }
  }
 
  void SetFinishLimit(size_t limit) {
    std::unique_lock lock(FinishMutex);
    FinishLimit = limit;
    FinishCount = 0;
  }
 
 private:
  inline void NotifyOne() { Condition.notify_one(); }
 
  inline void NotifyAll() { Condition.notify_all(); }
 
  inline bool HasWork() const { return !JobsQueue.empty(); }
 
  void NotifyFinish() { ConditionFinish.notify_all(); }
 
  std::mutex Mutex; 
  std::condition_variable
      Condition; 
  std::mutex FinishMutex; 
  std::condition_variable ConditionFinish; 
  size_t FinishCount = 0; 
  size_t FinishLimit =
      std::numeric_limits<size_t>::max(); 
  std::queue<std::shared_ptr<Job>>
      JobsQueue; 
  std::vector<std::unique_ptr<JobWorkerThread>>
      Threads; 
};
 
}  // namespace Utils
 
#endif  // __THREADS_POOL_H_