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_