ThreadPool.hpp

Go to the documentation of this file.

#pragma once
#include "AtomicQueue.hpp"
#include "Thread.hpp"
#include <cmath>
namespace Foundation::Core
{
    struct ThreadPoolJob
    {
        virtual ~ThreadPoolJob() = default;
        virtual void Execute(size_t id) noexcept = 0;
    };
    template <typename Lambda, typename ReturnType, typename... Args>
    struct ThreadPoolLambdaJob final : public ThreadPoolJob
    {
        Lambda mFunc;
        Promise<ReturnType> mPromise;
        ThreadPoolLambdaJob(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());
            }
        }
    };
    using JobQueue = MPMCQueue<UniquePtr<ThreadPoolJob>>;
    class ThreadPool
    {
        Allocator* mAllocator;
        String mName;
        Atomic<bool> mShutdown{false};
        Atomic<size_t> mComplete{0};
        Atomic<size_t> mTotal{0};
 
        JobQueue mJobs;
        // Ensure threads are joined first on destruction
        Vector<Thread> mThreads;
        void ThreadPoolWorker(size_t id);
 
    public:
        ThreadPool(size_t numThreads, size_t maxTasks, Allocator* alloc, StringView name = "ThreadPool");
        template <typename T, typename... Args>
            requires std::is_base_of_v<ThreadPoolJob, T>
        T* PushImpl(Args&&... args)
        {
            if (mShutdown)
                throw std::runtime_error("ThreadPool shutting down");
            auto task = ConstructUniqueBase<ThreadPoolJob, T>(mAllocator, std::forward<Args>(args)...);
            T* ptr = static_cast<T*>(task.get());
            if (!mJobs.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 Push(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());
            ThreadPoolLambdaJob<LambdaType, ReturnType> job(std::forward<LambdaType>(LambdaFn));
            auto fut = job.mPromise.get_future();
            PushImpl<ThreadPoolLambdaJob<LambdaType, ReturnType>>(std::move(job));
            return std::move(fut);
        }
        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); }
 
        constexpr static size_t getTaskSize(size_t size) { return 1ULL << static_cast<size_t>(std::ceil(std::log2f(size))); }
    };
} // namespace Foundation::Core