thread_pool.hpp

Go to the documentation of this file.

#pragma once
 
#define THREAD_POOL_VERSION "v2.0.0 (2021-08-14)"
 
#include <atomic>      // std::atomic
#include <chrono>      // std::chrono
#include <cstdint>     // std::int_fast64_t, std::uint_fast32_t
#include <functional>  // std::function
#include <future>      // std::future, std::promise
#include <iostream>    // std::cout, std::ostream
#include <memory>      // std::shared_ptr, std::unique_ptr
#include <mutex>       // std::mutex, std::scoped_lock
#include <queue>       // std::queue
#include <thread>      // std::this_thread, std::thread
#include <type_traits> // std::common_type_t, std::decay_t, std::enable_if_t, std::is_void_v, std::invoke_result_t
#include <utility>     // std::move
 
// ============================================================================================= //
//                                    Begin class thread_pool                                    //
 
class thread_pool {
    typedef std::uint_fast32_t ui32;
    typedef std::uint_fast64_t ui64;
 
  public:
    // ============================
    // Constructors and destructors
    // ============================
 
    thread_pool(const ui32 &_thread_count = std::thread::hardware_concurrency())
        : thread_count(_thread_count ? _thread_count : std::thread::hardware_concurrency()), threads(new std::thread[_thread_count ? _thread_count : std::thread::hardware_concurrency()]) {
        create_threads();
    }
 
    ~thread_pool() {
        wait_for_tasks();
        running = false;
        destroy_threads();
    }
 
    // =======================
    // Public member functions
    // =======================
 
    ui64 get_tasks_queued() const {
        const std::scoped_lock lock(queue_mutex);
        return tasks.size();
    }
 
    ui32 get_tasks_running() const {
        return tasks_total - (ui32)get_tasks_queued();
    }
 
    ui32 get_tasks_total() const {
        return tasks_total;
    }
 
    ui32 get_thread_count() const {
        return thread_count;
    }
 
    template <typename T1, typename T2, typename F>
    void parallelize_loop(const T1 &first_index, const T2 &index_after_last, const F &loop, ui32 num_blocks = 0) {
        typedef std::common_type_t<T1, T2> T;
        T the_first_index = (T)first_index;
        T last_index = (T)index_after_last;
        if (the_first_index == last_index)
            return;
        if (last_index < the_first_index) {
            T temp = last_index;
            last_index = the_first_index;
            the_first_index = temp;
        }
        last_index--;
        if (num_blocks == 0)
            num_blocks = thread_count;
        ui64 total_size = (ui64)(last_index - the_first_index + 1);
        ui64 block_size = (ui64)(total_size / num_blocks);
        if (block_size == 0) {
            block_size = 1;
            num_blocks = (ui32)total_size > 1 ? (ui32)total_size : 1;
        }
        std::atomic<ui32> blocks_running = 0;
        for (ui32 t = 0; t < num_blocks; t++) {
            T start = ((T)(t * block_size) + the_first_index);
            T end = (t == num_blocks - 1) ? last_index + 1 : ((T)((t + 1) * block_size) + the_first_index);
            blocks_running++;
            push_task([start, end, &loop, &blocks_running] {
                loop(start, end);
                blocks_running--;
            });
        }
        while (blocks_running != 0) {
            sleep_or_yield();
        }
    }
 
    template <typename F>
    void push_task(const F &task) {
        tasks_total++;
        {
            const std::scoped_lock lock(queue_mutex);
            tasks.push(std::function<void()>(task));
        }
    }
 
    template <typename F, typename... A>
    void push_task(const F &task, const A &...args) {
        push_task([task, args...] { task(args...); });
    }
 
    void reset(const ui32 &_thread_count = std::thread::hardware_concurrency()) {
        bool was_paused = paused;
        paused = true;
        wait_for_tasks();
        running = false;
        destroy_threads();
        thread_count = _thread_count ? _thread_count : std::thread::hardware_concurrency();
        threads.reset(new std::thread[thread_count]);
        paused = was_paused;
        running = true;
        create_threads();
    }
 
