ThreadPool.hpp

Go to the documentation of this file.

#pragma once
#include "AtomicQueue.hpp"
#include "Thread.hpp"
#include <algorithm>
#include <array>
#include <cmath>
#include <iterator>
#include <type_traits>
namespace Foundation::Core
{
    enum class JobPriority : size_t
    {
        Low,
        Normal,
        High,
    };
    inline constexpr size_t kJobPriorityCount = static_cast<size_t>(JobPriority::High) + 1;
 
    enum class ExecutionPolicy
    {
        Seq, // run inline on the calling thread, in index order
        Par, // fork-join across the pool (default)
    };
 
    struct Job
    {
        virtual ~Job() = default;
        virtual void Execute(size_t id) noexcept = 0;
    };
    template <typename Lambda, typename ReturnType, typename... Args>
    struct LambdaJob final : public Job
    {
        Lambda mFunc;
        Promise<ReturnType> mPromise;
        LambdaJob(Lambda&& func) : mFunc(std::forward<Lambda>(func)) {}
        void Execute(size_t) noexcept override
        {
            try
            {
                if constexpr (std::is_same_v<ReturnType, void>)
                {
                    mFunc();
                    mPromise.set_value();
                }
                else
                    mPromise.set_value(mFunc());
            }
            catch (...)
            {
                mPromise.set_exception(std::current_exception());
            }
        }
    };
    template <typename Fn, typename Arg>
    inline void ParallelForInvoke(Fn& fn, Arg&& arg, size_t workerId)
    {
        if constexpr (std::is_invocable_v<Fn&, Arg&&, size_t>)
            fn(std::forward<Arg>(arg), workerId);
        else
            fn(std::forward<Arg>(arg));
    }
 
    template <typename Fn>
    struct ParallelForJob final : Job
    {
        Fn* fn{nullptr};
        size_t total{0};
        Atomic<size_t> cursor{0};
        Atomic<size_t> remaining{0}; // co-invocations still running (fork-join latch)
        void Execute(size_t workerId) noexcept override
        {
            for (size_t i; (i = cursor.fetch_add(1, std::memory_order_relaxed)) < total;)
                ParallelForInvoke(*fn, i, workerId);
            if (remaining.fetch_sub(1, std::memory_order_acq_rel) == 1)
                remaining.notify_all();
        }
    };
 
    template <typename Fn>
    struct ParallelForAsyncState
    {
        Fn fn;
        size_t total{0};
        Atomic<size_t> cursor{0};
        Atomic<size_t> remaining{0}; // pushed worker jobs still running (fork-join latch)
        Promise<void> promise;
        Allocator* alloc{nullptr};
        ParallelForAsyncState(Fn&& f, size_t t, Allocator* a) : fn(std::move(f)), total(t), alloc(a) {}
    };
 
    template <typename Fn>
    struct ParallelForAsyncJob final : Job
    {
        ParallelForAsyncState<Fn>* state{nullptr};
        explicit ParallelForAsyncJob(ParallelForAsyncState<Fn>* s) : state(s) {}
        void Execute(size_t workerId) noexcept override
        {
            for (size_t i; (i = state->cursor.fetch_add(1, std::memory_order_relaxed)) < state->total;)
                ParallelForInvoke(state->fn, i, workerId);
            // The worker that observes the latch reach zero is the last one touching `state`.
            if (state->remaining.fetch_sub(1, std::memory_order_acq_rel) == 1)
            {
                Allocator* alloc = state->alloc;
                state->promise.set_value();
                Destruct(alloc, state);
            }
        }
    };
 
    using JobQueue = MPMCQueue<UniquePtr<Job>>;
    using JobQueues = std::array<JobQueue, kJobPriorityCount>;
    class ThreadPool
    {
        Allocator* mAllocator;
        String mName;
        Atomic<bool> mShutdown{false};
        Atomic<size_t> mComplete{0};
        Atomic<size_t> mTotal{0};
 
