JobSystemThreadPool.h

Go to the documentation of this file.

// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
// SPDX-FileCopyrightText: 2021 Jorrit Rouwe
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <Jolt/Core/JobSystemWithBarrier.h>
#include <Jolt/Core/FixedSizeFreeList.h>
#include <Jolt/Core/Semaphore.h>
 
JPH_SUPPRESS_WARNINGS_STD_BEGIN
#include <thread>
JPH_SUPPRESS_WARNINGS_STD_END
 
JPH_NAMESPACE_BEGIN
 
// Things we're using from STL
using std::thread;
 
class JPH_EXPORT JobSystemThreadPool final : public JobSystemWithBarrier
{
public:
    JPH_OVERRIDE_NEW_DELETE
 
                            JobSystemThreadPool(uint inMaxJobs, uint inMaxBarriers, int inNumThreads = -1);
                            JobSystemThreadPool() = default;
    virtual                 ~JobSystemThreadPool() override;
 
    using InitExitFunction = function<void(int)>;
    void                    SetThreadInitFunction(const InitExitFunction &inInitFunction)   { mThreadInitFunction = inInitFunction; }
    void                    SetThreadExitFunction(const InitExitFunction &inExitFunction)   { mThreadExitFunction = inExitFunction; }
 
    void                    Init(uint inMaxJobs, uint inMaxBarriers, int inNumThreads = -1);
 
    // See JobSystem
    virtual int             GetMaxConcurrency() const override              { return int(mThreads.size()) + 1; }
    virtual JobHandle       CreateJob(const char *inName, ColorArg inColor, const JobFunction &inJobFunction, uint32 inNumDependencies = 0) override;
    virtual void            CreateJobs(const char *inName, ColorArg inColor, const JobFunction &inJobFunction, JobHandle *outJobs, uint32 inNumJobs, uint32 inNumDependencies = 0) override;
 
    void                    SetNumThreads(int inNumThreads)                 { StopThreads(); StartThreads(inNumThreads); }
 
protected:
    // See JobSystem
    virtual void            QueueJob(Job *inJob) override;
    virtual void            QueueJobs(Job **inJobs, uint inNumJobs) override;
    virtual void            FreeJob(Job *inJob) override;
 
private:
    class Node : public Job
    {
    public:
                            Node(const char *inJobName, ColorArg inColor, JobSystem *inJobSystem, const JobFunction &inJobFunction, uint32 inNumDependencies)
            : Job(inJobName, inColor, inJobSystem, inJobFunction, inNumDependencies)
        {
            // Work around TSAN false positive. We can read mNext from a freed and then recyled node. So instead of constructing, we atomically intialize mNext to 0
            mNext.store(0, std::memory_order_relaxed);
        }
 
        atomic<uint64>      mNext;                                          
    };
 
    static constexpr uint32 cNullIndex = 0xffffffff;
    static uint64           sMakePointer(uint32 inIndex, uint32 inCount)    { return (uint64(inCount) << 32) | inIndex; }
    static uint32           sGetIndex(uint64 inPointer)                     { return uint32(inPointer); }
    static uint32           sGetCount(uint64 inPointer)                     { return uint32(inPointer >> 32); }
 
    Job *                   AllocateJob(const char *inName, ColorArg inColor, const JobFunction &inJobFunction, uint32 inNumDependencies);
 
    void                    Enqueue(Job *inJob);
 
    Job *                   Dequeue();
 
    void                    StartThreads(int inNumThreads);
    void                    StopThreads();
 
    void                    ThreadMain(int inThreadIndex);
 
    void                    ExecuteJobs();
 
    InitExitFunction        mThreadInitFunction = [](int) { };
    InitExitFunction        mThreadExitFunction = [](int) { };
 
    using AvailableJobs = FixedSizeFreeList<Node>;
    AvailableJobs           mJobs;
 
    Array<thread>           mThreads;
 
    // Head and tail of the queue, these are index + counters (see sMakePointer)
    alignas(JPH_CACHE_LINE_SIZE) atomic<uint64> mHead;
    alignas(JPH_CACHE_LINE_SIZE) atomic<uint64> mTail;
 
    // Semaphore used to signal worker threads that there is new work
    Semaphore               mSemaphore;
 
    atomic<bool>            mQuit = false;
};
 
JPH_NAMESPACE_END