thread_pool.h

#pragma once
 
#include <algorithm>
#include <atomic>
#include <condition_variable>
#include <deque>
#include <exception>
#include <functional>
#include <mutex>
#include <thread>
#include <type_traits>
#include <vector>
 
namespace Vkm::Engine {

class ThreadPool {
    public:
        class Batch {
            public:
                template<typename Fn>
                Batch(size_t count, const Fn& task)
                    : m_run([](const void* body, size_t index) { (*static_cast<const Fn*>(body))(index); })
                    , m_task(&task)
                    , m_count(count)
                {}

                template<typename Fn>
                Batch(size_t count, const Fn && task) = delete;
 
                ~Batch() = default;
 
                Batch(const Batch& other) = delete;
                Batch& operator=(const Batch& other) = delete;
 
                Batch(Batch && other) = delete;
                Batch& operator=(Batch && other) = delete;
 
            private:
                friend class ThreadPool;
 
            private:
                void (*m_run)(const void* task, size_t index);
                const void* m_task;
                size_t      m_count;
                size_t      m_claimed = 0;          
                std::exception_ptr m_error;         
 
                std::atomic<size_t> m_pending{0};   
        };
 
    public:
        ThreadPool(const ThreadPool& other) = delete;
        ThreadPool& operator=(const ThreadPool& other) = delete;
 
        ThreadPool(ThreadPool && other) = delete;
        ThreadPool& operator=(ThreadPool && other) = delete;
 
    public:
        static ThreadPool& get();

        void addTask(std::function<void()> && task);

        void addBatch(Batch& batch);

        void waitForBatch(Batch& batch);

        void shutdown();

        static bool isWorkerThread();
 
        size_t threadCount() const { return m_threads.size(); }
 
    private:
        struct QueuedTask {
            std::function<void()> function;           
            Batch*                batch   = nullptr;
            size_t                index   = 0;        
        };
 
    private:
        ThreadPool(size_t threadCount);
        ~ThreadPool();

        void process();

        void runIndex(Batch& batch, size_t index);

        void retire(std::atomic<size_t>& pending);
 
    private:
        std::atomic<bool> m_running;
 
        std::vector<std::thread> m_threads;
 
        std::deque<QueuedTask> m_frameTasks;       
        std::deque<QueuedTask> m_backgroundTasks;  
 
        std::mutex m_tasksMutex;
        std::condition_variable m_tasksCV;
        std::condition_variable m_doneCV;   
};

template<class Function>
void parallelFor(size_t count, size_t grain, Function&& function) {
    if (count == 0) {
        return;
    }
 
    grain = std::max(grain, size_t(1));
 
    auto invokeAt = [&](size_t i) {
        if constexpr (std::is_invocable_v<Function, size_t>) {
            function(i);
        } else {
            function();
        }
    };
 
    if (ThreadPool::isWorkerThread()) {
        for (size_t i = 0; i < count; ++i) invokeAt(i);
        return;
    }
 
    auto& pool = ThreadPool::get();
 
    // No workers after shutdown (or with no cores reported): a submission would never retire.
    if (pool.threadCount() == 0) {
        for (size_t i = 0; i < count; ++i) invokeAt(i);
        return;
    }
 
    // Chunk 0 is the caller's; batch index i is chunk i + 1.
    const size_t chunks = (count - 1) / grain + 1;
    const auto runChunk = [&](size_t chunk) {
        const size_t end = std::min(count, (chunk + 1) * grain);
        for (size_t index = chunk * grain; index < end; ++index) invokeAt(index);
    };
    const auto runQueued = [&](size_t task) { runChunk(task + 1); };
    ThreadPool::Batch batch(chunks - 1, runQueued);
 
    pool.addBatch(batch);
 
    // Queued indices reach into this frame, so join before the caller's exception leaves it;
    // that exception, not a queued chunk's, propagates.
    try {
        runChunk(0);
    } catch (...) {
        try {
            pool.waitForBatch(batch);
        } catch (...) {
        }
        throw;
    }
    pool.waitForBatch(batch);
}

template<class Function>
void parallelFor(size_t count, Function&& function) {
    auto& pool = ThreadPool::get();
 
    // Below this the dispatch cost (mutex, notify_all, done-CV round trip) dwarfs the work;
    // grain == count sweeps serially.
    constexpr size_t MIN_PARALLEL = 2048;
 
    // The +1 is the calling thread, which runs a chunk too.
    const size_t grain = (count < MIN_PARALLEL)
        ? count
        : count / (pool.threadCount() + 1);
 
    parallelFor(count, grain, function);
}
 
} // namespace Vkm::Engine