        JobQueues mJobs;
        // Ensure threads are joined first on destruction
        Vector<Thread> mThreads;
        void ThreadPoolWorker(size_t id);
        static constexpr size_t PriorityIndex(JobPriority priority) noexcept { return static_cast<size_t>(priority); }
        template <typename T, typename... Args>
            requires std::is_base_of_v<Job, T>
        T* PushImplInternal(JobPriority priority, Allocator* jobAllocator, Args&&... args)
        {
            if (mShutdown)
                throw std::runtime_error("ThreadPool shutting down");
            if (PriorityIndex(priority) >= kJobPriorityCount)
                throw std::runtime_error("Invalid job priority");
            Allocator* allocator = jobAllocator ? jobAllocator : mAllocator;
            auto task = ConstructUniqueBase<Job, T>(allocator, std::forward<Args>(args)...);
            T* ptr = static_cast<T*>(task.get());
            if (!mJobs[PriorityIndex(priority)].Push(std::move(task)))
                throw std::runtime_error("Jobs full");
            mTotal.fetch_add(1, std::memory_order_relaxed);
            mTotal.notify_one();
            return ptr;
        }
        template <typename Lambda, typename... Args>
        auto PushLambdaInternal(JobPriority priority, Allocator* jobAllocator, Lambda&& func, Args const&... args)
        {
            auto LambdaFn = [func = std::forward<Lambda>(func), ... args = args] { return func(args...); };
            using LambdaType = decltype(LambdaFn);
            using ReturnType = decltype(LambdaFn());
            LambdaJob<LambdaType, ReturnType> job(std::forward<LambdaType>(LambdaFn));
            auto fut = job.mPromise.get_future();
            PushImplInternal<LambdaJob<LambdaType, ReturnType>>(priority, jobAllocator, std::move(job));
            return std::move(fut);
        }
 
    public:
        ThreadPool(size_t numThreads, size_t maxTasks, Allocator* alloc, StringView name = "ThreadPool");
        template <typename T, typename... Args>
            requires std::is_base_of_v<Job, T>
        T* PushImpl(JobPriority priority, Args&&... args)
        {
            return PushImplInternal<T>(priority, nullptr, std::forward<Args>(args)...);
        }
        template <typename T, typename... Args>
            requires std::is_base_of_v<Job, T>
        T* PushImpl(Args&&... args)
        {
            return PushImplInternal<T>(JobPriority::Normal, nullptr, std::forward<Args>(args)...);
        }
        template <typename T, typename... Args>
            requires std::is_base_of_v<Job, T>
        T* PushImplAlloc(Allocator* jobAllocator, Args&&... args)
        {
            return PushImplInternal<T>(JobPriority::Normal, jobAllocator, std::forward<Args>(args)...);
        }
        template <typename T, typename... Args>
            requires std::is_base_of_v<Job, T>
        T* PushImplAlloc(JobPriority priority, Allocator* jobAllocator, Args&&... args)
        {
            return PushImplInternal<T>(priority, jobAllocator, std::forward<Args>(args)...);
        }
        template <typename Lambda, typename... Args>
        auto Push(JobPriority priority, Lambda&& func, Args const&... args)
        {
            return PushLambdaInternal(priority, nullptr, std::forward<Lambda>(func), args...);
        }
        template <typename Lambda, typename... Args>
        auto Push(Lambda&& func, Args const&... args)
        {
            return PushLambdaInternal(JobPriority::Normal, nullptr, std::forward<Lambda>(func), args...);
        }
        template <typename Lambda, typename... Args>
        auto PushAlloc(Allocator* jobAllocator, Lambda&& func, Args const&... args)
        {
            return PushLambdaInternal(JobPriority::Normal, jobAllocator, std::forward<Lambda>(func), args...);
        }
        template <typename Lambda, typename... Args>
        auto PushAlloc(JobPriority priority, Allocator* jobAllocator, Lambda&& func, Args const&... args)
        {
            return PushLambdaInternal(priority, jobAllocator, std::forward<Lambda>(func), args...);
        }
        void CoInvoke(Job& job, size_t count, JobPriority priority = JobPriority::Normal);
 
        [[nodiscard]] size_t GetWorkerCount() const noexcept { return mThreads.size(); }
 
        [[nodiscard]] size_t GetParallelForConcurrency() const noexcept { return mThreads.size() + 1; }
 
        template <typename Fn>
        void ParallelFor(ExecutionPolicy policy, size_t count, Fn&& fn)
        {
            if (count == 0)
                return;
            size_t const workers = mThreads.size();
            if (policy == ExecutionPolicy::Seq || workers == 0)
            {
                for (size_t i = 0; i < count; ++i)
                    ParallelForInvoke(fn, i, size_t{0}); // inline: caller is worker 0
                return;
            }
            ParallelForJob<std::remove_reference_t<Fn>> job;
            job.fn = &fn;
            job.total = count;
            size_t const helpers = std::min(workers, count);
            job.remaining.store(helpers + 1, std::memory_order_relaxed); // + the participating caller
            CoInvoke(job, helpers);
            // Participate on the calling thread with id == workers (its own scratch slot).
            for (size_t i; (i = job.cursor.fetch_add(1, std::memory_order_relaxed)) < count;)
                ParallelForInvoke(fn, i, workers);
            // Fork-join: every co-invocation must return before the stack job dies.
            if (job.remaining.fetch_sub(1, std::memory_order_acq_rel) != 1)
            {
                size_t r;
                while ((r = job.remaining.load(std::memory_order_acquire)) != 0)
                    job.remaining.wait(r, std::memory_order_acquire);
            }
        }
        template <typename Fn>
        void ParallelFor(size_t count, Fn&& fn)
        {
            ParallelFor(ExecutionPolicy::Par, count, std::forward<Fn>(fn));
        }
 
