JobSystem.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/Reference.h>
#include <Jolt/Core/Color.h>
#include <Jolt/Core/Profiler.h>
#include <Jolt/Core/NonCopyable.h>
#include <Jolt/Core/StaticArray.h>
#include <Jolt/Core/Atomics.h>
 
JPH_NAMESPACE_BEGIN
 
class JPH_EXPORT JobSystem : public NonCopyable
{
protected:
    class Job;
 
public:
    JPH_OVERRIDE_NEW_DELETE
 
    class JobHandle : private Ref<Job>
    {
    public:
        inline              JobHandle()                                 = default;
        inline              JobHandle(const JobHandle &inHandle)        = default;
        inline              JobHandle(JobHandle &&inHandle) noexcept    : Ref<Job>(std::move(inHandle)) { }
 
        inline explicit     JobHandle(Job *inJob)                       : Ref<Job>(inJob) { }
 
        inline JobHandle &  operator = (const JobHandle &inHandle)      = default;
        inline JobHandle &  operator = (JobHandle &&inHandle) noexcept  = default;
 
        inline bool         IsValid() const                             { return GetPtr() != nullptr; }
 
        inline bool         IsDone() const                              { return GetPtr() != nullptr && GetPtr()->IsDone(); }
 
        inline void         AddDependency(int inCount = 1) const        { GetPtr()->AddDependency(inCount); }
 
        inline void         RemoveDependency(int inCount = 1) const     { GetPtr()->RemoveDependencyAndQueue(inCount); }
 
        static inline void  sRemoveDependencies(const JobHandle *inHandles, uint inNumHandles, int inCount = 1);
 
        template <uint N>
        static inline void  sRemoveDependencies(StaticArray<JobHandle, N> &inHandles, int inCount = 1)
        {
            sRemoveDependencies(inHandles.data(), inHandles.size(), inCount);
        }
 
        using Ref<Job>::GetPtr;
    };
 
    class Barrier : public NonCopyable
    {
    public:
        JPH_OVERRIDE_NEW_DELETE
 
        virtual void        AddJob(const JobHandle &inJob) = 0;
 
        virtual void        AddJobs(const JobHandle *inHandles, uint inNumHandles) = 0;
 
    protected:
        friend class Job;
 
        virtual             ~Barrier() = default;
 
        virtual void        OnJobFinished(Job *inJob) = 0;
    };
 
    using JobFunction = function<void()>;
 
    virtual                 ~JobSystem() = default;
 
    virtual int             GetMaxConcurrency() const = 0;
 
    virtual JobHandle       CreateJob(const char *inName, ColorArg inColor, const JobFunction &inJobFunction, uint32 inNumDependencies = 0) = 0;
 
    virtual void            CreateJobs(const char *inName, ColorArg inColor, const JobFunction &inJobFunction, JobHandle *outJobs, uint32 inNumJobs, uint32 inNumDependencies = 0) = 0;
 
    virtual Barrier *       CreateBarrier() = 0;
 
    virtual void            DestroyBarrier(Barrier *inBarrier) = 0;
 
    virtual void            WaitForJobs(Barrier *inBarrier) = 0;
 
protected:
    class Job
    {
    public:
        JPH_OVERRIDE_NEW_DELETE
 
                            Job([[maybe_unused]] const char *inJobName, [[maybe_unused]] ColorArg inColor, JobSystem *inJobSystem, const JobFunction &inJobFunction, uint32 inNumDependencies) :
            mJobFunction(inJobFunction),
        #if defined(JPH_EXTERNAL_PROFILE) || defined(JPH_PROFILE_ENABLED)
            mJobName(inJobName),
            mColor(inColor),
        #endif // defined(JPH_EXTERNAL_PROFILE) || defined(JPH_PROFILE_ENABLED)
            mJobSystem(inJobSystem),
            mNumDependencies(inNumDependencies)
        {
        }
 
        inline JobSystem *  GetJobSystem()                              { return mJobSystem; }
 
        inline void         AddRef(uint32 inCount = 1)
        {
            // Adding a reference can use relaxed memory ordering
            mReferenceCount.fetch_add(inCount, memory_order_relaxed);
        }
        inline void         Release()
        {
        #ifndef JPH_TSAN_ENABLED
            // Releasing a reference must use release semantics...
            if (mReferenceCount.fetch_sub(1, memory_order_release) == 1)
            {
                // ... so that we can use acquire to ensure that we see any updates from other threads that released a ref before freeing the job
                atomic_thread_fence(memory_order_acquire);
                mJobSystem->FreeJob(this);
            }
        #else
            // But under TSAN, we cannot use atomic_thread_fence, so we use an acq_rel operation unconditionally instead
            if (mReferenceCount.fetch_sub(1, memory_order_acq_rel) == 1)
                mJobSystem->FreeJob(this);
        #endif
        }
 
