ThreadPool.hpp

#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
 
 
namespace mltk {
// ============================================================================================= //
//                                    Begin class thread_pool                                    //
 
class ThreadPool
{
    typedef std::uint_fast32_t ui32;
    typedef std::uint_fast64_t ui64;
 
public:
    // ============================
    // Constructors and destructors
    // ============================
 
    ThreadPool(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();
    }
 
    ~ThreadPool()
    {
        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(&ThreadPool::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                                        //
// ============================================================================================= //
 
} // namespace mltk