        template <typename It, typename Fn>
            requires std::random_access_iterator<It>
        void ParallelFor(ExecutionPolicy policy, It first, It last, Fn&& fn)
        {
            auto const count = last - first;
            if (count <= 0)
                return;
            ParallelFor(policy, static_cast<size_t>(count), [&](size_t i, size_t worker)
                        { ParallelForInvoke(fn, first[static_cast<std::iter_difference_t<It>>(i)], worker); });
        }
        template <typename It, typename Fn>
            requires std::random_access_iterator<It>
        void ParallelFor(It first, It last, Fn&& fn)
        {
            ParallelFor(ExecutionPolicy::Par, first, last, std::forward<Fn>(fn));
        }
 
        [[nodiscard]] static Future<void> MakeReadyFuture()
        {
            Promise<void> p;
            p.set_value();
            return p.get_future();
        }
 
        template <typename Fn>
        [[nodiscard]] Future<void> ParallelForAsync(ExecutionPolicy policy, size_t count, Fn&& fn)
        {
            size_t const workers = mThreads.size();
            if (count == 0)
                return MakeReadyFuture();
            if (policy == ExecutionPolicy::Seq || workers == 0)
            {
                for (size_t i = 0; i < count; ++i)
                    ParallelForInvoke(fn, i, size_t{0}); // inline: caller is worker 0
                return MakeReadyFuture();
            }
            using State = ParallelForAsyncState<std::remove_reference_t<Fn>>;
            State* state = Construct<State>(mAllocator, std::forward<Fn>(fn), count, mAllocator);
            size_t const helpers = std::min(workers, count);
            // Arm the latch and grab the future BEFORE pushing: a pushed worker may run (and free
            // the state) before this call returns.
            state->remaining.store(helpers, std::memory_order_relaxed);
            Future<void> fut = state->promise.get_future();
            size_t pushed = 0;
            try
            {
                for (; pushed < helpers; ++pushed)
                    PushImpl<ParallelForAsyncJob<std::remove_reference_t<Fn>>>(JobPriority::Normal, state);
            }
            catch (...)
            {
                // Some worker jobs failed to enqueue: retire their latch slots so the remaining
                // (successfully pushed) workers can still satisfy the promise. If none were pushed
                // or they have all already finished, finish the promise here.
                size_t const missing = helpers - pushed;
                if (missing != 0 && state->remaining.fetch_sub(missing, std::memory_order_acq_rel) == missing)
                {
                    Allocator* alloc = state->alloc;
                    state->promise.set_value();
                    Destruct(alloc, state);
                }
                throw;
            }
            return fut;
        }
        template <typename Fn>
        [[nodiscard]] Future<void> ParallelForAsync(size_t count, Fn&& fn)
        {
            return ParallelForAsync(ExecutionPolicy::Par, count, std::forward<Fn>(fn));
        }
        template <typename It, typename Fn>
            requires std::random_access_iterator<It>
        [[nodiscard]] Future<void> ParallelForAsync(ExecutionPolicy policy, It first, It last, Fn&& fn)
        {
            auto const count = last - first;
            if (count <= 0)
                return MakeReadyFuture();
            return ParallelForAsync(policy, static_cast<size_t>(count),
                [first, fn = std::forward<Fn>(fn)](size_t i, size_t worker) mutable
                { ParallelForInvoke(fn, first[static_cast<std::iter_difference_t<It>>(i)], worker); });
        }
        template <typename It, typename Fn>
            requires std::random_access_iterator<It>
        [[nodiscard]] Future<void> ParallelForAsync(It first, It last, Fn&& fn)
        {
            return ParallelForAsync(ExecutionPolicy::Par, first, last, std::forward<Fn>(fn));
        }
 
        void Shutdown();
        void Join();
        ~ThreadPool();
 
        [[nodiscard]] size_t GetPendingJobCount() const noexcept
        {
            return mTotal.load(std::memory_order_relaxed) - mComplete.load(std::memory_order_relaxed);
        }
        [[nodiscard]] size_t GetCompletedJobCount() const noexcept { return mComplete.load(std::memory_order_relaxed); }
        [[nodiscard]] size_t GetTotalJobCount() const noexcept { return mTotal.load(std::memory_order_relaxed); }
 
        const static size_t CalcTaskSize(size_t size) { return 1ULL << static_cast<size_t>(std::ceil(std::log2f(size))); }
    };
} // namespace Foundation::Core