        inline void         AddDependency(int inCount);
 
        inline bool         RemoveDependency(int inCount);
 
        inline void         RemoveDependencyAndQueue(int inCount);
 
        inline bool         SetBarrier(Barrier *inBarrier)
        {
            // Acquire on failure so that, if the job already finished, the caller sees all side effects of the job function
            intptr_t barrier = 0;
            if (mBarrier.compare_exchange_strong(barrier, reinterpret_cast<intptr_t>(inBarrier), memory_order_release, memory_order_acquire))
                return true;
            JPH_ASSERT(barrier == cBarrierDoneState, "A job can only belong to 1 barrier");
            return false;
        }
 
        inline uint32       Execute()
        {
            // Transition job to executing state
            uint32 state = 0; // We can only start running with a dependency counter of 0
            if (!mNumDependencies.compare_exchange_strong(state, cExecutingState, memory_order_acquire))
                return state; // state is updated by compare_exchange_strong to the current value
 
            // Run the job function
            {
                JPH_PROFILE(mJobName, mColor.GetUInt32());
                mJobFunction();
            }
 
            // Clear the job function. It is a lambda with captured state that we want to destruct in a controlled fashion.
            // If we release it after notifying the barrier then that state might access deleted data.
            mJobFunction = {};
 
            // Fetch the barrier pointer and exchange it for the done state, so we're sure that no barrier gets set after we want to call the callback
            intptr_t barrier = mBarrier.load(memory_order_relaxed);
            for (;;)
            {
                // Acquire to pair with the release in SetBarrier, release so that a failing SetBarrier sees the side effects of the job function
                if (mBarrier.compare_exchange_weak(barrier, cBarrierDoneState, memory_order_acq_rel, memory_order_relaxed))
                    break;
            }
            JPH_ASSERT(barrier != cBarrierDoneState);
 
            // Mark job as done
            state = cExecutingState;
            mNumDependencies.compare_exchange_strong(state, cDoneState, memory_order_relaxed);
            JPH_ASSERT(state == cExecutingState);
 
            // Notify the barrier after we've changed the job to the done state so that any thread reading the state after receiving the callback will see that the job has finished
            if (barrier != 0)
                reinterpret_cast<Barrier *>(barrier)->OnJobFinished(this);
 
            return cDoneState;
        }
 
        inline bool         CanBeExecuted() const                       { return mNumDependencies.load(memory_order_relaxed) == 0; }
 
        inline bool         IsDone() const                              { return mNumDependencies.load(memory_order_relaxed) == cDoneState; }
 
    #if defined(JPH_EXTERNAL_PROFILE) || defined(JPH_PROFILE_ENABLED)
        const char *        GetName() const                             { return mJobName; }
    #endif // defined(JPH_EXTERNAL_PROFILE) || defined(JPH_PROFILE_ENABLED)
 
        static constexpr uint32 cExecutingState = 0xe0e0e0e0;           
        static constexpr uint32 cDoneState      = 0xd0d0d0d0;           
 
        static constexpr intptr_t cBarrierDoneState = ~intptr_t(0);     
 
private:
        JobFunction         mJobFunction;                               
    #if defined(JPH_EXTERNAL_PROFILE) || defined(JPH_PROFILE_ENABLED)
        const char *        mJobName;                                   
        Color               mColor;                                     
    #endif // defined(JPH_EXTERNAL_PROFILE) || defined(JPH_PROFILE_ENABLED)
        JobSystem *         mJobSystem;                                 
        atomic<intptr_t>    mBarrier = 0;                               
        atomic<uint32>      mReferenceCount = 0;                        
        atomic<uint32>      mNumDependencies;                           
    };
 
    virtual void            QueueJob(Job *inJob) = 0;
 
    virtual void            QueueJobs(Job **inJobs, uint inNumJobs) = 0;
 
    virtual void            FreeJob(Job *inJob) = 0;
};
 
using JobHandle = JobSystem::JobHandle;
 
JPH_NAMESPACE_END
 
#include "JobSystem.inl"