ThreadsPool.h

Go to the documentation of this file.

 
#pragma once
 
#ifndef __THREADS_POOL_H_
#define __THREADS_POOL_H_
 
#include <algorithm>
#include <condition_variable>
#include <limits>
#include <mutex>
#include <new>
#include <queue>
#include <vector>
 
#include <fstream>
 
#if defined(__unix__) || defined(__APPLE__)
#include <pthread.h>
#define MAESTRO_POOL_FORK_HANDLERS 1
#endif
 
#include "WorkerThread.h"
 
namespace Utils
{
 
#ifdef MAESTRO_POOL_FORK_HANDLERS
namespace Private
{
// Pools keep their threads between executions, but threads do not survive
// fork(): a child would queue jobs for workers that do not exist. Each pool
// holds its locks across fork(); in the child an idle pool drops its stale
// workers and creates new ones on its next Resize(). Threads are not stopped
// before fork(): their exit would wait on OpenMP's own fork locks.
class ForkAwarePool
{
  public:
    virtual ~ForkAwarePool() = default;
    virtual void PrepareFork() = 0;
    virtual void AfterForkParent() = 0;
    virtual void AfterForkChild() = 0;
};
 
class PoolForkRegistry
{
  public:
    static PoolForkRegistry &Instance()
    {
        // Never destroyed: pools owned by static objects unregister at exit.
        static auto *registry = new PoolForkRegistry;
        return *registry;
    }
 
    void Add(ForkAwarePool *pool)
    {
        std::call_once(installed, [] { pthread_atfork(&Prepare, &Parent, &Child); });
        const std::lock_guard lock(mutex);
        pools.push_back(pool);
    }
 
    void Remove(ForkAwarePool *pool)
    {
        const std::lock_guard lock(mutex);
        pools.erase(std::remove(pools.begin(), pools.end(), pool), pools.end());
    }
 
  private:
    // The registry lock is held from Prepare() until Parent()/Child().
    static void Prepare()
    {
        auto &registry = Instance();
        registry.mutex.lock();
        for (auto *pool : registry.pools)
            pool->PrepareFork();
    }
 
    static void Parent()
    {
        auto &registry = Instance();
        for (auto *pool : registry.pools)
            pool->AfterForkParent();
        registry.mutex.unlock();
    }
 
    static void Child()
    {
        auto &registry = Instance();
        for (auto *pool : registry.pools)
            pool->AfterForkChild();
        registry.mutex.unlock();
    }
 
    std::mutex mutex;
    std::vector<ForkAwarePool *> pools;
    std::once_flag installed;
};
} // namespace Private
#endif

template <class Job>
class ThreadsPool
#ifdef MAESTRO_POOL_FORK_HANDLERS
    : private Private::ForkAwarePool
#endif
{
    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));
#ifdef MAESTRO_POOL_FORK_HANDLERS
        Private::PoolForkRegistry::Instance().Add(this);
#endif
    }
 
    ThreadsPool(const ThreadsPool &) = delete;
    ThreadsPool &operator=(const ThreadsPool &) = delete;

    ~ThreadsPool()
    {
#ifdef MAESTRO_POOL_FORK_HANDLERS
        Private::PoolForkRegistry::Instance().Remove(this);
#endif
        Stop();
    }

    void Stop()
    {
        const std::lock_guard lifecycle(LifecycleMutex);
        {
            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;
 
        const std::lock_guard lifecycle(LifecycleMutex);
        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 = nrThreads; i < oldSize; ++i)
                    Threads[i]->SetStopUnlocked();
            }
 
            NotifyAll();
 
            for (size_t i = nrThreads; i < oldSize; ++i)
                Threads[i]->Join();
            Threads.resize(nrThreads);
        }
    }

    void SetFinishLimit(size_t limit)
    {
        std::unique_lock lock(FinishMutex);
        FinishLimit = limit;
        FinishCount = 0;
    }
 
  private:
#ifdef MAESTRO_POOL_FORK_HANDLERS
    void PrepareFork() override
    {
        // The locks are released by AfterForkParent() or AfterForkChild().
        LifecycleMutex.lock();
        Mutex.lock();
        FinishMutex.lock();
        idleAtFork = JobsQueue.empty() && Running == 0;
    }
 
    void AfterForkParent() override
    {
        FinishMutex.unlock();
        Mutex.unlock();
        LifecycleMutex.unlock();
    }
 
    void AfterForkChild() override
    {
        // Only the forking thread exists here. A pool in the middle of an
        // execution is left as it is; an idle one forgets its workers, whose
        // waits also left the copied condition variable unusable. Threads
        // that were waiting for jobs to finish are gone too.
        if (idleAtFork)
        {
            for (auto &worker : Threads)
                worker->Abandon();
            Threads.clear();
            new (&Condition) std::condition_variable;
        }
        new (&ConditionFinish) std::condition_variable;
        new (&FinishMutex) std::mutex;
        new (&Mutex) std::mutex;
        new (&LifecycleMutex) std::mutex;
    }
#endif

    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 LifecycleMutex;         
    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; 
    size_t Running = 0;                         
    bool idleAtFork = false;                    
 
    std::vector<std::unique_ptr<JobWorkerThread>> Threads; 
};
 
} // namespace Utils
 
#endif // __THREADS_POOL_H_