thread_pool.hpp

Go to the documentation of this file.

#pragma once

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <exception>
#include <functional>
#include <future>
#include <iostream>
#include <memory>
#include <mutex>
#include <queue>
#include <thread>
#include <type_traits>
#include <utility>
#include <vector>

class thread_pool {
    using concurrency_t = std::invoke_result_t<decltype(std::thread::hardware_concurrency)>;

    std::atomic<bool> running = false;

    std::condition_variable task_available_cv = {};

    std::condition_variable task_done_cv = {};

    std::queue<std::function<void()>> tasks = {};

    std::atomic<size_t> tasks_total = 0;

    mutable std::mutex tasks_mutex = {};

    concurrency_t thread_count = 0;

    std::unique_ptr<std::thread[]> threads = nullptr;

    std::atomic<bool> waiting = false;

    void create_threads() {
        running = true;
        for(concurrency_t i = 0; i < thread_count; ++i) threads[i] = std::thread(&thread_pool::worker, this);
    }

    void destroy_threads() {
        running = false;
        task_available_cv.notify_all();
        for(concurrency_t i = 0; i < thread_count; ++i) threads[i].join();
    }

    void worker() {
        while(running) {
            std::function<void()> task;
            std::unique_lock<std::mutex> tasks_lock(tasks_mutex);
            task_available_cv.wait(tasks_lock, [&] { return !tasks.empty() || !running; });
            if(running && !paused) {
                task = std::move(tasks.front());
                tasks.pop();
                tasks_lock.unlock();
                task();
                --tasks_total;
                if(waiting) task_done_cv.notify_one();
            }
        }
    }

public:
    explicit thread_pool(const concurrency_t thread_count_ = std::thread::hardware_concurrency())
        : thread_count(thread_count_ > 1 ? thread_count_ : 1), threads(std::make_unique<std::thread[]>(thread_count)) { create_threads(); }

    ~thread_pool() {
        wait_for_tasks();
        destroy_threads();
    }

    size_t get_tasks_queued() const {
        const std::scoped_lock tasks_lock(tasks_mutex);
        return tasks.size();
    }

    size_t get_tasks_running() const {
        const std::scoped_lock tasks_lock(tasks_mutex);
        return tasks_total - tasks.size();
    }

    size_t get_tasks_total() const { return tasks_total; }

    concurrency_t get_thread_count() const { return thread_count; }

    template<typename F, typename... A>
    void push_task(const F& task, const A&... args) {
        {
            const std::scoped_lock tasks_lock(tasks_mutex);
            if constexpr(sizeof...(args) == 0)
                tasks.push(std::function<void()>(task));
            else
                tasks.push(std::function<void()>([task, args...] { task(args...); }));
        }
        ++tasks_total;
        task_available_cv.notify_one();
    }

    void reset(const concurrency_t thread_count_ = std::thread::hardware_concurrency()) {
        const bool was_paused = paused;
        paused = true;
        wait_for_tasks();
        destroy_threads();
        thread_count = thread_count_ > 1 ? thread_count_ : 1;
        threads = std::make_unique<std::thread[]>(thread_count);
        paused = was_paused;
        create_threads();
    }

    template<typename F, typename... A, typename R = std::invoke_result_t<std::decay_t<F>, std::decay_t<A>...>>
    std::future<R> submit(const F& task, const A&... args) {
        std::shared_ptr<std::promise<R>> task_promise = std::make_shared<std::promise<R>>();
        push_task(
            [task, args..., task_promise] {
                try {
                    if constexpr(std::is_void_v<R>) {
                        task(args...);
                        task_promise->set_value();
                    }
                    else {
                        task_promise->set_value(task(args...));
                    }
                }
                catch(...) {
                    try {
                        task_promise->set_exception(std::current_exception());
                    }
                    catch(...) {
                    }
                }
            });
        return task_promise->get_future();
    }

    void wait_for_tasks() {
        waiting = true;
        std::unique_lock<std::mutex> tasks_lock(tasks_mutex);
        task_done_cv.wait(tasks_lock, [this] { return (tasks_total == (paused ? tasks.size() : 0)); });
        waiting = false;
    }

    std::atomic<bool> paused = false;
};