ThreadPool.hpp

Go to the documentation of this file.

#pragma once
#include <Atomics/Queue.hpp>
#include <Bits/Ranges.hpp>
#include <Core/Core.hpp>
#include "Future.hpp"
#include "Thread.hpp"
namespace Foundation::Async
{
    using namespace Core;
    using namespace Async;
    using namespace Atomics;
    struct ThreadPoolJob
    {
        virtual ~ThreadPoolJob() = default;
        virtual void Execute(size_t id) noexcept = 0;
    };
    template <typename Lambda, typename ReturnType>
    class ThreadPoolLambdaJob final : public ThreadPoolJob
    {
        Lambda mFunc;
        SharedPromise<ReturnType> mPromise;
    public:
        ThreadPoolLambdaJob(SharedPromise<ReturnType> promise, Lambda&& func) : mFunc(func), mPromise(promise) {}
        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;
        JobQueue::Writer mJobsWriter;
        // 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>
        void PushImpl(Args&&... args)
        {
            CHECK_MSG(!mShutdown, "ThreadPool shutting down");
            CHECK_MSG(mJobsWriter.Push(
                ConstructUniqueBase<ThreadPoolJob, T>(mAllocator, std::forward<Args>(args)...)), "Jobs full");
            mTotal.fetch_add(1, std::memory_order_relaxed);
            mTotal.notify_one();
        }
        template <typename Lambda, typename... Args>
        auto Push(Lambda&& func, Args&&... args)
        {
            CHECK_MSG(!mShutdown, "ThreadPool shutting down");
            auto ThreadPoolPackagedLambda = [](Lambda&& fn, Args&&... fargs)
            {
                return [func = std::forward<Lambda>(fn), ... fargs = std::forward<Args>(fargs)]
                { return func(std::forward<Args>(fargs)...); };
            };
            auto packaged = ThreadPoolPackagedLambda(std::forward<Lambda>(func), std::forward<Args>(args)...);
            using ReturnType = decltype(func(args...));
            using PackagedType = decltype(packaged);
            // Use the wrapped lambda type for the job
            using LambdaType = ThreadPoolLambdaJob<PackagedType, ReturnType>;
            SharedPromise<ReturnType> promise = ConstructShared<std::promise<ReturnType>>(mAllocator);
            CHECK_MSG(mJobsWriter.Push(
                ConstructUniqueBase<ThreadPoolJob, LambdaType>(
                    mAllocator, promise,
                    std::forward<PackagedType>(packaged))), "Jobs full");
            mTotal.fetch_add(1, std::memory_order_relaxed);
            mTotal.notify_one();
            return promise;
        }
        void Shutdown();
        void Join();
        ~ThreadPool();
    };
} // namespace Foundation::Async