thread_pool.hpp

Go to the documentation of this file.

#pragma once
 
#include "threading/jthread_compat.hpp"
 
#include <condition_variable>
#include <cstddef>
#include <functional>
#include <future>
#include <memory>
#include <mutex>
#include <queue>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace libstp::threading
{
    class ThreadPool
    {
    public:
        explicit ThreadPool(std::size_t worker_count = 0);
        ~ThreadPool();
 
        ThreadPool(const ThreadPool&) = delete;
        ThreadPool& operator=(const ThreadPool&) = delete;

        template <typename F, typename... Args>
        auto submit(F&& fn, Args&&... args)
            -> std::future<std::invoke_result_t<std::decay_t<F>, std::decay_t<Args>...>>
        {
            using R = std::invoke_result_t<std::decay_t<F>, std::decay_t<Args>...>;
            auto task = std::make_shared<std::packaged_task<R()>>(
                [fn = std::forward<F>(fn),
                 tup = std::make_tuple(std::forward<Args>(args)...)]() mutable
                {
                    return std::apply(std::move(fn), std::move(tup));
                });
            auto fut = task->get_future();
            enqueue([task]() { (*task)(); });
            return fut;
        }

        void wait_idle();

        void shutdown();
 
        [[nodiscard]] std::size_t worker_count() const noexcept { return workers_.size(); }
 
    private:
        void enqueue(std::function<void()> job);
        void worker_loop(stop_token stop);
 
        mutable std::mutex                 mu_;
        std::condition_variable_any        work_cv_;
        std::condition_variable_any        idle_cv_;
        std::queue<std::function<void()>>  queue_;
        std::size_t                        in_flight_ = 0;
        bool                               stopping_  = false;
        std::vector<jthread>          workers_;
    };
}