    template <typename F, typename... A, typename = std::enable_if_t<std::is_void_v<std::invoke_result_t<std::decay_t<F>, std::decay_t<A>...>>>>
    std::future<bool> submit(const F &task, const A &...args) {
        std::shared_ptr<std::promise<bool>> task_promise(new std::promise<bool>);
        std::future<bool> future = task_promise->get_future();
        push_task([task, args..., task_promise] {
            try {
                task(args...);
                task_promise->set_value(true);
            } catch (...) {
                try {
                    task_promise->set_exception(std::current_exception());
                } catch (...) {
                }
            }
        });
        return future;
    }
 
    template <typename F, typename... A, typename R = std::invoke_result_t<std::decay_t<F>, std::decay_t<A>...>, typename = std::enable_if_t<!std::is_void_v<R>>>
    std::future<R> submit(const F &task, const A &...args) {
        std::shared_ptr<std::promise<R>> task_promise(new std::promise<R>);
        std::future<R> future = task_promise->get_future();
        push_task([task, args..., task_promise] {
            try {
                task_promise->set_value(task(args...));
            } catch (...) {
                try {
                    task_promise->set_exception(std::current_exception());
                } catch (...) {
                }
            }
        });
        return future;
    }
 
    void wait_for_tasks() {
        while (true) {
            if (!paused) {
                if (tasks_total == 0)
                    break;
            } else {
                if (get_tasks_running() == 0)
                    break;
            }
            sleep_or_yield();
        }
    }
 
    // ===========
    // Public data
    // ===========
 
    std::atomic<bool> paused = false;
 
    ui32 sleep_duration = 1000;
 
  private:
    // ========================
    // Private member functions
    // ========================
 
    void create_threads() {
        for (ui32 i = 0; i < thread_count; i++) {
            threads[i] = std::thread(&thread_pool::worker, this);
        }
    }
 
    void destroy_threads() {
        for (ui32 i = 0; i < thread_count; i++) {
            threads[i].join();
        }
    }
 
    bool pop_task(std::function<void()> &task) {
        const std::scoped_lock lock(queue_mutex);
        if (tasks.empty())
            return false;
        else {
            task = std::move(tasks.front());
            tasks.pop();
            return true;
        }
    }
 
    void sleep_or_yield() {
        if (sleep_duration)
            std::this_thread::sleep_for(std::chrono::microseconds(sleep_duration));
        else
            std::this_thread::yield();
    }
 
    void worker() {
        while (running) {
            std::function<void()> task;
            if (!paused && pop_task(task)) {
                task();
                tasks_total--;
            } else {
                sleep_or_yield();
            }
        }
    }
 
    // ============
    // Private data
    // ============
 
    mutable std::mutex queue_mutex = {};
 
    std::atomic<bool> running = true;
 
    std::queue<std::function<void()>> tasks = {};
 
    ui32 thread_count;
 
    std::unique_ptr<std::thread[]> threads;
 
    std::atomic<ui32> tasks_total = 0;
};
 
//                                     End class thread_pool                                     //
// ============================================================================================= //
 
// ============================================================================================= //
//                                   Begin class synced_stream                                   //
 
class synced_stream {
  public:
    synced_stream(std::ostream &_out_stream = std::cout)
        : out_stream(_out_stream){};
 
    template <typename... T>
    void print(const T &...items) {
        const std::scoped_lock lock(stream_mutex);
        (out_stream << ... << items);
    }
 
    template <typename... T>
    void println(const T &...items) {
        print(items..., '\n');
    }
 
  private:
    mutable std::mutex stream_mutex = {};
 
    std::ostream &out_stream;
};
 
//                                    End class synced_stream                                    //
// ============================================================================================= //
 
// ============================================================================================= //
//                                       Begin class timer                                       //
 
class timer {
    typedef std::int_fast64_t i64;
 
  public:
    void start() {
        start_time = std::chrono::steady_clock::now();
    }
 
    void stop() {
        elapsed_time = std::chrono::steady_clock::now() - start_time;
    }
 
    i64 ms() const {
        return (std::chrono::duration_cast<std::chrono::milliseconds>(elapsed_time)).count();
    }
 
  private:
    std::chrono::time_point<std::chrono::steady_clock> start_time = std::chrono::steady_clock::now();
 
    std::chrono::duration<double> elapsed_time = std::chrono::duration<double>::zero();
};
 
//                                        End class timer                                        //
// ============================================================================================= //